#!/usr/bin/env python3
"""Materialize the stable code window selected by audit_code_trace.py."""

from __future__ import annotations

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--audit", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--sample-seed", type=int, default=20260723)
    return parser.parse_args()


def session_uniform(seed: int, window_id: str, session_root: Any) -> float:
    payload = json.dumps(
        {"seed": seed, "window_id": window_id, "session_root": session_root},
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    return int.from_bytes(hashlib.blake2b(payload, digest_size=8).digest(), "big") / (
        1 << 64
    )


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    args = parse_args()
    audit = json.loads(args.audit.read_text())
    if audit["data_gate"] != "PASS":
        raise ValueError(f"trace data gate is not PASS: {audit['data_gate']}")
    selected = audit["selected"]
    source = Path(selected["source"])
    window = selected["stable_window"]
    start = float(window["start_timestamp"])
    end = float(window["end_timestamp"])
    if args.output_root.exists():
        raise ValueError(f"refusing to overwrite {args.output_root}")
    args.output_root.mkdir(parents=True)
    destination = args.output_root / "code-raw-window.jsonl"
    root_of: dict[Any, Any] = {}
    request_count = 0
    with source.open() as input_stream, destination.open("w") as output_stream:
        for source_index, line in enumerate(input_stream):
            if not line.strip():
                continue
            row = json.loads(line)
            timestamp = float(row["timestamp"])
            if timestamp < start:
                continue
            if timestamp >= end:
                break
            if int(row["input_length"]) <= 0 or int(row["output_length"]) <= 0:
                continue
            chat = row.get("chat_id", source_index)
            parent = row.get("parent_chat_id")
            has_parent = parent not in (None, "", -1, "-1")
            session_root = root_of.get(parent, parent) if has_parent else chat
            root_of[chat] = session_root
            materialized = {
                **row,
                "source_index": source_index,
                "session_root": session_root,
                "sampling_u": session_uniform(
                    args.sample_seed,
                    f"code-{start:.6f}-{end:.6f}",
                    session_root,
                ),
            }
            output_stream.write(
                json.dumps(materialized, ensure_ascii=False, separators=(",", ":"))
                + "\n"
            )
            request_count += 1
    expected = int(selected["selected_window_stats"]["requests"])
    if request_count != expected:
        raise ValueError(f"window request mismatch: materialized={request_count}, audit={expected}")
    manifest = {
        "schema": "frontier-code-window-v1",
        "audit": str(args.audit.resolve()),
        "audit_sha256": sha256(args.audit),
        "source": str(source.resolve()),
        "source_block_size": selected["hash_contract"]["exact_source_block_size"],
        "target_block_size": 16,
        "start_timestamp": start,
        "end_timestamp": end,
        "duration_s": end - start,
        "requests": request_count,
        "sample_seed": args.sample_seed,
        "sampling_rule": "session-coherent deterministic sampling_u",
        "max_model_len": audit["max_model_len_recommendation"],
        "window_stats": selected["selected_window_stats"],
        "raw_window": str(destination.resolve()),
        "raw_window_sha256": sha256(destination),
    }
    (args.output_root / "window-manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n"
    )
    print(json.dumps({"requests": request_count, "output_root": str(args.output_root)}))


if __name__ == "__main__":
    main()