#!/usr/bin/env python3
"""Prepare exact Frontier fixtures for the P1 primary low/high probes.

Prompt-bearing band traces remain under ``--private-root``.  The emitted
fixtures and public manifest contain token IDs, block IDs, hashes, and
aggregate metadata, but no prompt text.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import math
import subprocess
import sys
from pathlib import Path
from typing import Any

from transformers import AutoTokenizer


HERE = Path(__file__).resolve().parent
AITUNER_ROOT = HERE.parents[1]
sys.path.insert(0, str(HERE))

import prepare_pilot as pilot  # noqa: E402


PRIMARY_ROLES = ("low1", "high1")


def load_module(path: Path):
    module_root = str(path.parent.resolve())
    if module_root not in sys.path:
        sys.path.insert(0, module_root)
    spec = importlib.util.spec_from_file_location("simfid_s2rb_prepare", path)
    if spec is None or spec.loader is None:
        raise ImportError(path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def order_hash(values: list[str]) -> str:
    return hashlib.sha256("\n".join(values).encode()).hexdigest()


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


def git_capture(root: Path, *arguments: str) -> str:
    return subprocess.run(
        ["git", "-C", str(root), *arguments],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    ).stdout


def raw_rows(path: Path) -> dict[int, dict[str, Any]]:
    result = {}
    with path.open(encoding="utf-8") as source:
        for index, line in enumerate(source):
            if line.strip():
                result[index] = json.loads(line)
    return result


def selected_hashes(
    selected: list[Any], rows: dict[int, dict[str, Any]]
) -> dict[str, str]:
    identifiers = []
    arrivals = []
    lengths = []
    for item in selected:
        row = rows[item.row_index]
        identifiers.append(str(row.get("request_id") or row.get("id") or item.row_index))
        arrivals.append(f"{float(item.timestamp) * 0.1:.12f}")
        lengths.append(str(int(item.input_length)))
    return {
        "request_id_order_sha256": order_hash(identifiers),
        "arrival_order_sha256": order_hash(arrivals),
        "input_length_order_sha256": order_hash(lengths),
    }


def kv_blocks(raw_root: Path, cell: str) -> int:
    stream = next((raw_root / cell / "opprof").glob("*.jsonl"))
    with stream.open(encoding="utf-8") as source:
        for line in source:
            record = json.loads(line)
            if "step_index" in record:
                return int(record["kv"]["total_blocks"])
    raise ValueError(f"no Layer-1 record for {cell}")


def source_window(windows_path: Path, window_id: str) -> tuple[dict[str, Any], Path]:
    return pilot.resolve_source_trace(windows_path, window_id)


def prepare(args: argparse.Namespace) -> dict[str, Any]:
    simulator = load_module(args.replayserve_root / "tools/simfid_s2rb_prepare.py")
    manifest = json.loads(args.pilot_manifest.read_text(encoding="utf-8"))
    window, trace = source_window(args.source_windows, args.source_window_id)
    if args.band_root is not None:
        role_paths = {
            role: (args.band_root / f"{role}.jsonl").resolve()
            for role in PRIMARY_ROLES
        }
        band_stats = {
            role: manifest["private"]["band_stats"][role]
            for role in PRIMARY_ROLES
        }
        for role, path in role_paths.items():
            if sha256_file(path) != band_stats[role]["sha256"]:
                raise ValueError(f"pre-materialized band hash mismatch: {role}")
        private_windows = None
    else:
        private_windows, all_band_stats = pilot.materialize_bands(
            trace, window, args.private_root
        )
        private_payload = json.loads(private_windows.read_text(encoding="utf-8"))
        role_paths = {
            item["fidelity_pilot_role"]: (
                private_windows.parent / item["trace_file"]
            ).resolve()
            for item in private_payload["windows"]
        }
        band_stats = {
            role: all_band_stats[role]
            for role in PRIMARY_ROLES
        }

    tokenizer = AutoTokenizer.from_pretrained(
        args.tokenizer, local_files_only=True, use_fast=True
    )
    fixture_root = args.output / "fixtures"
    config_root = args.output / "configs"
    fixture_root.mkdir(parents=True, exist_ok=True)
    config_root.mkdir(parents=True, exist_ok=True)
    entries = []
    red_flags = []
    for role in PRIMARY_ROLES:
        trace_path = role_paths[role]
        retained, trace_stats = simulator.scan_trace(trace_path)
        rows = raw_rows(trace_path)
        primary_pool = [retained[(index * len(retained)) // 512] for index in range(512)]
        selections: dict[str, list[Any]] = {}
        selected_union: set[int] = set()
        for cell, cell_manifest in sorted(manifest["cells"].items()):
            level = "low" if role.startswith("low") else "high"
            expected = cell_manifest["targets"][level]["selections"][role]
            pool = retained if int(cell_manifest["tp"]) == 4 else primary_pool
            selected = [item for item in pool if item.sampling_u <= float(expected["anchor"])]
            selections[cell] = selected
            selected_union.update(item.row_index for item in selected)
            hashes = selected_hashes(selected, rows)
            if len(selected) != int(expected["selected_count"]):
                red_flags.append(f"selection_count_{cell}_{role}")
            for key, value in hashes.items():
                if value != expected[key]:
                    red_flags.append(f"selection_hash_{cell}_{role}_{key}")

        token_gates, selected_records, block_stats = simulator.tokenize_and_hash(
            trace=trace_path,
            tokenizer=tokenizer,
            retained=retained,
            selected_union=selected_union,
        )
        if any(gate["status"] != "pass" for gate in token_gates.values()):
            red_flags.append(f"token_gate_{role}")
        for cell, selected in selections.items():
            cell_manifest = manifest["cells"][cell]
            level = "low" if role.startswith("low") else "high"
            expected = cell_manifest["targets"][level]["selections"][role]
            fixture_id = f"fidelity_p1_{cell}_{role}"
            cell_record = {
                "cell_id": cell,
                "tensor_parallel_size": int(cell_manifest["tp"]),
                "max_num_seqs": int(cell_manifest["mns"]),
                "store_role": "companion" if int(cell_manifest["tp"]) == 4 else "primary",
                "kv_capacity": {
                    "block_size_tokens": 16,
                    "num_blocks": kv_blocks(args.phase6_raw_root, cell),
                },
            }
            probe = {
                "probe_index": 0 if role == "low1" else 1,
                "sampling_u": float(expected["anchor"]),
            }
            fixture = simulator.create_fixture(
                fixture_root=fixture_root,
                fixture_id=fixture_id,
                cell=cell_record,
                probe=probe,
                row_indexes=[item.row_index for item in selected],
                meta_by_index={item.row_index: item for item in retained},
                selected_records=selected_records,
            )
            config_path = config_root / f"{fixture_id}.json"
            config = simulator.build_config(
                path=config_path,
                cell=cell_record,
                mode="frozen-calibrated",
                fixture_ids=[fixture_id],
                frontier_root=args.frontier_root,
                cache_dir=args.cache_dir,
            )
            entries.append(
                {
                    "cell": cell,
                    "role": role,
                    "level": level,
                    "anchor": expected["anchor"],
                    "selected_count": len(selected),
                    "fixture_id": fixture_id,
                    "fixture_manifest": str(
                        (fixture_root / fixture_id / "fixture_manifest.json").resolve()
                    ),
                    "frontier_csv": fixture["frontier_csv"]["path"],
                    "sidecar": fixture["sidecar_jsonl"]["path"],
                    "config": str(config_path.resolve()),
                    "calibration_scale": config["calibration"]["a_tp"],
                }
            )
        if block_stats["selected_union_records"] != len(selected_union):
            red_flags.append(f"selected_union_{role}")
        if trace_stats["retained_inclusive_0_8192"] < 512:
            red_flags.append(f"retained_too_small_{role}")

    selected_counts = [int(entry["selected_count"]) for entry in entries]
    calibration = [float(entry["calibration_scale"]) for entry in entries]
    result = {
        "schema": "fidelity-p1-frontier-prepared-v1",
        "status": "PASS" if not red_flags else "STOP",
        "source": {
            "pilot_manifest": str(args.pilot_manifest.resolve()),
            "source_windows": str(args.source_windows.resolve()),
            "source_window_id": args.source_window_id,
            "source_trace": str(trace.resolve()),
            "private_windows": (
                str(private_windows.resolve()) if private_windows is not None else None
            ),
            "pre_materialized_band_root": (
                str(args.band_root.resolve()) if args.band_root is not None else None
            ),
            "band_stats": band_stats,
        },
        "simulator": {
            "replayserve_root": str(args.replayserve_root.resolve()),
            "frontier_root": str(args.frontier_root.resolve()),
            "tokenizer": str(args.tokenizer.resolve()),
            "mode": "frozen-calibrated",
        },
        "generator": {
            "script": str(Path(__file__).resolve()),
            "script_sha256": sha256_file(Path(__file__).resolve()),
            "aituner_git_head": git_capture(AITUNER_ROOT, "rev-parse", "HEAD").strip(),
            "aituner_git_status_short": git_capture(AITUNER_ROOT, "status", "--short"),
        },
        "entries": entries,
        "sanity": {
            "red_flags": red_flags,
            "n": len(entries),
            "selected_count": {
                "n": len(selected_counts),
                "min": min(selected_counts),
                "max": max(selected_counts),
                "distinct_n": len(set(selected_counts)),
            },
            "calibration_scale": {
                "n": len(calibration),
                "min": min(calibration),
                "max": max(calibration),
                "distinct_n": len(set(calibration)),
            },
            "invariants": {
                "entries_12": len(entries) == 12,
                "roles_2": {entry["role"] for entry in entries} == set(PRIMARY_ROLES),
                "cells_6": len({entry["cell"] for entry in entries}) == 6,
                "selected_nonnegative": all(value > 0 for value in selected_counts),
                "per_config_not_identical": len(set(selected_counts)) > 1,
            },
        },
    }
    args.public_manifest.parent.mkdir(parents=True, exist_ok=True)
    args.public_manifest.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    return result


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser()
    result.add_argument("--pilot-manifest", type=Path, required=True)
    result.add_argument("--source-windows", type=Path, required=True)
    result.add_argument("--source-window-id", required=True)
    result.add_argument("--private-root", type=Path, required=True)
    result.add_argument("--band-root", type=Path)
    result.add_argument("--output", type=Path, required=True)
    result.add_argument("--public-manifest", type=Path, required=True)
    result.add_argument("--phase6-raw-root", type=Path, required=True)
    result.add_argument("--replayserve-root", type=Path, required=True)
    result.add_argument("--frontier-root", type=Path, required=True)
    result.add_argument("--cache-dir", type=Path, required=True)
    result.add_argument("--tokenizer", type=Path, required=True)
    return result


def main() -> None:
    result = prepare(parser().parse_args())
    print(
        json.dumps(
            {
                "status": result["status"],
                "entries": len(result["entries"]),
                "red_flags": result["sanity"]["red_flags"],
            },
            sort_keys=True,
        )
    )
    if result["status"] != "PASS":
        raise RuntimeError(result["sanity"]["red_flags"])


if __name__ == "__main__":
    main()