#!/usr/bin/env python3
"""anchor-check: host-side check of a COPIED anchor-recorder session (runs on the Mac, not the board).

    tools/anchor-check <session_dir> [...]          decode + counts + content screen, write
                                                    SHA256SUMS and check.json into the session
    tools/anchor-check --verify-copy <session_dir>  re-hash files against SHA256SUMS (and sizes
                                                    recorded in check.json) after another copy

Needs Python 3.9+, ffmpeg and ffprobe. Stdlib only. Exit 0 = no failures, 1 = failures, 2 = usage.

What it checks, per chunk and eye:
  a) ffmpeg -v error decode of every L/R segment: exit code, stderr errors, decoded frame count.
  b) ffprobe packet count per eye vs frames.csv rows. The last segment of an INTERRUPTED session
     may be truncated (power removed), so its shortfall is a warning; anywhere else it is a failure.
  c) 1 fps luma screen (64x48 gray): left == right (identical eyes), repeated samples per eye
     (frozen eye), mean luma below 16 ("black: lens cap?"). A static scene can repeat without a
     fault and a live sensor can be black in the dark: these are flags to look at, not verdicts.
  d) SHA256SUMS (sha256sum format, paths relative to the session) and check.json.
Finding codes: DECODE_ERROR, NO_VIDEO, COUNT_MISMATCH, LR_IDENTICAL, FROZEN_EYE, BLACK_FRAMES,
MANIFEST_MISSING; copy verification: SIZE_MISMATCH, HASH_MISMATCH, MISSING_FILE, PATH_ESCAPE,
SUMS_MISSING (same codes as Kihyun's validate_session where they overlap).
"""
import argparse, hashlib, json, os, subprocess, sys, tempfile, time
from pathlib import Path

VERSION = "1"
SAMPLE_W, SAMPLE_H = 64, 48
BLACK_LUMA = 16            # full-range video: a capped lens reads ~0-10
SKIP = {"SHA256SUMS", "check.json"}


def finding(out, code, severity, where, detail):
    out.append({"code": code, "severity": severity, "where": where, "detail": detail})


def packets(seg):
    r = subprocess.run(["ffprobe", "-v", "error", "-count_packets", "-select_streams", "v:0", "-show_entries",
                        "stream=nb_read_packets", "-of", "csv=p=0", str(seg)], capture_output=True, text=True)
    digits = "".join(ch for ch in r.stdout.split("\n")[0] if ch.isdigit())
    return int(digits) if digits else 0


def decode(seg):
    """One decode pass: error lines, exit code, decoded frame count, 1 fps gray samples."""
    with tempfile.NamedTemporaryFile(suffix=".crc", delete=False) as t:
        crc = Path(t.name)
    fc = f"[0:v]split[a][b];[b]fps=1,scale={SAMPLE_W}:{SAMPLE_H},format=gray[s]"
    p = subprocess.run(["ffmpeg", "-nostdin", "-hide_banner", "-v", "error", "-copyts", "-i", str(seg),
                        "-filter_complex", fc, "-map", "[a]", "-fps_mode", "passthrough", "-f", "framecrc", "-y",
                        str(crc), "-map", "[s]", "-f", "rawvideo", "pipe:1"], capture_output=True)
    try:
        decoded = sum(1 for ln in crc.read_text().splitlines() if ln.startswith("0,"))
    finally:
        crc.unlink(missing_ok=True)
    n = SAMPLE_W * SAMPLE_H
    samples = [p.stdout[i:i + n] for i in range(0, len(p.stdout) - n + 1, n)]
    errors = [ln for ln in p.stderr.decode("utf-8", "replace").splitlines() if ln.strip()]
    return {"exit": p.returncode, "errors": errors[:20], "error_lines": len(errors), "decoded": decoded}, samples


def frames_rows(chunk):
    f = chunk / "meta" / "frames.csv"
    if not f.exists():
        return 0
    return sum(1 for ln in f.read_text(errors="replace").splitlines()[1:] if len(ln.split(",")) >= 11)


def luma(samples):
    means = [sum(s) / len(s) for s in samples]
    rep = sum(1 for a, b in zip(samples, samples[1:]) if a == b)
    return {"samples": len(samples), "mean_min": round(min(means), 1) if means else None,
            "mean_avg": round(sum(means) / len(means), 1) if means else None,
            "black_samples": sum(1 for m in means if m < BLACK_LUMA), "repeated_samples": rep}


def check_chunk(chunk, last_chunk, interrupted, out):
    res = {"chunk": chunk.name, "frames_csv_rows": frames_rows(chunk), "eyes": {}}
    eye_samples = {}
    for eye in ("L", "R"):
        segs = sorted((chunk / eye).glob("[0-9]*.ts"))
        e = {"segments": [], "packets": 0, "decoded": 0}
        eye_samples[eye] = []
        if not segs:
            finding(out, "NO_VIDEO", "fail", f"{chunk.name}/{eye}", "no .ts segments")
        for i, seg in enumerate(segs):
            d, samples = decode(seg)
            d.update({"file": seg.name, "bytes": seg.stat().st_size, "packets": packets(seg)})
            d["luma"] = luma(samples)
            eye_samples[eye] += samples
            e["segments"].append(d)
            e["packets"] += d["packets"]; e["decoded"] += d["decoded"]
            tail = last_chunk and i == len(segs) - 1
            if d["exit"] != 0 or d["error_lines"] or d["decoded"] < d["packets"]:
                sev = "warn" if tail and interrupted else "fail"
                finding(out, "DECODE_ERROR", sev, f"{chunk.name}/{eye}/{seg.name}",
                        f"exit {d['exit']}, {d['error_lines']} error lines, decoded {d['decoded']} of "
                        f"{d['packets']} packets" + (" (last segment of an interrupted session)" if sev == "warn" else ""))
        e["luma"] = luma(eye_samples[eye])
        res["eyes"][eye] = e
    L, R, rows = res["eyes"]["L"]["packets"], res["eyes"]["R"]["packets"], res["frames_csv_rows"]
    if not (L == R == rows):
        # A power cut can only shorten the END: tolerate it when every earlier segment pair agrees.
        ls_, rs_ = res["eyes"]["L"]["segments"], res["eyes"]["R"]["segments"]
        earlier_ok = len(ls_) == len(rs_) and all(a["packets"] == b["packets"] and a["decoded"] == a["packets"]
                                                  and b["decoded"] == b["packets"] for a, b in zip(ls_[:-1], rs_[:-1]))
        sev = "warn" if last_chunk and interrupted and earlier_ok else "fail"
        finding(out, "COUNT_MISMATCH", sev, chunk.name,
                f"TS packets L {L} / R {R}, decoded L {res['eyes']['L']['decoded']} / R {res['eyes']['R']['decoded']}, "
                f"frames.csv {rows}" + (" (power-cut tail)" if sev == "warn" else ""))
    ls, rs = eye_samples["L"], eye_samples["R"]
    same = sum(1 for a, b in zip(ls, rs) if a == b)
    res["lr_identical_samples"] = same
    if ls and same and same >= max(1, len(ls) // 10):
        finding(out, "LR_IDENTICAL", "fail", chunk.name,
                f"{same} of {min(len(ls), len(rs))} 1 fps samples identical in left and right (one eye copied?)")
    for eye, smp in (("L", ls), ("R", rs)):
        lu = res["eyes"][eye]["luma"]
        if lu["samples"] >= 3 and lu["black_samples"] >= lu["samples"] // 2:
            finding(out, "BLACK_FRAMES", "fail", f"{chunk.name}/{eye}",
                    f"black: lens cap? {lu['black_samples']} of {lu['samples']} samples below luma {BLACK_LUMA}")
        elif lu["samples"] >= 3 and lu["repeated_samples"] >= 0.9 * (lu["samples"] - 1):
            finding(out, "FROZEN_EYE", "fail", f"{chunk.name}/{eye}",
                    f"{lu['repeated_samples']} of {lu['samples'] - 1} consecutive 1 fps samples identical")
    return res


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""):
            h.update(b)
    return h.hexdigest()


def session_files(s):
    return sorted(p for p in s.rglob("*") if p.is_file() and p.relative_to(s).as_posix() not in SKIP
                  and not p.name.endswith(".tmp"))


def check(session):
    s = Path(session)
    out, t0 = [], time.time()
    man = {}
    if (s / "manifest.json").exists():
        man = json.loads((s / "manifest.json").read_text())
    else:
        finding(out, "MANIFEST_MISSING", "fail", s.name, "no manifest.json")
    interrupted = man.get("status") == "interrupted"
    chunks = sorted(p for p in s.glob("c[0-9][0-9][0-9]") if p.is_dir())
    res = {"tool": "anchor-check", "version": VERSION, "session": s.name, "format": man.get("format"),
           "status": man.get("status"), "checked_host_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
           "chunks": [check_chunk(c, c == chunks[-1], interrupted, out) for c in chunks]}
    files = session_files(s)
    sums = [(sha256(p), p.relative_to(s).as_posix(), p.stat().st_size) for p in files]
    (s / "SHA256SUMS").write_text("".join(f"{h}  {rel}\n" for h, rel, _ in sums))
    res["files"] = {rel: size for _, rel, size in sums}
    res["findings"] = out
    res["pass"] = not any(f["severity"] == "fail" for f in out)
    res["seconds"] = round(time.time() - t0, 1)
    tmp = s / "check.json.tmp"
    tmp.write_text(json.dumps(res, indent=1)); os.replace(tmp, s / "check.json")
    return res


def verify_copy(session):
    s = Path(session).resolve()
    out = []
    sums = s / "SHA256SUMS"
    if not sums.exists():
        finding(out, "SUMS_MISSING", "fail", s.name, "run anchor-check on the source copy first")
        return {"session": s.name, "findings": out, "pass": False}
    sizes = {}
    if (s / "check.json").exists():
        sizes = json.loads((s / "check.json").read_text()).get("files", {})
    listed = set()
    for ln in sums.read_text().splitlines():
        if not ln.strip():
            continue
        h, _, rel = ln.partition("  ")
        listed.add(rel)
        p = (s / rel)
        if os.path.isabs(rel) or ".." in Path(rel).parts or not p.resolve().is_relative_to(s):
            finding(out, "PATH_ESCAPE", "fail", rel, "path leaves the session directory")
            continue
        if not p.is_file():
            finding(out, "MISSING_FILE", "fail", rel, "listed in SHA256SUMS, not present")
            continue
        if rel in sizes and p.stat().st_size != sizes[rel]:
            finding(out, "SIZE_MISMATCH", "fail", rel, f"{p.stat().st_size} bytes, expected {sizes[rel]}")
            continue
        if sha256(p) != h:
            finding(out, "HASH_MISMATCH", "fail", rel, "content differs from SHA256SUMS")
    extra = [p.relative_to(s).as_posix() for p in session_files(s) if p.relative_to(s).as_posix() not in listed]
    for rel in extra:
        finding(out, "EXTRA_FILE", "warn", rel, "not in SHA256SUMS (written after the check?)")
    return {"session": s.name, "files_listed": len(listed), "findings": out,
            "pass": not any(f["severity"] == "fail" for f in out)}


def main(argv=None):
    ap = argparse.ArgumentParser(prog="anchor-check", description=__doc__.split("\n")[0])
    ap.add_argument("sessions", nargs="+")
    ap.add_argument("--verify-copy", action="store_true", help="only re-check files against SHA256SUMS")
    ap.add_argument("--json", action="store_true", help="print the full result as JSON")
    a = ap.parse_args(argv)
    ok = True
    for sess in a.sessions:
        if not Path(sess).is_dir():
            print(f"anchor-check: {sess}: not a directory", file=sys.stderr); return 2
        r = verify_copy(sess) if a.verify_copy else check(sess)
        ok &= r["pass"]
        if a.json:
            print(json.dumps(r, indent=1)); continue
        head = f"{r['session']}: {'PASS' if r['pass'] else 'FAIL'}"
        if not a.verify_copy:
            eyes = "; ".join(f"{c['chunk']} L {c['eyes']['L']['decoded']}/{c['eyes']['L']['packets']} "
                             f"R {c['eyes']['R']['decoded']}/{c['eyes']['R']['packets']} "
                             f"meta {c['frames_csv_rows']}" for c in r["chunks"])
            head += f"  (decoded/packets: {eyes}; status {r['status']}; {r['seconds']} s)"
        else:
            head += f"  ({r.get('files_listed', 0)} files)"
        print(head)
        for f in r["findings"]:
            print(f"  {f['severity'].upper():4} {f['code']:15} {f['where']}: {f['detail']}")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())
