#!/usr/bin/env python3
"""Hands-on P04 — an LSM storage engine, assembled from seven lego blocks."""
import hashlib, math, os, random, struct, tempfile, zlib
from _harness import block, run_all

rng = random.Random(0)
DIR = tempfile.mkdtemp(prefix="h04-")

@block(1, "The record format", "a length-prefixed, checksummed record is the atom of durability")
def b1(s, show):
    def pack(k, v):
        body = struct.pack("<II", len(k), len(v)) + k + v
        return struct.pack("<I", zlib.crc32(body)) + struct.pack("<I", len(body)) + body
    def unpack(buf, off):
        if off + 8 > len(buf): return None, off
        crc, blen = struct.unpack_from("<II", buf, off)
        if off + 8 + blen > len(buf): return None, off        # torn tail
        body = buf[off+8: off+8+blen]
        if zlib.crc32(body) != crc: return None, off          # corruption
        kl, vl = struct.unpack_from("<II", body, 0)
        return (body[8:8+kl], body[8+kl:8+kl+vl]), off + 8 + blen
    r = pack(b"alpha", b"1")
    got, _ = unpack(r, 0)
    assert got == (b"alpha", b"1")
    torn, _ = unpack(r[:-3], 0)
    bad = bytearray(r); bad[-1] ^= 0xFF
    corrupt, _ = unpack(bytes(bad), 0)
    if show:
        print(f"  record for ('alpha','1') = {len(r)} bytes: crc|len|klen|vlen|k|v")
        print(f"  round-trip: {got}")
        print(f"  truncated tail -> {torn}   (normal after a crash, not corruption)")
        print(f"  flipped bit    -> {corrupt}   (detected by CRC, never served)")
    return {"pack": pack, "unpack": unpack}

@block(2, "Write-ahead log", "log BEFORE data, or a crash leaves neither version")
def b2(s, show):
    class WAL:
        def __init__(self, path): self.path = path; self.f = open(path, "ab")
        def append(self, k, v, sync=False):
            self.f.write(s["pack"](k, v))
            self.f.flush()
            if sync: os.fsync(self.f.fileno())
        def replay(self):
            buf = open(self.path, "rb").read(); off = 0; out = []
            while off < len(buf):
                rec, noff = s["unpack"](buf, off)
                if rec is None: break                     # stop at first bad record
                out.append(rec); off = noff
            return out, len(buf) - off
    w = WAL(os.path.join(DIR, "wal"))
    for i in range(50): w.append(f"k{i:03}".encode(), f"v{i}".encode())
    w.f.flush()
    with open(w.path, "ab") as f: f.write(b"\x01\x02\x03")   # simulate a torn write
    recs, leftover = w.replay()
    if show:
        print(f"  50 records written, then 3 junk bytes appended (a torn write)")
        print(f"  replay recovered {len(recs)} records, stopped with {leftover} bytes left")
        print("  a partial record at the tail is EXPECTED after a crash. Stopping")
        print("  there is correct; scanning past it is how you serve garbage.")
    return {"WAL": WAL}

@block(3, "Memtable and flush", "sorted in memory, immutable on disk")
def b3(s, show):
    def flush(memtable, path):
        items = sorted(memtable.items())
        with open(path, "wb") as f:
            offsets = []
            for k, v in items:
                offsets.append((k, f.tell()))
                f.write(s["pack"](k, v))
        return items, offsets
    mt = {f"key{i:04}".encode(): f"val{i}".encode() for i in rng.sample(range(500), 200)}
    items, offs = flush(mt, os.path.join(DIR, "sst0"))
    if show:
        print(f"  memtable {len(mt)} keys (a dict) -> SSTable, sorted on disk")
        print(f"  first three keys: {[k.decode() for k,_ in items[:3]]}")
        print(f"  sorted: {items == sorted(items)}")
        print("  sortedness makes range scans a merge and lookups a binary search;")
        print("  immutability makes concurrent reads lock-free.")
    return {"flush": flush}

@block(4, "Sparse index", "one entry per BLOCK, not per key -- that is what fits in RAM")
def b4(s, show):
    def build_sparse(offsets, every=16):
        return offsets[::every]
    def seek(sparse, key):
        lo = None
        for k, off in sparse:
            if k <= key: lo = off
            else: break
        return lo or 0
    mt = {f"key{i:04}".encode(): f"v{i}".encode() for i in range(1000)}
    items, offs = s["flush"](mt, os.path.join(DIR, "sst1"))
    sp = build_sparse(offs, 16)
    if show:
        print(f"  {len(offs)} keys -> {len(sp)} sparse entries ({len(offs)//len(sp)}x smaller)")
        print(f"  lookup 'key0500': scan starts at byte {seek(sp, b'key0500')}, "
              f"not byte 0")
        print("  a dense index over a billion keys does not fit in memory. A sparse")
        print("  one narrows to a block, and you scan the block.")
    return {"build_sparse": build_sparse, "seek": seek}

@block(5, "Bloom filter", "the read path's whole viability, for 10 bits per key")
def b5(s, show):
    class Bloom:
        def __init__(self, n, bpk=10):
            self.m = max(8, int(n * bpk)); self.k = max(1, round(bpk * math.log(2)))
            self.bits = bytearray((self.m + 7) // 8)
        def _p(self, key):
            d = hashlib.blake2b(key, digest_size=16).digest()
            h1 = int.from_bytes(d[:8], "little"); h2 = int.from_bytes(d[8:], "little") | 1
            for i in range(self.k): yield (h1 + i * h2) % self.m
        def add(self, key):
            for p in self._p(key): self.bits[p >> 3] |= 1 << (p & 7)
        def __contains__(self, key):
            return all(self.bits[p >> 3] >> (p & 7) & 1 for p in self._p(key))
    n = 20000
    keys = [f"present{i}".encode() for i in range(n)]
    if show:
        print(f"  {'bits/key':>9}{'k':>4}{'theory':>10}{'measured':>10}{'RAM':>10}")
        for bpk in (4, 8, 10, 16):
            bf = Bloom(n, bpk)
            for k in keys: bf.add(k)
            assert all(k in bf for k in keys), "FALSE NEGATIVE -- not a Bloom filter"
            absent = [f"absent{i}".encode() for i in range(40000)]
            meas = sum(1 for k in absent if k in bf) / len(absent)
            th = (1 - math.exp(-bf.k * n / bf.m)) ** bf.k
            print(f"  {bpk:>9}{bf.k:>4}{th:>10.5f}{meas:>10.5f}"
                  f"{len(bf.bits)/1024:>9.0f}K")
        print("  never a false negative -- that is the one guarantee. Theory tracks")
        print("  measurement within a few percent at every setting.")
    return {"Bloom": Bloom}

@block(6, "The read path", "newest run first, and count every block you touch")
def b6(s, show):
    class Run:
        def __init__(self, items, bpk=10):
            self.d = dict(items)
            self.bloom = s["Bloom"](max(1, len(items)), bpk) if bpk else None
            if self.bloom:
                for k in self.d: self.bloom.add(k)
        def get(self, key, stats, use_bloom=True):
            if use_bloom and self.bloom is not None and key not in self.bloom:
                stats["skipped"] += 1; return None
            stats["block_reads"] += 1
            return self.d.get(key)
    def make_db(n_runs=30, per_run=1500, bpk=10):
        runs, all_keys = [], []
        for r in range(n_runs):
            items = [(f"k:{r}:{i}".encode(), f"v{i}".encode()) for i in range(per_run)]
            all_keys += [k for k, _ in items]
            runs.append(Run(items, bpk))
        return runs, all_keys
    def get(runs, key, use_bloom=True):
        st = {"block_reads": 0, "skipped": 0}
        for run in reversed(runs):
            v = run.get(key, st, use_bloom)
            if v is not None: return v, st
        return None, st
    if show:
        runs, keys = make_db()
        st = get(runs, b"nope", True)[1]
        print(f"  30 runs. absent key WITH bloom: {st['block_reads']} block reads, "
              f"{st['skipped']} skipped")
        st = get(runs, b"nope", False)[1]
        print(f"  absent key WITHOUT bloom:      {st['block_reads']} block reads")
        print("  that ratio IS the read path. Everything else is bookkeeping.")
    return {"make_db": make_db, "get": get, "Run": Run}

@block(7, "Compaction, and the three amplifications", "you are choosing which cost to pay")
def b7(s, show):
    def amp(T, L):
        return dict(lev=(T*L+1, L+1, 1+1/T), tier=(L+1, T*L, 2.0))
    if show:
        print(f"  {'data':>8}{'levels':>8}{'leveled W/R/S':>20}{'tiered W/R/S':>18}")
        for gb in (1, 8, 64, 512):
            L = max(1, math.ceil(math.log(gb*1e9/64e6, 10)))
            a = amp(10, L)
            print(f"  {gb:>6}GB{L:>8}"
                  f"{a['lev'][0]:>10.0f}/{a['lev'][1]:.0f}/{a['lev'][2]:.2f}"
                  f"{a['tier'][0]:>12.0f}/{a['tier'][1]:.0f}/{a['tier'][2]:.2f}")
        print("  leveled writes each byte ~31x to keep reads at 4 runs and space at")
        print("  1.1x. Size-tiered writes 4x and pays with 30 runs and 2x the disk.")
        print("  No third option wins both. That is the RUM conjecture.")
    return {"amp": amp}

def assembly(s):
    print("\nSeven blocks = a storage engine. Now measure what it costs.\n")
    for bpk, label in ((0, "no filter"), (4, "4 bits/key"), (10, "10 bits/key"), (16, "16 bits/key")):
        runs, keys = s["make_db"](30, 1500, bpk)
        use = bpk > 0
        absent = [f"missing{i}".encode() for i in range(1500)]
        present = [rng.choice(keys) for _ in range(1500)]
        ra = sum(s["get"](runs, k, use)[1]["block_reads"] for k in absent) / len(absent)
        rp = sum(s["get"](runs, k, use)[1]["block_reads"] for k in present) / len(present)
        if bpk == 0:
            print(f"  {'config':<13}{'absent reads':>14}{'present reads':>15}")
        print(f"  {label:<13}{ra:>14.3f}{rp:>15.2f}")
    print("\n  Absent keys: 30 reads -> 0.25. That is the Bloom filter earning its")
    print("  125 KB per 100k keys, and it is why LSM reads are viable at all.")
    print("  Present keys: 15.5 -> 1.1. The folklore says filters only help misses;")
    print("  measured, they help hits too, because a hit must SKIP the newer runs.")
    print("\n  Crash test -- the property that matters more than any throughput number:")
    w = s["WAL"](os.path.join(DIR, "wal2"))
    for i in range(200): w.append(f"key{i}".encode(), f"val{i}".encode(), sync=(i % 50 == 0))
    w.f.flush()
    with open(w.path, "r+b") as f:                        # simulate kill -9 mid-write
        f.seek(0, 2); size = f.tell(); f.truncate(size - 7)
    recs, left = w.replay()
    print(f"    wrote 200, truncated 7 bytes off the tail, recovered {len(recs)}")
    print(f"    every recovered record is intact: "
          f"{all(k.startswith(b'key') for k, _ in recs)}")
    print("\n  Built: record format -> WAL -> memtable/flush -> sparse index ->")
    print("  bloom -> read path -> amplification.")
    print("  Missing, on the project page: real SSTable blocks (m4), tombstones (m7),")
    print("  merging iterators (m8), both compaction strategies (m9/m10), and the")
    print("  crossover figure (E3) that is the project's headline deliverable.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P04 — LSM storage engine, block by block")