#!/usr/bin/env python3
"""
mofe_blackbox_decode.py — decode a MOFE "black-box" SCREEN RECORDING into a telemetry stream and merge it
with the coexisting MofePipelineRecorder ndjson.

The androidApp overlay (debug-gated) draws the engine's per-frame state — a BUFFER of raw IMU samples since
the previous frame (the ndjson does NOT log IMU), the fused pose, and a `tEngineMicros` sync key — as a
luma-only cell grid. A phone SCREEN RECORDING of a walk (captured WHILE the ndjson recorder runs) then decodes
here into per-frame telemetry, aligned to the ndjson's authoritative ranges by the shared engine clock. That
gives the #69b-i observer solve both the ranges (ndjson) and the IMU-odometry position (overlay) from a single
recording, with no on-device NDJSON change.

Wire format is byte-identical to :common `com.aether.mofe.diag.{BlackBoxCodec,BlackBoxSchema}` and locked by
the Kotlin golden `BlackBoxWireFormatTest`. Verified end-to-end through real H.264 (see --selftest).

Deps: numpy + the ffmpeg bundled by `pip install imageio-ffmpeg` (or any ffmpeg on PATH via --ffmpeg).
Usage:
  mofe_blackbox_decode.py CAPTURE.mp4 --ndjson TRACE.ndjson --auto --out records.jsonl
  mofe_blackbox_decode.py CAPTURE.mp4 --cols 80 --rows 130 --repeat 5 --cell 12 --x0 60 --y0 180
  mofe_blackbox_decode.py --selftest
"""
import argparse, json, os, shutil, struct, subprocess, sys, zlib
import numpy as np

MAGIC = b"\xA5\x3C"; VERSION = 1; HEADER = 5; TRAILER = 4; OVERHEAD = HEADER + TRAILER

# ── transport: bytes ⇄ framed bytes ⇄ luma grid (mirror of BlackBoxCodec.kt) ─────────────
def crc32(b): return zlib.crc32(b) & 0xFFFFFFFF

def frame(payload):
    body = bytes([VERSION]) + struct.pack(">H", len(payload)) + payload
    return MAGIC + body + struct.pack(">I", crc32(body))

def deframe(b):
    if len(b) < OVERHEAD or b[0:2] != MAGIC: return None
    ln = struct.unpack(">H", b[3:5])[0]
    if HEADER + ln + TRAILER > len(b): return None
    if crc32(b[2:HEADER + ln]) != struct.unpack(">I", b[HEADER + ln:HEADER + ln + 4])[0]: return None
    return bytes(b[HEADER:HEADER + ln])

def capacity_bytes(cols, rows, repeat=3): return max(0, (cols * rows) // repeat // 8 - OVERHEAD)

def encode_to_luma(payload, cols, rows, repeat=3):
    f = frame(payload); total = cols * rows; usable = total // repeat
    assert len(f) * 8 <= usable, f"payload too big: {len(f)*8} > {usable}"
    luma = np.zeros(total, np.uint8)
    for i in range(len(f) * 8):
        if (f[i // 8] >> (7 - (i % 8))) & 1:
            for c in range(repeat): luma[i + c * usable] = 255
    return luma

def decode_from_luma(luma, cols, rows, repeat=3):
    total = cols * rows
    if luma.size != total: return None
    usable = total // repeat
    thr = (int(luma.min()) + int(luma.max())) // 2
    bits = (luma[:usable * repeat].reshape(repeat, usable) > thr)
    maj = (bits.sum(axis=0) * 2 >= repeat)
    out = bytearray(usable // 8)
    for i in np.nonzero(maj[:len(out) * 8])[0]:
        out[i // 8] |= 1 << (7 - (i % 8))
    return deframe(out)

# ── schema: BlackBoxState ⇄ bytes (mirror of BlackBoxSchema.kt) ──────────────────────────
def encode_state(s):
    o = bytearray([VERSION]); o += struct.pack(">I", s["frameIdx"]); o += struct.pack(">Q", s["tEngineMicros"])
    o += bytes([1 if s.get("pose") else 0]); o += bytes([len(s["ranges"])])
    for r in s["ranges"]:
        o += struct.pack(">I", int(r["anchorTag"][-6:], 16))[1:]
        o += struct.pack(">H", r["distanceMm"] & 0xFFFF) + bytes([r["sq"] & 0xFF]) + struct.pack(">b", r["rssiDbm"])
        o += struct.pack(">H", r["azCentiDeg"] & 0xFFFF) + struct.pack(">h", r["elCentiDeg"])
    o += bytes([len(s["imu"])])
    for m in s["imu"]:
        o += struct.pack(">H", m["dtMicros"] & 0xFFFF) + struct.pack(">hhhhhh", m["axMilli"], m["ayMilli"], m["azMilli"], m["gxMilli"], m["gyMilli"], m["gzMilli"])
    if s.get("pose"):
        p = s["pose"]; o += struct.pack(">hhhhhhh", p["xMm"], p["yMm"], p["zMm"], p["qwE4"], p["qxE4"], p["qyE4"], p["qzE4"])
    return bytes(o)

def decode_state(b):
    try:
        p = 0
        def take(n):
            nonlocal p; v = b[p:p + n]; p += n
            assert len(v) == n; return v
        if take(1)[0] != VERSION: return None
        frameIdx = struct.unpack(">I", take(4))[0]; t = struct.unpack(">Q", take(8))[0]; flags = take(1)[0]
        ranges = []
        for _ in range(take(1)[0]):
            tag = struct.unpack(">I", b"\x00" + take(3))[0]
            dist, = struct.unpack(">H", take(2)); sq = take(1)[0]; rssi, = struct.unpack(">b", take(1))
            az, = struct.unpack(">H", take(2)); el, = struct.unpack(">h", take(2))
            ranges.append(dict(anchorTag=f"{tag:06x}", distanceMm=dist, sq=sq, rssiDbm=rssi, azCentiDeg=az, elCentiDeg=el))
        imu = []
        for _ in range(take(1)[0]):
            dt, = struct.unpack(">H", take(2)); v = struct.unpack(">hhhhhh", take(12))
            imu.append(dict(dtMicros=dt, axMilli=v[0], ayMilli=v[1], azMilli=v[2], gxMilli=v[3], gyMilli=v[4], gzMilli=v[5]))
        pose = None
        if flags & 1:
            v = struct.unpack(">hhhhhhh", take(14))
            pose = dict(xMm=v[0], yMm=v[1], zMm=v[2], qwE4=v[3], qxE4=v[4], qyE4=v[5], qzE4=v[6])
        return dict(frameIdx=frameIdx, tEngineMicros=t, ranges=ranges, imu=imu, pose=pose)
    except (AssertionError, struct.error, IndexError):
        return None

# ── ffmpeg glue ─────────────────────────────────────────────────────────────────────────
def ffmpeg_bin(explicit=None):
    if explicit: return explicit
    if shutil.which("ffmpeg"): return "ffmpeg"
    import imageio_ffmpeg; return imageio_ffmpeg.get_ffmpeg_exe()

def mp4_to_frames(path, ff=None):
    import re
    info = subprocess.run([ffmpeg_bin(ff), "-hide_banner", "-i", path], stderr=subprocess.PIPE).stderr.decode("utf8", "ignore")
    W = H = 0
    for line in info.splitlines():
        if "Video:" in line:  # avoid matching the "0x1" stream-id token on other lines
            for a, b in re.findall(r"(\d{2,5})x(\d{2,5})", line):
                if int(a) >= 16 and int(b) >= 16:
                    W, H = int(a), int(b); break
            if W: break
    if not (W and H):
        raise RuntimeError("could not parse video resolution from ffmpeg -i")
    buf = subprocess.run([ffmpeg_bin(ff), "-loglevel", "error", "-i", path, "-f", "rawvideo", "-pix_fmt", "gray",
                          "-vsync", "0", "-"], stdout=subprocess.PIPE).stdout
    n = len(buf) // (W * H)
    return np.frombuffer(buf[:n * W * H], np.uint8).reshape(n, H, W), W, H

# ── grid sampling ───────────────────────────────────────────────────────────────────────
def sample_grid(img, x0, y0, cols, rows, cellW, cellH):
    out = np.zeros(rows * cols, np.uint8)
    hw = max(1, int(cellW) // 4); hh = max(1, int(cellH) // 4)
    H, W = img.shape
    for r in range(rows):
        cy = min(H - 1, max(0, int(y0 + (r + 0.5) * cellH)))
        for c in range(cols):
            cx = min(W - 1, max(0, int(x0 + (c + 0.5) * cellW)))
            out[r * cols + c] = int(np.median(img[cy - hh:cy + hh + 1, cx - hw:cx + hw + 1]))
    return out

def overlay_geometry(W, H, cols, rows, margin_frac=0.04):
    """EXACT grid geometry for an UNCROPPED screen recording, from the overlay's known layout
    (BlackBoxOverlay.kt: margin = margin_frac*min(w,h); the cols×rows grid fills the rest with non-square
    cells). Ratio-based, so it stays exact even if the recorder downscales the whole frame. Cropping breaks
    it → use --fiducial."""
    m = margin_frac * min(W, H)
    return m, m, (W - 2 * m) / cols, (H - 2 * m) / rows

def detect_geometry(img, cols, rows):
    """Fallback auto-registration from the 3 corner fiducials (solid bright squares at the grid's outer
    TL/TR/BL corners, per BlackBoxOverlay.kt). Returns (x0, y0, cellW, cellH). Use for a FILMED/cropped
    screen (where overlay_geometry can't apply). Best-effort — refine against the first such capture."""
    H, W = img.shape
    bright = img > (int(img.max()) * 0.7)
    def blob_corner(ys, xs, want):
        pts = np.argwhere(bright[ys[0]:ys[1], xs[0]:xs[1]])
        if len(pts) < 20: raise RuntimeError("fiducial not found; pass explicit --x0/--y0/--cell")
        pts = pts + [ys[0], xs[0]]; yy, xx = pts[:, 0], pts[:, 1]
        return (yy.min() if want[0] == 'min' else yy.max(), xx.min() if want[1] == 'min' else xx.max())
    tl = blob_corner((0, H // 2), (0, W // 2), ('max', 'max'))   # inner (bottom-right) corner of TL fiducial
    tr = blob_corner((0, H // 2), (W // 2, W), ('max', 'min'))   # inner (bottom-left) corner of TR fiducial
    bl = blob_corner((H // 2, H), (0, W // 2), ('min', 'max'))   # inner (top-right) corner of BL fiducial
    x0 = (tl[1] + bl[1]) / 2; y0 = (tl[0] + tr[0]) / 2
    return x0, y0, (tr[1] - x0) / cols, (bl[0] - y0) / rows

# ── ndjson merge (align overlay IMU to authoritative ndjson ranges by tEngineMicros) ──────
def parse_ndjson_ranges(path):
    by_t = {}
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if not line: continue
            try: r = json.loads(line)
            except json.JSONDecodeError: continue
            if r.get("ev") == "raw": by_t.setdefault(r["t"], []).append(r)
    return sorted(by_t.items())

def merge(states, nd, tol_us=100_000):
    ts = [t for t, _ in nd]; out = []
    for s in states:
        if s is None: continue
        t = s["tEngineMicros"]; best = None
        if ts:
            j = int(np.searchsorted(ts, t))
            for k in (j - 1, j):
                if 0 <= k < len(ts) and (best is None or abs(ts[k] - t) < abs(best[0] - t)): best = (ts[k], nd[k][1])
        rng = None
        if best and abs(best[0] - t) <= tol_us:
            rng = [dict(anc=r["anc"], d=r["d"], sq=r.get("sq"), rssi=r.get("rssi")) for r in best[1]]
        out.append(dict(t=t, frameIdx=s["frameIdx"], imu=s["imu"], pose=s["pose"],
                        ndjson_t=(best[0] if best else None), ranges=rng))
    return out

# ── selftest: real H.264 end-to-end ───────────────────────────────────────────────────────
def selftest(ff=None):
    import tempfile
    COLS, ROWS, REP, W, H = 80, 130, 5, 1080, 2400  # phone-portrait; overlay layout + resolution-exact decode
    def synth(n):
        tags = ["264b1b", "a6080f", "85b92d", "c77cc1"]; t = 1_700_000_000_000; out = []
        for k in range(n):
            t += 33_000
            out.append(dict(frameIdx=k, tEngineMicros=t,
                ranges=[dict(anchorTag=tags[i], distanceMm=1000 + 400 * i + k, sq=200 - 10 * i, rssiDbm=-70 - 5 * i,
                             azCentiDeg=(4000 + 900 * i + k) % 36000, elCentiDeg=-300 + 50 * i) for i in range(4)],
                imu=[dict(dtMicros=5000, axMilli=100 + j - k, ayMilli=-20 + j, azMilli=9810, gxMilli=3 * j, gyMilli=-2 * j, gzMilli=j + k) for j in range(7)],
                pose=dict(xMm=800 + k, yMm=-1500 + 2 * k, zMm=1100, qwE4=9990, qxE4=100 - k, qyE4=-40, qzE4=k)))
        return out
    assert decode_state(encode_state(synth(1)[0])) == synth(1)[0], "schema round-trip"
    states = synth(30)
    # render each frame EXACTLY as BlackBoxOverlay.kt does: black bg, margin=0.04*min, non-square cells, 3 fiducials
    gx0, gy0, gcw, gch = overlay_geometry(W, H, COLS, ROWS)
    mm = int(gx0); fid = int(gx0 * 0.72); gap = int(gx0 * 0.14)
    frames = []
    for s in states:
        luma = encode_to_luma(encode_state(s), COLS, ROWS, REP).reshape(ROWS, COLS)
        img = np.zeros((H, W), np.uint8)
        for r in range(ROWS):
            y0i, y1i = int(gy0 + r * gch), int(gy0 + (r + 1) * gch)
            row = luma[r]
            for c in range(COLS):
                if row[c]:
                    img[y0i:y1i, int(gx0 + c * gcw):int(gx0 + (c + 1) * gcw)] = 255
        img[mm - gap - fid:mm - gap, mm - gap - fid:mm - gap] = 255                 # TL fiducial
        img[mm - gap - fid:mm - gap, W - mm + gap:W - mm + gap + fid] = 255         # TR
        img[H - mm + gap:H - mm + gap + fid, mm - gap - fid:mm - gap] = 255         # BL
        frames.append(img)
    rc = 1
    for crf in (23, 28):
        with tempfile.TemporaryDirectory() as d:
            mp4 = os.path.join(d, "c.mp4")
            p = subprocess.Popen([ffmpeg_bin(ff), "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "gray",
                                  "-s", f"{W}x{H}", "-r", "30", "-i", "-", "-c:v", "libx264", "-pix_fmt", "yuv420p",
                                  "-crf", str(crf), mp4], stdin=subprocess.PIPE)
            for f in frames: p.stdin.write(f.tobytes())
            p.stdin.close(); p.wait()
            dec, dW, dH = mp4_to_frames(mp4, ff)
            ax0, ay0, acw, ach = overlay_geometry(dW, dH, COLS, ROWS)  # resolution-exact auto-registration
            ok = wrong = 0
            for i in range(min(len(dec), len(states))):
                st = decode_state(decode_from_luma(sample_grid(dec[i], ax0, ay0, COLS, ROWS, acw, ach), COLS, ROWS, REP))
                if st == states[i]: ok += 1
                elif st is not None: wrong += 1
            print(f"  H.264 crf={crf}: {ok}/{len(states)} exact, wrong={wrong} (overlay layout + resolution-exact decode)")
            if wrong != 0 or (crf == 23 and ok != len(states)) or (crf == 28 and ok < 0.9 * len(states)): rc = 0
    nd = [(s["tEngineMicros"] + 1200, [dict(anc=r["anchorTag"], d=r["distanceMm"] / 1000.0) for r in s["ranges"]]) for s in states[:5]]
    mg = merge(states[:5], nd)
    assert all(m["ranges"] and len(m["imu"]) == 7 for m in mg), "merge lost data"
    print(f"  ndjson merge: OK ({len(mg)} frames)")
    print("SELFTEST PASS" if rc else "SELFTEST FAIL"); return rc

def main():
    ap = argparse.ArgumentParser(description="Decode a MOFE black-box screen recording + merge with ndjson.")
    ap.add_argument("mp4", nargs="?"); ap.add_argument("--ndjson"); ap.add_argument("--out")
    ap.add_argument("--cols", type=int, default=80); ap.add_argument("--rows", type=int, default=130)
    ap.add_argument("--repeat", type=int, default=5); ap.add_argument("--cell", type=float)
    ap.add_argument("--x0", type=int); ap.add_argument("--y0", type=int)
    ap.add_argument("--fiducial", action="store_true", help="register grid from corner fiducials (for a FILMED/cropped screen); default = exact resolution geometry")
    ap.add_argument("--ffmpeg"); ap.add_argument("--selftest", action="store_true")
    a = ap.parse_args()
    if a.selftest: sys.exit(0 if selftest(a.ffmpeg) else 1)
    if not a.mp4: ap.error("need CAPTURE.mp4 (or --selftest)")
    frames, W, H = mp4_to_frames(a.mp4, a.ffmpeg)
    if a.x0 is not None and a.y0 is not None and a.cell is not None:
        x0, y0, cw, ch = a.x0, a.y0, a.cell, a.cell
    elif a.fiducial:
        x0, y0, cw, ch = detect_geometry(frames[len(frames) // 2], a.cols, a.rows)
        sys.stderr.write(f"fiducial geometry: x0={x0:.1f} y0={y0:.1f} cellW={cw:.2f} cellH={ch:.2f}\n")
    else:
        x0, y0, cw, ch = overlay_geometry(W, H, a.cols, a.rows)   # exact for an uncropped screen recording
        sys.stderr.write(f"resolution geometry ({W}x{H}): x0={x0:.1f} y0={y0:.1f} cellW={cw:.2f} cellH={ch:.2f}\n")
    states = [decode_state(decode_from_luma(sample_grid(f, x0, y0, a.cols, a.rows, cw, ch), a.cols, a.rows, a.repeat)) for f in frames]
    good = [s for s in states if s is not None]
    sys.stderr.write(f"decoded {len(good)}/{len(frames)} frames\n")
    recs = merge(good, parse_ndjson_ranges(a.ndjson)) if a.ndjson else [dict(t=s["tEngineMicros"], frameIdx=s["frameIdx"], imu=s["imu"], pose=s["pose"]) for s in good]
    sink = open(a.out, "w") if a.out else sys.stdout
    for r in recs: sink.write(json.dumps(r) + "\n")
    if a.out: sink.close(); sys.stderr.write(f"wrote {len(recs)} records → {a.out}\n")

if __name__ == "__main__":
    main()