from __future__ import annotations

from hashlib import sha256
import json
from pathlib import Path
from types import SimpleNamespace
import tempfile
import unittest

from agent_harness.components import Candidate
from agent_harness.live_agent_experiment import LiveToolHarness
from agent_harness.protocol_experiment import ProtocolWorkspace
from agent_harness.specs import (
    load_embeddings,
    load_harnesses,
    load_models,
    load_task_split,
    load_tasks,
)


ROOT = Path(__file__).resolve().parents[1]


class Study5ExperimentTests(unittest.TestCase):
    def test_frozen_manifest_counts_and_hashes(self) -> None:
        expected = {"E13": 1440, "E14": 540, "E15": 540, "E16": 306}
        for experiment_id, count in expected.items():
            path = ROOT / "configs" / "study5" / f"{experiment_id}_cells.json"
            value = json.loads(path.read_text(encoding="utf-8"))
            digest = value.pop("design_sha256")
            observed = sha256(
                json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")
            ).hexdigest()
            self.assertEqual(digest, observed)
            self.assertEqual(len(value["cells"]), count)
            identities = {
                (item["task_id"], item["harness_id"], item["interface_id"], item["model_id"])
                for item in value["cells"]
            }
            self.assertEqual(len(identities), count)

    def test_e16_selection_is_frozen_before_outcomes(self) -> None:
        selection = json.loads(
            (ROOT / "configs" / "study5" / "E16_selection.json").read_text(encoding="utf-8")
        )
        digest = selection.pop("selection_sha256")
        observed = sha256(
            json.dumps(selection, sort_keys=True, separators=(",", ":")).encode("utf-8")
        ).hexdigest()
        self.assertEqual(digest, observed)
        self.assertEqual(
            selection["selected_harnesses"],
            ["H000", "H001", "H002", "H007", "H008", "H014"],
        )
        self.assertTrue(selection["selection_is_outcome_blind_to_e16"])

    def test_fresh_validation_split_is_unique_and_disjoint(self) -> None:
        tasks = load_tasks(ROOT)
        split = load_task_split(ROOT / "tasks" / "splits" / "study5_fresh.txt")
        self.assertEqual(len(split), 17)
        commits = {tasks[item].gold_commit for item in split}
        prior = {task.gold_commit for task_id, task in tasks.items() if task_id not in split}
        self.assertEqual(len(commits), 17)
        self.assertFalse(commits & prior)

    def _tools(self, harness_id: str, tree: Path) -> LiveToolHarness:
        task = load_tasks(ROOT)["TASK_CR_001"]
        workspace = ProtocolWorkspace(tree, ("example.go",), task, 2)
        return LiveToolHarness(
            load_harnesses(ROOT)[harness_id],
            SimpleNamespace(),
            workspace,
            SimpleNamespace(),
            load_models(ROOT)["M002"],
            load_embeddings(ROOT)["EMB002"],
            lambda transition: None,
        )

    def test_one_shot_and_iterative_search_policies_are_operational(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            tree = Path(temporary)
            (tree / "example.go").write_text("package example\nfunc Work() {}\n", encoding="utf-8")
            one_shot = self._tools("H008", tree)
            one_shot._begin_search()
            with self.assertRaisesRegex(ValueError, "exactly one"):
                one_shot._begin_search()
            iterative = self._tools("H010", tree)
            iterative._begin_search()
            iterative._begin_search()
            self.assertEqual(iterative.search_call_count, 2)

    def test_search_observation_uses_declared_packing(self) -> None:
        candidate = Candidate(
            path="example.go",
            line_start=1,
            line_end=2,
            text="package example\nfunc Work() {}",
            source="test",
            score=1.0,
        )
        with tempfile.TemporaryDirectory() as temporary:
            tree = Path(temporary)
            (tree / "example.go").write_text("package example\nfunc Work() {}\n", encoding="utf-8")
            skeleton = self._tools("H013", tree)._packed_records((candidate,))[0]
            summary = self._tools("H015", tree)._packed_records((candidate,))[0]
            whole = self._tools("H014", tree)._packed_records((candidate,))[0]
            self.assertEqual(skeleton["packing"], "skeletons")
            self.assertIn("Work", skeleton["snippet"])
            self.assertEqual(summary["packing"], "role_summaries")
            self.assertIn("Work", summary["snippet"])
            self.assertEqual(whole["packing"], "whole_files")
            self.assertIn("package example", whole["snippet"])


if __name__ == "__main__":
    unittest.main()