#!/usr/bin/env python3
"""Freeze all 256 Chat-generation prompt hashes into the Round 08 manifest."""

from __future__ import annotations

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


ROOT = Path(__file__).resolve().parents[1]
DEEPSEEK_EXPERIMENTS = ROOT / "experiments" / "deepseek"
sys.path.insert(0, str(DEEPSEEK_EXPERIMENTS))

from transformers import AutoTokenizer  # noqa: E402

import v2_lite_chat_special_token_behavior_probe as behavior  # noqa: E402
import v2_lite_routing_special_token_family_control as special  # noqa: E402


PROTOCOL_ID = "llm-atlas-deepseek-chat-task-bootstrap-crn-v1"
CONDITIONS = (
    "s0_eos",
    "s1_eos",
    "s0_period",
    "s1_period",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--artifact-dir", type=Path, required=True)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument(
        "--reference-routing-json",
        type=Path,
        required=True,
    )
    parser.add_argument("--human-eval", type=Path, required=True)
    parser.add_argument("--gsm8k", type=Path, required=True)
    parser.add_argument("--tnews", type=Path, required=True)
    parser.add_argument("--tnews-archive", type=Path, required=True)
    parser.add_argument("--wikitext", type=Path, required=True)
    return parser.parse_args()


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


def install_condition_contract() -> None:
    factors = {
        condition: special.FACTORS[condition]
        for condition in CONDITIONS
    }
    special.BOUNDARY_LEVELS = ("eos", "period")
    special.CONDITIONS = CONDITIONS
    special.FACTORS = factors
    special.SYSTEM_CELLS = {
        "eos": ("s0_eos", "s1_eos"),
        "period": ("s0_period", "s1_period"),
    }
    special.SYSTEM_EDGE_CONTRASTS = {
        "period_minus_eos": ("eos", "period"),
    }
    special.install_control_contract()


def main() -> None:
    args = parse_args()
    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    if manifest["protocol_id"] != PROTOCOL_ID:
        raise RuntimeError("manifest protocol ID differs")
    install_condition_contract()
    tokenizer = AutoTokenizer.from_pretrained(
        args.artifact_dir,
        trust_remote_code=True,
        local_files_only=True,
        use_fast=True,
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
    args.domains = ["code", "math"]
    args.per_domain = 32
    source_rows, _ = behavior.selected_sources(args, tokenizer)
    frozen_rows = {
        row["id"]: row for row in manifest["sources"]
    }
    audit = []
    for source in source_rows:
        frozen = frozen_rows.get(source["id"])
        if frozen is None:
            raise RuntimeError(f"{source['id']} is absent from manifest")
        hashes = {}
        tokens = {}
        for condition in CONDITIONS:
            rendered = special.prior.render_boundary_variant(
                tokenizer,
                source["content"],
                condition,
            )
            hashes[condition] = rendered["token_ids_sha256"]
            tokens[condition] = rendered["tokens"]
            audit.append(
                {
                    "source_id": source["id"],
                    "condition": condition,
                    "tokens": rendered["tokens"],
                    "sha256": rendered["token_ids_sha256"],
                }
            )
        frozen["chat_generation_prompt_token_ids_sha256"] = hashes
        frozen["chat_generation_prompt_tokens"] = tokens
    if len(audit) != 256 or len(source_rows) != 64:
        raise RuntimeError(
            f"expected 64 sources / 256 cells, got "
            f"{len(source_rows)} / {len(audit)}"
        )
    manifest["status"] = (
        "preregistered_and_chat_prompt_hashes_corrected_before_"
        "any_protocol_output"
    )
    manifest["prompt_contract"] = {
        "routing_probe_hashes_are_not_chat_generation_hashes": True,
        "chat_generation_hashes_frozen": True,
        "chat_generation_prompt_cells": len(audit),
        "chat_generation_prompt_hashes_sha256": canonical_hash(audit),
        "renderer": (
            "official apply_chat_template followed by the frozen "
            "single boundary-ID edit for period cells"
        ),
    }
    manifest.pop("canonical_content_sha256", None)
    manifest["canonical_content_sha256"] = canonical_hash(manifest)
    args.manifest.write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    payload = args.manifest.read_bytes()
    print(
        json.dumps(
            {
                "manifest": str(args.manifest),
                "bytes": len(payload),
                "sha256": hashlib.sha256(payload).hexdigest(),
                "canonical_content_sha256": manifest[
                    "canonical_content_sha256"
                ],
                "sources": len(source_rows),
                "prompt_cells": len(audit),
                "prompt_hashes_sha256": manifest["prompt_contract"][
                    "chat_generation_prompt_hashes_sha256"
                ],
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()