"""Read-only audit of the frozen A/B/C payloads; writes statistics, not training data.

Uses the existing tokenizer and admission metadata, never reads gold/tests or runs
the factory. Temporary SQLite indexes keep the five tokenized files off the heap.
"""
from __future__ import annotations

import argparse
from array import array
from collections import Counter
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import re
import sqlite3
import sys
import tempfile
import zlib


def digest(path):
    with path.open("rb") as handle:
        return hashlib.file_digest(handle, "sha256").hexdigest()


def quantiles(values):
    values = sorted(values)
    if not values:
        return {"n": 0}
    def q(f):
        x = (len(values) - 1) * f
        i = int(x)
        return round(values[i] + (values[min(i + 1, len(values) - 1)] - values[i]) * (x - i), 2)
    return dict(n=len(values), total=sum(values), min=values[0], p25=q(.25), p50=q(.5), p75=q(.75), p90=q(.9), max=values[-1])


def read_lines(path):
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                yield json.loads(line)


def parse_sequence(ids, mask, tokenizer, sr):
    """Inspect serialized assistant/observation blocks, without executing content."""
    start_id = tokenizer.token_to_id("<|im_start|>")
    starts = [i for i, token in enumerate(ids) if token == start_id]
    bounds = starts + [len(ids)]
    blocks = []
    for left, right in zip(bounds, bounds[1:]):
        text = tokenizer.decode(ids[left:right], skip_special_tokens=False)
        match = re.match(r"<\|im_start\|>(\w+)\n(.*)", text, re.S)
        if not match:
            continue
        role, body = match.groups()
        blocks.append(dict(start=left, end=right, role=role, body=body,
                           target=sum(mask[left:right])))
    events, thoughts, final = [], [], None
    summary = Counter()
    spans = []
    for bi, block in enumerate(blocks):
        if block["role"] != "assistant":
            summary["nonassistant_target_tokens"] += block["target"]
            continue
        body, target = block["body"], block["target"] > 0
        summary["assistant_blocks"] += 1
        summary["supervised_blocks"] += target
        thought = re.search(r"<think>\n?(.*?)</think>", body, re.S)
        thought_text = thought.group(1) if thought else ""
        thought_ids = tokenizer.encode(thought_text, add_special_tokens=False).ids
        if target:
            thoughts.append(len(thought_ids))
            summary["empty_think_blocks"] += not thought_text.strip()
            if len(thought_text) >= 3000:
                ratio = len(zlib.compress(thought_text.encode())) / len(thought_text.encode())
                summary["low_deflate_long_thoughts"] += ratio < .1
        calls = []
        for name, content in re.findall(r"<function=([^>]+)>(.*?)</function>", body, re.S):
            args = {k: v.strip() for k, v in re.findall(r"<parameter=([^>]+)>\n?(.*?)\n?</parameter>", content, re.S)}
            command = args.get("command", "") if name == "bash" else ""
            obs = ""
            if bi + 1 < len(blocks) and blocks[bi + 1]["role"] == "user" and "<tool_response>" in blocks[bi + 1]["body"]:
                obs = sr.strip_ansi(blocks[bi + 1]["body"])
            test = sr.parse_test_result(command, obs) if name == "bash" and sr.is_testy(command) else None
            test = test.as_dict() if test else None
            event = dict(block=bi, name=name, command=command, args=args, target=target,
                         observation=obs, test=test,
                         mutation=name == "edit" or (name == "bash" and sr.is_mutation(command)))
            calls.append(event)
            events.append(event)
        tail = re.sub(r"<think>.*?</think>", "", body, flags=re.S)
        tail = re.sub(r"<\|im_end\|>", "", tail).strip()
        if not calls and tail:
            final = dict(block=bi, target=target, text=tail, tokens=block["target"])
        spans.append(dict(block=bi, start=block["start"], end=block["end"],
                          target=block["target"], calls=len(calls), thought_tokens=len(thought_ids)))
    last_mut = max((e["block"] for e in events if e["mutation"]), default=-1)
    first_mut = min((e["block"] for e in events if e["mutation"]), default=None)
    tests = [e for e in events if e["test"] and e["test"].get("test_ran")]
    passes = [e for e in tests if e["test"]["state"] == "pass"]
    failures = [e for e in tests if e["test"]["state"] == "fail"]
    last_pass = max((e["block"] for e in passes), default=-1)
    summary["has_parsed_pass"] = bool(passes)
    summary["has_parsed_fail"] = bool(failures)
    summary["parsed_pass_after_last_mutation"] = last_pass >= 0 and last_pass >= last_mut
    summary["final_present"] = final is not None
    summary["final_supervised"] = final is not None and final["target"]
    summary["final_masked"] = final is not None and not final["target"]
    summary["final_tokens"] = final["tokens"] if final else 0
    summary["multi_call_blocks"] = sum(s["calls"] > 1 for s in spans)
    summary["first_mutation_step_proxy"] = sum(b["role"] == "assistant" for b in blocks[:first_mut + 1]) if first_mut is not None else 0
    signatures = set()
    for event in events:
        prefix = "target_" if event["target"] else "masked_"
        summary[prefix + "tool_" + event["name"]] += 1
        command = event["command"]
        history = bool(re.search(r"\bgit\s+(?:log|show|blame|rev-list|reflog)\b", command))
        summary[prefix + "git_history"] += history
        summary[prefix + "test_commands"] += bool(event["test"])
        if history:
            summary[prefix + "history_tasks"] = 1
        key = (event["name"], json.dumps(event["args"], sort_keys=True), sr.observation_digest(event["observation"]))
        if key in signatures and event["observation"]:
            summary[prefix + "exact_repeat_calls"] += 1
        signatures.add(key)
        if event["test"]:
            summary[prefix + "test_state_" + event["test"]["state"]] += 1
    return dict(metrics=dict(summary), thoughts=thoughts, spans=spans,
                final_excerpt=(final["text"][:1600] if final else ""),
                history_commands=[e["command"][:500] for e in events if e["target"] and re.search(r"\bgit\s+(?:log|show|blame|rev-list|reflog)\b", e["command"])][:5])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-root", type=Path, required=True)
    parser.add_argument("--reference", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    root = args.source_root.resolve()
    repo = Path(__file__).resolve().parents[2]
    source = root / "swefactory_self/results"
    sys.path.insert(0, str(source))
    sys.path.insert(0, str(repo / "training/swe_world"))
    import segment_review as sr
    from tokenized_sft import validate_row
    from tokenizers import Tokenizer

    files = {"A": root / "analysis/swe/family_recovery_20261024/cap3_release/pool_a_cap3_combined.jsonl"}
    files.update({p: source / f"pool_bc/pool_{p.lower()}.jsonl" for p in ["B", "C1", "C2", "C3"]})
    frozen = json.loads((repo / "storage/audits/swe50-comparison-20261028/training-supervision-audit.json").read_text(encoding="utf-8"))
    expected = {Path(f["file"].replace("\\", "/")).name: f["sha256"] for f in frozen["files"]}
    tokenizer_path = root / "model/Qwen3.5-4B/tokenizer.json"
    tokenizer = Tokenizer.from_file(str(tokenizer_path))
    source_db = source / "pool_a_formal/index.sqlite3"
    original = sqlite3.connect(source_db.as_uri() + "?mode=ro", uri=True)
    original.row_factory = sqlite3.Row
    metadata = {r["instance_id"]: dict(r) for r in original.execute("SELECT * FROM samples")}
    ledgers = {}
    for path in [source / "pool_bc/a_reviewed_ledger.jsonl", source / "pool_bc/bc_ledger.jsonl", root / "analysis/swe/family_recovery_20261024/cap3_release/reviewed_ledger.jsonl"]:
        for row in read_lines(path):
            ledgers[row["sample_id"]] = row
    protected_path = root / "swefactory_self/configs/data/protected_swe_families.json"
    protected = json.loads(protected_path.read_text(encoding="utf-8"))["protected_families"]
    protected = {p.casefold() for p in protected}
    report = dict(schema_version=1, generated_at=datetime.now(timezone.utc).isoformat(),
                  scope="Static full-payload audit and hypothetical merge accounting; no training/judge rerun.",
                  training_eligible=False, benchmark_training_eligible=False,
                  reference=dict(path=str(args.reference), sha256=digest(args.reference)),
                  tokenizer_sha256=digest(tokenizer_path), sources={}, pools={}, pairs={}, scenarios={},
                  validation_errors=[], instances=[], b_cases=[], c3_cases=[])
    with tempfile.TemporaryDirectory(prefix="acb-quality-") as temp:
        db = sqlite3.connect(str(Path(temp) / "audit.sqlite3"))
        db.row_factory = sqlite3.Row
        db.executescript("""
        CREATE TABLE sequences (iid TEXT PRIMARY KEY, sh TEXT, ids BLOB);
        CREATE TABLE views (sid TEXT PRIMARY KEY, pool TEXT, iid TEXT, sh TEXT,
            mh TEXT, n INTEGER, sup INTEGER, mask BLOB, keeps TEXT, spans TEXT);
        CREATE INDEX by_instance ON views(iid);
        """)
        for pool, path in files.items():
            sha = digest(path)
            assert sha == expected[path.name], f"Frozen source changed: {path}"
            report["sources"][pool] = dict(path=str(path), sha256=sha, bytes=path.stat().st_size, matches_frozen=True)
            for row in read_lines(path):
                validate_row(row, 65536, tokenizer.get_vocab_size())
                admission = ledgers.get(row["sample_id"], {})
                assert admission.get("training_eligible"), f"Missing eligible ledger: {row['sample_id']}"
                assert admission["keep_msgs"] == row["keep_msgs"], "Ledger/payload supervision differs"
                assert admission["seq_tokens"] == len(row["ids"]), "Ledger/payload length differs"
                assert admission["supervised_tokens"] == sum(row["mask"]), "Ledger/payload target count differs"
                ids = array("I", row["ids"]).tobytes()
                mask = bytes(row["mask"])
                sh, mh = hashlib.sha256(ids).hexdigest(), hashlib.sha256(mask).hexdigest()
                previous = db.execute("SELECT sh FROM sequences WHERE iid=?", (row["instance_id"],)).fetchone()
                assert previous is None or previous["sh"] == sh, "Same instance has different contexts"
                db.execute("INSERT OR IGNORE INTO sequences VALUES (?,?,?)", (row["instance_id"], sh, zlib.compress(ids)))
                db.execute("INSERT INTO views VALUES (?,?,?,?,?,?,?,?,?,?)", (row["sample_id"], pool, row["instance_id"], sh, mh, len(row["ids"]), sum(mask), zlib.compress(mask), json.dumps(row["keep_msgs"]), json.dumps(row["spans"])))
            db.commit()
            rows = db.execute("SELECT n,sup FROM views WHERE pool=?", (pool,)).fetchall()
            report["pools"][pool] = dict(rows=len(rows), context=quantiles([r["n"] for r in rows]), supervised=quantiles([r["sup"] for r in rows]))
            print(f"Indexed {pool}: {len(rows)} rows", flush=True)
        # Pool overlap is measured on token positions, not text similarity.
        for other in ["B", "C1", "C2", "C3"]:
            counts = Counter()
            for row in db.execute("SELECT b.*,a.mask AS amask,a.sup AS asup FROM views b LEFT JOIN views a ON a.iid=b.iid AND a.pool='A' WHERE b.pool=?", (other,)):
                if row["amask"] is None:
                    counts["not_in_A"] += 1
                    continue
                amask, bmask = zlib.decompress(row["amask"]), zlib.decompress(row["mask"])
                common = sum(a and b for a, b in zip(amask, bmask))
                counts["shared"] += 1
                counts["identical_mask" if amask == bmask else "subset_mask" if common == row["sup"] else "other_mask"] += 1
                counts["overlap_tokens"] += common
                counts["novel_tokens_over_A"] += row["sup"] - common
                counts["A_tokens_outside_view"] += row["asup"] - common
            report["pairs"]["A/" + other] = dict(counts)
        all_views = list(db.execute("SELECT sid,pool,iid,sh,mh,n,sup FROM views"))
        def scenario(rows):
            return dict(rows=len(rows), instances=len({r["iid"] for r in rows}), context_tokens=sum(r["n"] for r in rows), supervised_tokens=sum(r["sup"] for r in rows))
        report["scenarios"]["original"] = scenario(all_views)
        seen, exact = set(), []
        for row in all_views:
            key = (row["iid"], row["sh"], row["mh"])
            if row["pool"] == "B" or key not in seen:
                exact.append(row)
                seen.add(key)
        report["scenarios"]["exact_AC_then_B"] = scenario(exact)
        main = [r for r in all_views if r["pool"] in ["A", "B", "C3"]]
        report["scenarios"]["union_AC_then_B"] = scenario(main)
        report["scenarios"]["union_ABC"] = scenario([r for r in main if r["pool"] != "B"])
        c2ids = {r["iid"] for r in all_views if r["pool"] == "C2"}
        c2prefer = [r for r in all_views if r["pool"] in ["B", "C2", "C3"] or (r["pool"] == "A" and r["iid"] not in c2ids)]
        report["scenarios"]["prefer_C2_then_B"] = scenario(c2prefer)
        report["scenarios"]["union_AC12_then_B"] = scenario([r for r in main if r["pool"] != "C3"])
        b_masks = {r["iid"]: zlib.decompress(r["mask"]) for r in db.execute("SELECT iid,mask FROM views WHERE pool='B'")}
        scenario_rows = dict(original=all_views, exact_AC_then_B=exact, union_AC_then_B=main,
                             union_ABC=[r for r in main if r["pool"] != "B"], prefer_C2_then_B=c2prefer,
                             union_AC12_then_B=[r for r in main if r["pool"] != "C3"])
        c1ids = {r["iid"] for r in all_views if r["pool"] == "C1"}
        for key, rows in scenario_rows.items():
            recovery_exposures = 0
            for row in rows:
                if row["iid"] in b_masks:
                    mask = zlib.decompress(db.execute("SELECT mask FROM views WHERE sid=?", (row["sid"],)).fetchone()[0])
                    recovery_exposures += sum(a and b for a, b in zip(mask, b_masks[row["iid"]]))
            report["scenarios"][key].update(
                B_region_exposures=recovery_exposures,
                B_region_share=recovery_exposures / report["scenarios"][key]["supervised_tokens"],
                C1_member_rows=sum(r["iid"] in c1ids for r in rows),
                C2_member_rows=sum(r["iid"] in c2ids for r in rows))
        family_counts, family_tokens, meta_counts = Counter(), Counter(), Counter()
        tool_metrics, thought_lengths, lengths, steps, files_changed = Counter(), [], [], [], []
        for index, seq in enumerate(db.execute("SELECT * FROM sequences ORDER BY iid")):
            iid = seq["iid"]
            views = list(db.execute("SELECT * FROM views WHERE iid=?", (iid,)))
            base = next((r for r in views if r["pool"] == "A"), next(r for r in views if r["pool"] == "C3") if all(r["pool"] != "A" for r in views) else views[0])
            ids = array("I")
            ids.frombytes(zlib.decompress(seq["ids"]))
            mask = zlib.decompress(base["mask"])
            info = parse_sequence(list(ids), mask, tokenizer, sr)
            assert info["metrics"]["nonassistant_target_tokens"] == 0, "Non-assistant target found"
            meta = metadata.get(iid, {})
            ledger = ledgers.get(base["sid"], {})
            family = ledger.get("family") or meta.get("family")
            assert family, f"Missing family: {iid}"
            assert family.casefold() not in protected, f"Protected family: {family}"
            family_counts[family] += 1
            family_tokens[family] += base["sup"]
            metrics = info["metrics"]
            record = dict(instance_id=iid, family=family, pools=[r["pool"] for r in views],
                          source_group="cap3" if "CAP3_FAMILY_RECOVERY" in ledger.get("reasons", []) else "baseline" if base["pool"] == "A" else "C3",
                          context_tokens=base["n"], supervised_tokens=base["sup"],
                          ledger_eligible=ledger.get("training_eligible"), ledger_masks=len(ledger.get("masked_msgs", [])),
                          evaluator_success=meta.get("success"), node_evidence=meta.get("node_evidence"),
                          f2p_passed=meta.get("f2p_passed"), p2p_passed=meta.get("p2p_passed"),
                          patch_chars=meta.get("patch_chars"), files_changed=meta.get("files_changed"),
                          steps=meta.get("steps") or ledger.get("steps"), finish_reason=meta.get("finish_reason"),
                          teacher_path_label=Path(str(meta.get("traj_path", "")).replace("\\", "/")).parent.name,
                          metrics=metrics)
            report["instances"].append(record)
            if base["pool"] == "A":
                tool_metrics.update(metrics)
                thought_lengths.extend(info["thoughts"])
                lengths.append(base["n"])
                if record["steps"] is not None:
                    steps.append(record["steps"])
                if record["files_changed"] is not None:
                    files_changed.append(record["files_changed"])
                meta_counts["node_" + str(meta.get("node_evidence"))] += 1
                meta_counts["finish_" + str(meta.get("finish_reason"))] += 1
                meta_counts["evaluator_success"] += bool(meta.get("success"))
                meta_counts["sql_verified_after_last_edit"] += bool(meta.get("verified_after_last_edit"))
                meta_counts["masked_trajectories"] += bool(ledger.get("masked_msgs"))
                meta_counts["masked_messages"] += len(ledger.get("masked_msgs", []))
            if "B" in record["pools"]:
                b = next(r for r in views if r["pool"] == "B")
                b_ledger = ledgers[b["sid"]]
                evidence = b_ledger.get("evidence", {})
                binfo = parse_sequence(list(ids), zlib.decompress(b["mask"]), tokenizer, sr)
                assert binfo["metrics"]["nonassistant_target_tokens"] == 0, "Non-assistant B target found"
                keeps = json.loads(b["keeps"])
                report["b_cases"].append(dict(instance_id=iid, family=family, supervised_tokens=b["sup"], context_tokens=b["n"],
                    kept_messages=len(keeps), before_or_at_failure=sum(k <= evidence.get("fail_obs_msg", -1) for k in keeps),
                    fail_state=evidence.get("fail_state"), pass_state=evidence.get("pass_state"), receipt=evidence.get("receipt"),
                    metrics=binfo["metrics"], final_excerpt=binfo["final_excerpt"]))
            if base["pool"] == "C3":
                report["c3_cases"].append(dict(instance_id=iid, family=family, metrics=metrics, evidence=ledger.get("evidence"), final_excerpt=info["final_excerpt"]))
            if index % 200 == 0:
                print(f"Decoded {index + 1} unique contexts", flush=True)
        report["quality"] = dict(A_metadata=dict(meta_counts), A_behavior=dict(tool_metrics),
                                  thought_lengths=quantiles(thought_lengths), context_lengths=quantiles(lengths),
                                  steps=quantiles(steps), changed_files=quantiles(files_changed),
                                  families=len(family_counts), family_sizes=dict(Counter(family_counts.values())),
                                  top_families=family_tokens.most_common(15), protected_family_intersections=0)
        report["frozen_train_accounting"] = dict(original=frozen["train"], unique_positions=frozen["train_unique_positions"], pools=frozen["pools"],
            union_AC_then_B_rows=sum(frozen["pools"]["train/" + p]["rows"] for p in ["A", "B", "C3"]),
            union_AC_then_B_context=sum(frozen["pools"]["train/" + p]["context"] for p in ["A", "B", "C3"]),
            union_AC_then_B_supervised=sum(frozen["pools"]["train/" + p]["supervised"] for p in ["A", "B", "C3"]),
            note="Arithmetic projection from frozen split totals; original member-level split manifest unavailable in this checkout.")
        db.close()
    original.close()
    report["code_sha256"] = digest(Path(__file__))
    report["evidence_sources"] = {str(p.relative_to(root)): digest(p) for p in [source_db, protected_path, source / "formal_export.py", source / "bc_export.py", source / "segment_review.py", source / "pool_bc/a_reviewed_ledger.jsonl", source / "pool_bc/bc_ledger.jsonl", root / "analysis/swe/family_recovery_20261024/cap3_release/reviewed_ledger.jsonl"]}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({k: report[k] for k in ["scenarios", "pairs", "quality"]}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()