#!/usr/bin/env python3
"""Run pre-result Round 08 identity gates."""

from __future__ import annotations

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

import torch


def load_runner() -> Any:
    path = Path(__file__).resolve().parent / "train.py"
    spec = importlib.util.spec_from_file_location("k3_attnres_round08_train", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"cannot import Round 08 runner from {path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


runner = load_runner()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    step_zero = subparsers.add_parser("step-zero")
    step_zero.add_argument("--cache-dir", type=Path, required=True)
    step_zero.add_argument("--parent-manifest", type=Path, required=True)
    step_zero.add_argument("--study-manifest", type=Path, required=True)
    step_zero.add_argument("--output", type=Path, required=True)
    step_zero.add_argument("--seed", type=int, default=2026073001)

    smoke = subparsers.add_parser("smoke-compare")
    smoke.add_argument("--parent", type=Path, required=True)
    smoke.add_argument("--wrapper", type=Path, required=True)
    smoke.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def canonical_sha256(value: Any) -> str:
    return hashlib.sha256(
        json.dumps(
            value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
        ).encode()
    ).hexdigest()


def tensor_sha256(value: torch.Tensor) -> str:
    return hashlib.sha256(runner.parent.tensor_bytes(value)).hexdigest()


def read_and_verify(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text())
    expected = value["canonical_sha256_without_self"]
    payload = {
        key: item
        for key, item in value.items()
        if key != "canonical_sha256_without_self"
    }
    if canonical_sha256(payload) != expected:
        raise RuntimeError(f"canonical self-hash failed: {path}")
    return value


def smoke_compare(args: argparse.Namespace) -> None:
    parent_result = read_and_verify(args.parent)
    wrapper_result = read_and_verify(args.wrapper)
    fields = (
        "architecture",
        "depth",
        "seed",
        "steps",
        "batch_size",
        "target_bytes_seen",
        "manifest",
        "model",
        "optimizer",
        "hashes",
        "evaluations",
        "diagnostics",
        "training_history",
        "gradient_gate",
        "environment",
    )
    checks = {
        field: parent_result[field] == wrapper_result[field]
        for field in fields
    }
    wrapper_identity = {
        "protocol": wrapper_result["protocol_id"] == runner.PROTOCOL_ID,
        "parent_protocol": (
            wrapper_result["parent_protocol_id"]
            == runner.PARENT_PROTOCOL_ID
        ),
        "variant": wrapper_result["variant"] == "learned_reference",
        "forward_audit": wrapper_result["forward_intervention"]["passed"],
    }
    passed = all(checks.values()) and all(wrapper_identity.values())
    result = {
        "schema_version": 1,
        "protocol_id": runner.PROTOCOL_ID,
        "gate": "empty-selector-parent-equivalence",
        "passed": passed,
        "field_checks": checks,
        "wrapper_identity": wrapper_identity,
        "excluded_fields": [
            "protocol wrapper fields",
            "timing",
            "self hash",
            "parent runner self hash",
            "study manifest",
        ],
        "parent_file": str(args.parent.resolve()),
        "wrapper_file": str(args.wrapper.resolve()),
    }
    result["canonical_sha256_without_self"] = canonical_sha256(result)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
    )
    if not passed:
        raise RuntimeError(f"empty-selector parent equivalence failed: {checks}")


def exact_structure(value: Any) -> Any:
    return json.loads(
        json.dumps(value, ensure_ascii=False, sort_keys=True)
    )


def step_zero(args: argparse.Namespace) -> None:
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required by the frozen step-zero gate")
    study_manifest = json.loads(args.study_manifest.read_text())
    if study_manifest["protocol_id"] != runner.PROTOCOL_ID:
        raise RuntimeError("study manifest mismatch")
    parent_manifest = json.loads(args.parent_manifest.read_text())
    if parent_manifest["protocol_id"] != runner.PARENT_PROTOCOL_ID:
        raise RuntimeError("parent manifest mismatch")

    parent = runner.parent
    parent.configure_round04_globals(32)
    device = torch.device("cuda")
    corpus = parent.round04.ByteCorpus(
        args.cache_dir, parent_manifest, device
    )
    inputs, targets = corpus.fixed_batch(
        corpus.diagnostic_starts, 0, 16
    )

    variants = ("learned_reference",) + tuple(
        study_manifest["variants"].keys()
    )
    observations: dict[str, Any] = {}
    reference_payload: dict[str, Any] | None = None
    for variant in variants:
        parent.configure_determinism(args.seed)
        runner.ACTIVE_VARIANT = variant
        model = runner.ForwardInterventionLanguageModel("block").to(device)
        initial_public = parent.named_state_hash(
            model, include_mixers=False
        )
        initial_mixer = parent.named_state_hash(
            model, include_mixers=True
        )
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits, trace = model(inputs, capture=True)
            loss = parent.cross_entropy(logits, targets)
        if trace is None:
            raise RuntimeError("step-zero trace missing")
        evaluation = parent.evaluate(model, corpus, 64, 8)
        diagnostic = parent.diagnostic(model, corpus, 16)
        payload = {
            "initial_public_hash": initial_public,
            "initial_mixer_hash": initial_mixer,
            "logits_sha256": tensor_sha256(logits),
            "loss_nats": loss.detach().cpu().item(),
            "evaluation": exact_structure(evaluation),
            "diagnostic": exact_structure(diagnostic),
        }
        if reference_payload is None:
            reference_payload = payload
        exact_checks = {
            key: payload[key] == reference_payload[key]
            for key in payload
        }
        selected = tuple(runner.VARIANTS[variant])
        capture_checks = {}
        for index in selected:
            summary = trace.depth_weights[index]
            source_count = summary["sources"]
            capture_checks[str(index)] = {
                "source_count": source_count,
                "expected_source_count": study_manifest[
                    "selected_source_counts"
                ][str(index)],
                "capture_summary_exact_vs_learned": (
                    payload["diagnostic"]["depth_weights"][index]
                    == reference_payload["diagnostic"]["depth_weights"][index]
                ),
                "passed": (
                    source_count
                    == study_manifest["selected_source_counts"][str(index)]
                    and payload["diagnostic"]["depth_weights"][index]
                    == reference_payload["diagnostic"]["depth_weights"][index]
                ),
            }
        runtime_uniform_gate = (
            model.uniform_weight_max_abs_error
            <= study_manifest["thresholds"][
                "uniform_weight_max_abs_error"
            ]
        )
        observations[variant] = {
            "payload": payload,
            "exact_vs_learned_reference": exact_checks,
            "selected_capture_checks": capture_checks,
            "pre_reduction_uniform_weight_max_abs_error": (
                model.uniform_weight_max_abs_error
            ),
            "pre_reduction_uniform_weight_gate": runtime_uniform_gate,
            "passed": (
                all(exact_checks.values())
                and all(
                    item["passed"] for item in capture_checks.values()
                )
                and runtime_uniform_gate
            ),
        }
        del model, logits, loss, trace
        torch.cuda.empty_cache()

    passed = all(item["passed"] for item in observations.values())
    result = {
        "schema_version": 1,
        "protocol_id": runner.PROTOCOL_ID,
        "gate": "step-zero-cross-variant-byte-exact",
        "seed": args.seed,
        "passed": passed,
        "variants": observations,
        "parent_manifest_sha256": runner.parent.file_sha256(
            args.parent_manifest
        ),
        "study_manifest_sha256": runner.parent.file_sha256(
            args.study_manifest
        ),
        "environment": {
            "gpu": torch.cuda.get_device_name(0),
            "torch": torch.__version__,
            "cuda": torch.version.cuda,
            "cublas_workspace_config": os.environ.get(
                "CUBLAS_WORKSPACE_CONFIG"
            ),
        },
    }
    result["canonical_sha256_without_self"] = canonical_sha256(result)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
    )
    if not passed:
        failed = [
            name
            for name, value in observations.items()
            if not value["passed"]
        ]
        raise RuntimeError(f"step-zero exactness failed: {failed}")


def main() -> None:
    args = parse_args()
    if args.command == "step-zero":
        step_zero(args)
    else:
        smoke_compare(args)


if __name__ == "__main__":
    main()