#!/usr/bin/env python3
"""Hands-on P03 — a vector database, assembled from six lego blocks."""
import json, math, os, random, struct, tempfile, zlib
from _harness import block, run_all
rng = random.Random(0)
DIR = tempfile.mkdtemp(prefix="h03-")

@block(1, "Storage format, written before the code", "byte offsets are a design decision, not an implementation detail")
def b1(s, show):
    D = 8
    def enc(vid, vec, meta):
        m = json.dumps(meta, sort_keys=True).encode()
        body = struct.pack("<IHH", vid, D, len(m)) + struct.pack(f"<{D}f", *vec) + m
        return struct.pack("<I", zlib.crc32(body)) + struct.pack("<I", len(body)) + body
    def dec(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
        body = buf[off+8:off+8+blen]
        if zlib.crc32(body) != crc: return None, off
        vid, d, ml = struct.unpack_from("<IHH", body, 0)
        vec = list(struct.unpack_from(f"<{d}f", body, 8))
        meta = json.loads(body[8+4*d:8+4*d+ml])
        return (vid, vec, meta), off + 8 + blen
    r = enc(7, [0.1]*D, {"topic": "tech", "ts": 100})
    got, _ = dec(r, 0)
    if show:
        print(f"  layout: crc(4) len(4) id(4) dim(2) metalen(2) vec(4*d) meta(json)")
        print(f"  one record = {len(r)} bytes for d={D}")
        print(f"  round-trip id={got[0]} meta={got[2]}")
        print("  writing this table BEFORE the code is what stops the format drifting")
    return {"enc": enc, "dec": dec, "D": D}

@block(2, "Append-only log + full rebuild", "the naive design, measured -- this is what motivates P04")
def b2(s, show):
    class NaiveDB:
        def __init__(self, path):
            self.path = path; self.f = open(path, "ab"); self.idx = {}
        def insert(self, vid, vec, meta):
            self.f.write(s["enc"](vid, vec, meta)); self.idx[vid] = (vec, meta)
        def flush(self): self.f.flush(); os.fsync(self.f.fileno())
        def rebuild(self):
            buf = open(self.path, "rb").read(); off = 0; n = 0; self.idx = {}
            while off < len(buf):
                rec, off2 = s["dec"](buf, off)
                if rec is None: break
                self.idx[rec[0]] = (rec[1], rec[2]); off = off2; n += 1
            return n
    import time
    if show:
        print(f"  {'vectors':>9}{'rebuild ms':>13}{'per vector':>13}")
        for n in (2000, 8000, 32000):
            db = NaiveDB(os.path.join(DIR, f"n{n}"))
            for i in range(n):
                db.insert(i, [rng.random() for _ in range(s["D"])], {"topic": i % 8})
            db.flush()
            t0 = time.perf_counter(); cnt = db.rebuild(); ms = (time.perf_counter()-t0)*1e3
            print(f"  {n:>9}{ms:>13.1f}{ms/n*1000:>12.1f}us")
        print("  LINEAR in the corpus, with a large constant. Extrapolate to 1M")
        print("  vectors and startup is minutes. THAT number is why P04 exists.")
    return {"NaiveDB": NaiveDB}

@block(3, "Snapshot + WAL replay", "recovery time stops being a function of all history")
def b3(s, show):
    import pickle, time
    def snapshot(idx, path):
        with open(path, "wb") as f: pickle.dump(idx, f)
    def recover(snap_path, wal_path):
        idx = pickle.load(open(snap_path, "rb")) if os.path.exists(snap_path) else {}
        if os.path.exists(wal_path):
            buf = open(wal_path, "rb").read(); off = 0
            while off < len(buf):
                rec, off2 = s["dec"](buf, off)
                if rec is None: break
                idx[rec[0]] = (rec[1], rec[2]); off = off2
        return idx
    n = 32000
    db = s["NaiveDB"](os.path.join(DIR, "snapbase"))
    for i in range(n):
        db.insert(i, [rng.random() for _ in range(s["D"])], {"topic": i % 8})
    db.flush()
    snapshot(db.idx, os.path.join(DIR, "snap"))
    tail = s["NaiveDB"](os.path.join(DIR, "tailwal"))
    for i in range(n, n + 500):
        tail.insert(i, [rng.random() for _ in range(s["D"])], {"topic": 0})
    tail.flush()
    t0 = time.perf_counter(); full = db.rebuild(); t_full = (time.perf_counter()-t0)*1e3
    t0 = time.perf_counter()
    idx = recover(os.path.join(DIR, "snap"), os.path.join(DIR, "tailwal"))
    t_snap = (time.perf_counter()-t0)*1e3
    if show:
        print(f"  full rebuild of {n:,}:            {t_full:8.1f} ms")
        print(f"  snapshot + {500} WAL records: {t_snap:8.1f} ms   "
              f"({t_full/t_snap:.0f}x faster)")
        print(f"  recovered {len(idx):,} vectors")
        print("  recovery is now a function of the WAL TAIL, not of all history.")
    return {"recover": recover}

@block(4, "Metadata filtering, three ways", "the crossover is arithmetic, not taste")
def b4(s, show):
    def dot(a, b): return sum(x*y for x, y in zip(a, b))
    def brute_filtered(idx, q, pred, k=10):
        cands = [(vid, v) for vid, (v, m) in idx.items() if pred(m)]
        return sorted(cands, key=lambda t: -dot(t[1], q))[:k], len(cands)
    def post_filter(idx, q, pred, k=10, K=None):
        K = K or k * 10
        top = sorted(idx.items(), key=lambda t: -dot(t[1][0], q))[:K]
        kept = [(vid, v) for vid, (v, m) in top if pred(m)]
        return kept[:k], K
    if show:
        print("  post-filter needs K >= k/s candidates. Solve for the crossover:")
        HOP, DIST, N, k = 899.0, 13.3, 1_000_000, 10
        xo = math.sqrt(k * HOP / (N * DIST))
        print(f"    (k/s)*{HOP:.0f}ns  ==  s*N*{DIST}ns   ->   s = {xo:.4f}")
        print(f"  {'selectivity':>12}{'K needed':>11}{'post ms':>10}{'brute ms':>10}{'winner':>14}")
        for sel in (0.5, 0.1, 0.02, 0.001):
            K = math.ceil(k/sel); post = K*HOP/1e6; br = sel*N*DIST/1e6
            print(f"  {sel:>12}{K:>11,}{post:>10.2f}{br:>10.2f}"
                  f"{'post-filter' if post<br else 'BRUTE (exact)':>14}")
        print("  below ~2.6% selectivity an EXACT scan beats the approximate index.")
        print("  A planner without that third option falls off a cliff right here.")
    return {"brute_filtered": brute_filtered, "post_filter": post_filter, "dot": dot}

@block(5, "Tombstones", "you cannot delete from an immutable file; you write a marker")
def b5(s, show):
    if show:
        print(f"  {'churn':>7}{'live':>8}{'on disk':>9}{'space amp':>11}{'eff. ef=128':>13}")
        for churn in (0.0, 0.1, 0.3, 0.5):
            live, tomb = 10000, int(10000*churn)
            print(f"  {churn:>6.0%}{live-tomb:>8,}{live:>9,}"
                  f"{live/max(live-tomb,1):>11.2f}{128*(1-churn):>13.0f}")
        print("  deleted nodes stay in the graph as routing waypoints, so they occupy")
        print("  beam slots without producing results: recall drifts down with churn,")
        print("  and only compaction (a graph REBUILD) restores it.")
    return {}

@block(6, "Crash consistency", "the only property that distinguishes a database from an index")
def b6(s, show):
    def crash_test(n_writes, kill_at):
        path = os.path.join(DIR, f"crash{kill_at}")
        if os.path.exists(path): os.remove(path)
        db = s["NaiveDB"](path)
        acked = []
        for i in range(n_writes):
            db.insert(i, [float(i)]*s["D"], {"i": i})
            if i % 10 == 0:
                db.flush(); acked = list(range(i+1))       # only these are durable
            if i == kill_at: break
        db.f.flush()
        with open(path, "r+b") as f:                        # truncate mid-record
            sz = f.seek(0, 2); f.truncate(max(0, sz - 5))
        recovered = s["recover"]("/nonexistent", path)
        return acked, recovered
    if show:
        ok = True
        for kill in (37, 88, 155):
            acked, rec = crash_test(200, kill)
            lost = [i for i in acked if i not in rec]
            ok &= not lost
            print(f"  kill at write {kill:>4}: acked {len(acked):>4}, "
                  f"recovered {len(rec):>4}, ACKED LOST: {len(lost)}")
        print(f"  no acknowledged write lost at any kill point: {ok}")
        print("  automate this at 20+ random points. It is the highest-value test")
        print("  in the project and it finds bugs nothing else does.")
    return {}

def assembly(s):
    print("\nSix blocks = a vector database. The end-to-end query, with a planner.\n")
    idx = {}
    for i in range(4000):
        idx[i] = ([rng.random() for _ in range(s["D"])],
                  {"topic": i % 20, "ts": i, "pub": f"pub{i % 300}"})
    q = [rng.random() for _ in range(s["D"])]
    print(f"  {'filter':<22}{'matches':>9}{'strategy':>16}{'top-1 id':>10}")
    for name, pred, sel in (
            ("topic in 0..9",        lambda m: m["topic"] < 10,            0.50),
            ("topic == 3",           lambda m: m["topic"] == 3,            0.05),
            ("pub == pub7",          lambda m: m["pub"] == "pub7",         0.003),
            ("pub7 AND ts > 3800",   lambda m: m["pub"]=="pub7" and m["ts"]>3800, 0.0005)):
        nmatch = sum(1 for _, (v, m) in idx.items() if pred(m))
        actual_sel = nmatch / len(idx)
        strategy = "brute (exact)" if actual_sel < 0.026 else "post-filter"
        if strategy.startswith("brute"):
            res, _ = s["brute_filtered"](idx, q, pred)
        else:
            res, _ = s["post_filter"](idx, q, pred, K=math.ceil(10/max(actual_sel,1e-9)))
        print(f"  {name:<22}{nmatch:>9}{strategy:>16}{res[0][0] if res else '-':>10}")
    print("\n  The planner chose per query, from measured selectivity. It did not")
    print("  guess, and it did not use one strategy for everything.")
    print("\n  Recovery, the number that motivates the next project:")
    print("    naive rebuild is LINEAR in all history ever written")
    print("    snapshot + WAL tail is linear in the tail only -- ~100x here")
    print("\n  Built: record format -> append log -> snapshot/replay -> filtering")
    print("  planner -> tombstones -> crash consistency.")
    print("  Missing, on the project page: mmap segments (m7), compaction (m10),")
    print("  snapshot-isolated readers (m11), and E12 -- p99 DURING compaction,")
    print("  which is the only p99 that exists in production.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P03 — Vector database, block by block")