#!/usr/bin/env python3
"""Freeze and audit the prospective E13--E15 Study 5 cell manifests."""

from __future__ import annotations

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

from agent_harness.specs import (
    load_edit_interfaces,
    load_experiments,
    load_harnesses,
    load_models,
    load_task_split,
    load_tasks,
)


ROOT = Path(__file__).resolve().parents[1]
EXPERIMENTS = ("E13", "E14", "E15")
EXPECTED = {"E13": 1440, "E14": 540, "E15": 540}


def canonical_hash(value: Any) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"))
    return sha256(payload.encode("utf-8")).hexdigest()


def build_cells(root: Path, experiment_id: str) -> list[dict[str, Any]]:
    experiment = load_experiments(root)[experiment_id]
    tasks = load_tasks(root)
    harnesses = load_harnesses(root)
    interfaces = load_edit_interfaces(root)
    models = load_models(root)
    split = load_task_split(root / "tasks" / "splits" / f"{experiment.task_split}.txt")
    gate = json.loads(
        (root / "configs" / "gates" / "E09_model_interface_gate.json").read_text(
            encoding="utf-8"
        )
    )["selected"]
    enumerate_interfaces = experiment_id == "E14"
    cells: list[dict[str, Any]] = []
    for task_index, task_id in enumerate(split):
        task = tasks[task_id]
        if task.validation_status != "end_to_end_ready":
            raise RuntimeError(f"{task_id} is not end-to-end ready")
        for model_index, model_id in enumerate(experiment.model_ids):
            interface_ids = (
                experiment.edit_interface_ids
                if enumerate_interfaces
                else (str(gate[model_id]),)
            )
            treatments = [
                (harness_id, interface_id)
                for harness_id in experiment.harness_ids
                for interface_id in interface_ids
            ]
            offset = (task_index + model_index) % len(treatments)
            treatments = treatments[offset:] + treatments[:offset]
            for treatment_index, (harness_id, interface_id) in enumerate(treatments):
                harness = harnesses[harness_id]
                interface = interfaces[interface_id]
                model = models[model_id]
                cells.append(
                    {
                        "order": len(cells),
                        "task_id": task_id,
                        "repository_sha": task.base_commit,
                        "harness_id": harness_id,
                        "harness_hash": harness.config_hash,
                        "interface_id": interface_id,
                        "interface_hash": interface.config_hash,
                        "model_id": model_id,
                        "model_hash": model.config_hash,
                        "context_budget": experiment.context_budgets[0],
                        "seed": experiment.seeds[0],
                        "within_model_order": treatment_index,
                    }
                )
    return cells


def main() -> None:
    root = ROOT.resolve()
    output_dir = root / "configs" / "study5"
    output_dir.mkdir(parents=True, exist_ok=True)
    report: dict[str, Any] = {
        "schema_version": 1,
        "study": "Study 5 end-to-end harness behavior",
        "outcome_blind": True,
        "experiments": {},
    }
    for experiment_id in EXPERIMENTS:
        if (root / "results" / "raw" / experiment_id).exists():
            raise RuntimeError(
                f"{experiment_id} raw outcomes already exist; design cannot be re-frozen"
            )
        cells = build_cells(root, experiment_id)
        if len(cells) != EXPECTED[experiment_id]:
            raise RuntimeError(
                f"{experiment_id} expected {EXPECTED[experiment_id]} cells, got {len(cells)}"
            )
        identities = {
            (item["task_id"], item["harness_id"], item["interface_id"], item["model_id"])
            for item in cells
        }
        if len(identities) != len(cells):
            raise RuntimeError(f"{experiment_id} contains duplicate cell identities")
        manifest = {
            "schema_version": 1,
            "study": "Study 5 end-to-end harness behavior",
            "experiment_id": experiment_id,
            "outcome_blind": True,
            "planned_cells": len(cells),
            "cells": cells,
        }
        manifest["design_sha256"] = canonical_hash(manifest)
        path = output_dir / f"{experiment_id}_cells.json"
        path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        report["experiments"][experiment_id] = {
            "planned_cells": len(cells),
            "design_sha256": manifest["design_sha256"],
            "manifest": str(path.relative_to(root)),
        }
    report["audit_sha256"] = canonical_hash(report)
    (root / "docs" / "STUDY5_DESIGN_AUDIT.json").write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()