#!/usr/bin/env python3
"""CPU-only A-P6-1 checkpoint-sidecar adjudication for the four W1 cells."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any


def atomic_json(path: Path, value: Any) -> None:
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(value, sort_keys=True, indent=2) + "\n")
    os.replace(tmp, path)


def audit_cell(cell_dir: Path) -> dict[str, Any]:
    streams = sorted((cell_dir / "opprof").glob("*.jsonl"))
    sidecars = sorted((cell_dir / "opprof").glob("*.jsonl.footer.json"))
    if len(streams) != 1 or len(sidecars) != 1:
        raise RuntimeError(f"{cell_dir}: stream/sidecar count {len(streams)}/{len(sidecars)}")
    stream, sidecar_path = streams[0], sidecars[0]
    raw = stream.read_bytes()
    complete_newline = raw.endswith(b"\n")
    decoded = [json.loads(line) for line in raw.splitlines()]
    footers = [x for x in decoded if x.get("record_type") == "footer"]
    records = [x for x in decoded if "step_index" in x]
    sidecar = json.loads(sidecar_path.read_text())
    indices = [int(x["step_index"]) for x in records]
    checkpoint_stream_delta_s = abs(stream.stat().st_mtime_ns - int(sidecar["checkpoint_wall_ns"])) / 1e9
    anchor_results = []
    coverage = []
    for path in sorted(cell_dir.glob("anchor-*/result.json")):
        result = json.loads(path.read_text())
        lo = int(result["interval"]["start_mono_ns"])
        hi = int(result["interval"]["end_mono_ns"])
        selected = [x for x in records if lo <= int(x["submit_mono_ns"]) <= hi]
        # Layer-1 intervals are keyed by submit_mono_ns. A final async model step
        # may complete sub-millisecond after the client has received its last
        # response; the post-interval checkpoint proves that record is durable.
        covered = bool(selected) and max(int(x["submit_mono_ns"]) for x in selected) <= hi
        anchor_results.append({
            "anchor": result["anchor"], "result": str(path),
            "interval_start_mono_ns": lo, "interval_end_mono_ns": hi,
            "layer1_records": len(selected), "covered": covered,
        })
        coverage.append(covered)
    latest_anchor_wall_ns = max(
        int(json.loads(path.read_text())["interval"]["end_wall_ns"])
        for path in cell_dir.glob("anchor-*/result.json")
    )
    invariants = {
        "complete_final_newline": complete_newline,
        "all_schema_1": all(x.get("schema") == 1 for x in decoded) and sidecar.get("schema") == 1,
        "no_in_stream_footer": not footers,
        "checkpoint_sidecar": sidecar.get("final") is False,
        "steps_contiguous": indices == list(range(len(indices))),
        "written_matches_records": int(sidecar["written_records"]) == len(records),
        "encoded_balanced": int(sidecar["encoded_records"]) == int(sidecar["written_records"]) + int(sidecar["dropped_records"]),
        "last_step_matches": bool(records) and int(sidecar["last_step_index"]) == indices[-1],
        "zero_drops": int(sidecar["dropped_records"]) == 0 and all(int(x["dropped_records_before"]) == 0 for x in records),
        "checkpoint_within_flush_of_stream": checkpoint_stream_delta_s <= float(sidecar["flush_interval_seconds"]) + .1,
        "checkpoint_after_all_anchor_intervals": int(sidecar["checkpoint_wall_ns"]) >= latest_anchor_wall_ns,
        "two_anchor_intervals": len(anchor_results) == 2,
        "all_anchor_intervals_covered": len(coverage) == 2 and all(coverage),
    }
    return {
        "cell": cell_dir.name, "passed": all(invariants.values()),
        "stream": str(stream), "sidecar": str(sidecar_path), "records": len(records),
        "footer_count": len(footers), "checkpoint_stream_delta_s": checkpoint_stream_delta_s,
        "checkpoint_after_anchor_s": (int(sidecar["checkpoint_wall_ns"]) - latest_anchor_wall_ns) / 1e9,
        "counters": {key: sidecar[key] for key in ("encoded_records", "written_records", "dropped_records", "last_step_index")},
        "anchors": anchor_results, "invariants": invariants,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    cells = [audit_cell(path) for path in sorted((args.root / "cells").glob("tp1_mns*"))]
    result = {
        "schema": 1, "amendment": "A-P6-1", "mode": "checkpoint-sidecar",
        "cells": cells, "passed": len(cells) == 4 and all(x["passed"] for x in cells),
        "sanity": {
            "cell_count": len(cells), "anchor_count": sum(len(x["anchors"]) for x in cells),
            "records_min": min(x["records"] for x in cells), "records_max": max(x["records"] for x in cells),
            "records_distinct": len({x["records"] for x in cells}),
        },
    }
    for cell in cells:
        cell_dir = Path(cell["stream"]).parent.parent
        atomic_json(cell_dir / "cell-valid.json", {
            "cell": cell["cell"], "invariants": cell["invariants"],
            "layer1_records": cell["records"], "stream": cell["stream"],
            "accounting_mode": "A-P6-1-checkpoint-sidecar",
        })
    atomic_json(args.out, result)
    print(json.dumps({"passed": result["passed"], "sanity": result["sanity"], "cells": [{"cell": x["cell"], "passed": x["passed"], "invariants": x["invariants"]} for x in cells]}, sort_keys=True))
    if not result["passed"]:
        raise RuntimeError("W1 checkpoint-sidecar re-adjudication failed")


if __name__ == "__main__":
    main()