"""Runs one identical write trace three ways, changing only which segments
the compaction covers and whether it drops tombstones.

Two of the three answer get("apple") correctly. The third resurrects it.

Everything here is deterministic: every scenario gets its own directory
under `data/`, wiped and rebuilt from scratch, segment names come from a
counter, and no clock or RNG is consulted anywhere. Run it twice and diff.

The segment files are left behind on purpose -- `cat data/store/*.sst`
after a run and you are looking at the whole database.
"""

import os
import shutil

from lsm import DEL, LSMTree

DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")

# The trace. Threshold 2 means every second write seals a segment, so the
# tombstone for `apple` lands in seg0001 -- one segment newer than the
# `apple=red` it exists to hide. That separation is what makes the aha
# possible; see the commentary, section 6.4.
TRACE = [
    ("put", "apple", "red"),
    ("put", "banana", "yellow"),
    ("del", "apple", None),
    ("put", "cherry", "dark"),
    ("put", "date", "brown"),
    ("put", "elder", "black"),
]


def build(tag, flush_threshold=2):
    path = os.path.join(DATA, tag)
    shutil.rmtree(path, ignore_errors=True)
    tree = LSMTree(path, flush_threshold=flush_threshold)
    for op, key, value in TRACE:
        if op == "put":
            tree.put(key, value)
        else:
            tree.delete(key)
    return tree


def render(seg):
    cells = [f"{k}=<<tombstone>>" if op == DEL else f"{k}={v}"
             for k, (op, v) in seg.items()]
    return "  ".join(cells)


def show_segments(tree, label):
    print(f"  {label}")
    for seg in tree.segments:
        print(f"    {seg.name}  {render(seg)}")


def footprint(tree):
    n = len(tree.segments)
    return f"[{tree.store_bytes()} bytes in {n} segment{'' if n == 1 else 's'}]"


def scenario(tag, title, n, drop_tombstones):
    tree = build(tag)
    print(f"\n--- {title} ---")
    print(f"  get(apple) before compaction: {tree.get('apple')}"
          f"   {footprint(tree)}")
    tree.compact(n, drop_tombstones=drop_tombstones)
    show_segments(tree, "segments after:")
    after = tree.get("apple")
    verdict = "RESURRECTED" if after is not None else "still deleted"
    print(f"  get(apple) after  compaction: {after}"
          f"   {footprint(tree)}   <-- {verdict}")
    return after


def main():
    tree = build("store")
    print("=== The store after six writes (flush every 2 entries) ===\n")
    show_segments(tree, "segments, oldest first:")
    print(f"\n  get(apple) = {tree.get('apple')}"
          f"   (read {tree.last_get_segments} segments to say so)")
    print(f"  bytes on disk: {tree.store_bytes()}"
          f"   entries on disk: "
          f"{sum(1 for s in tree.segments for _ in s.items())}"
          f"   live keys: 4")

    print("\n\n=== The same trace, compacted three ways ===")
    a = scenario("partial-drop", "compact(newest 2, drop_tombstones=True)",
                 2, True)
    b = scenario("partial-keep", "compact(newest 2, drop_tombstones=False)",
                 2, False)
    c = scenario("full-drop", "compact(all 3, drop_tombstones=True)", 3, True)
    print(f"\n  partial + drop -> {a!r}   partial + keep -> {b!r}   "
          f"full + drop -> {c!r}")

    print("\n\n=== Supporting measurement 1: deleting data grows the store ===\n")
    path = os.path.join(DATA, "deletes")
    shutil.rmtree(path, ignore_errors=True)
    tree = LSMTree(path, flush_threshold=2)
    for key, value in [("apple", "red"), ("banana", "yellow"),
                       ("cherry", "dark"), ("date", "brown")]:
        tree.put(key, value)
    print(f"  4 keys written:  {len(tree.segments)} segments, "
          f"{tree.store_bytes()} bytes")
    for key in ["apple", "banana", "cherry", "date"]:
        tree.delete(key)
    print(f"  all 4 deleted:   {len(tree.segments)} segments, "
          f"{tree.store_bytes()} bytes")
    print(f"  the empty store holds "
          f"{sum(1 for s in tree.segments for _ in s.items())} entries "
          f"and answers get(apple) = {tree.get('apple')}")

    print("\n\n=== Supporting measurement 2: a miss is the costliest read ===\n")
    path = os.path.join(DATA, "amplification")
    shutil.rmtree(path, ignore_errors=True)
    tree = LSMTree(path, flush_threshold=2)
    for i in range(10):
        tree.put(f"k{i:02d}", f"v{i:02d}")
    print(f"  10 keys -> {len(tree.segments)} segments on disk\n")
    for key, note in [("k09", "newest key"), ("k00", "oldest key"),
                      ("zzz", "never written")]:
        value = tree.get(key)
        print(f"  get({key}) = {str(value):>6}   "
              f"segments read: {tree.last_get_segments}   "
              f"lines read: {tree.last_get_lines:>2}   ({note})")


if __name__ == "__main__":
    main()