from __future__ import annotations

import hashlib
import json
from pathlib import Path

from policyfaultbench.generalization_study import (
    extension_schedule_document,
)
from policyfaultbench.scenarios import load_scenarios


ROOT = Path(__file__).resolve().parents[1]
FREEZE_PATH = ROOT / "config" / "policy_generalization_freeze_v1.json"

FROZEN_FILES = (
    "config/policy_generalization_protocol_v1.json",
    "external_sources/agentdojo_v1_source_manifest.json",
    "external_sources/AGENTDOJO_NOTICE.txt",
    "scenarios/external_agentdojo_40.json",
    "policies/reference_policy.json",
    "policies/capability_policy_v1.json",
    "manifests/owned_assets.json",
    "src/policyfaultbench/providers/base.py",
    "src/policyfaultbench/providers/openai.py",
    "src/policyfaultbench/providers/anthropic.py",
    "src/policyfaultbench/scenarios.py",
    "src/policyfaultbench/conformance.py",
    "src/policyfaultbench/oracles.py",
    "src/policyfaultbench/sandbox.py",
    "src/policyfaultbench/capability_policy.py",
    "src/policyfaultbench/capability_mutations.py",
    "src/policyfaultbench/generalization_eval.py",
    "src/policyfaultbench/generalization_runner.py",
    "src/policyfaultbench/generalization_study.py",
    "scripts/build_external_agentdojo_corpus.py",
    "scripts/freeze_policy_generalization_extension.py",
    "scripts/analyze_policy_generalization.py",
    "scripts/run_policy_generalization_extension.py",
    "scripts/replay_core_under_capability_policy.py",
    "scripts/analyze_policy_generalization_extension.py",
    "scripts/combine_policy_generalization_results.py",
    "results/generalization_v1/policy_generalization_analysis.json",
    "results/generalization_v1/core_capability_replay_analysis.json",
)


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    protocol = json.loads(
        (
            ROOT / "config" / "policy_generalization_protocol_v1.json"
        ).read_text(encoding="utf-8")
    )
    if protocol["status"] != "frozen_before_live_execution":
        raise RuntimeError("Protocol must be frozen before creating the freeze")
    scenarios = load_scenarios(
        ROOT / protocol["design"]["external_corpus"]
    )
    schedule = extension_schedule_document(
        scenarios,
        int(protocol["design"]["repetitions_per_scenario"]),
        int(protocol["design"]["randomization_seed"]),
        list(protocol["preflight_gate"]["scenario_ids"]),
    )
    if (
        schedule["schedule_hash"]
        != protocol["design"]["expected_schedule_hash"]
    ):
        raise RuntimeError("Schedule hash does not match the protocol")
    for relative in FROZEN_FILES:
        if not (ROOT / relative).is_file():
            raise RuntimeError(f"Missing frozen component: {relative}")
    freeze = {
        "schema_version": "1.0",
        "protocol_id": protocol["protocol_id"],
        "freeze_date": protocol["freeze_date"],
        "schedule_hash": schedule["schedule_hash"],
        "scenario_count": schedule["scenario_count"],
        "planned_trials_per_provider": schedule["trial_count"],
        "gate_trial_ids": schedule["gate_trial_ids"],
        "frozen_file_sha256": {
            relative: sha256(ROOT / relative)
            for relative in FROZEN_FILES
        },
        "reporting_commitment": protocol["reporting_commitment"],
    }
    if FREEZE_PATH.is_file():
        existing = json.loads(FREEZE_PATH.read_text(encoding="utf-8"))
        if existing != freeze:
            raise RuntimeError(
                "Existing freeze does not match the current extension"
            )
    else:
        FREEZE_PATH.write_text(
            json.dumps(freeze, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
    print(
        json.dumps(
            {
                "freeze": str(FREEZE_PATH),
                "schedule_hash": schedule["schedule_hash"],
                "frozen_file_count": len(FROZEN_FILES),
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()
