from __future__ import annotations

import csv
import hashlib
import json
import unittest
from collections import Counter
from pathlib import Path

import build_checksums
import run_realworld_study as study

HERE = Path(__file__).resolve().parent


def read_csv(name: str) -> list[dict[str, str]]:
    with (HERE / "results" / name).open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


class SemanticTraceUnitTests(unittest.TestCase):
    def test_tool_classes_are_mechanistic(self) -> None:
        self.assertEqual(study.tool_class("get_order_details"), "read")
        self.assertEqual(study.tool_class("cancel_reservation"), "write")
        self.assertEqual(study.tool_class("transfer_to_human_agents"), "handoff")
        self.assertEqual(study.tool_class("$SYNTHETIC_TOOL_CALL"), "other")

    def test_post_write_read_closure_is_computed_from_ordered_calls(self) -> None:
        simulation = {
            "termination_reason": "user_stop",
            "messages": [
                {"role": "user", "content": "yes, please do it"},
                {
                    "role": "assistant",
                    "tool_calls": [
                        {"id": "w1", "name": "update_reservation_flights", "arguments": {"id": "A"}}
                    ],
                },
                {"role": "tool", "id": "w1", "error": False, "content": "ok"},
                {
                    "role": "assistant",
                    "tool_calls": [
                        {"id": "r1", "name": "get_reservation_details", "arguments": {"id": "A"}}
                    ],
                },
                {"role": "tool", "id": "r1", "error": False, "content": "state"},
            ],
        }
        features = study.extract_features(simulation, "airline", 32)
        self.assertEqual(features["raw_write_count"], 1)
        self.assertEqual(features["raw_verified_write_count"], 1)
        self.assertEqual(features["post_write_read_fraction"], 1.0)
        self.assertEqual(features["confirmed_write_fraction"], 1.0)

    def test_features_ignore_oracle_and_reward_fields(self) -> None:
        base = {
            "termination_reason": "user_stop",
            "messages": [
                {"role": "assistant", "content": "done", "tool_calls": None},
                {"role": "user", "content": "thanks", "tool_calls": None},
            ],
        }
        poisoned = {
            **base,
            "reward_info": {"reward": 0, "db_check": {"db_match": False}},
            "evaluation_criteria": {"actions": ["secret"]},
            "assertions": ["secret"],
        }
        self.assertEqual(
            study.extract_features(base, "retail", 32),
            study.extract_features(poisoned, "retail", 32),
        )

    def test_task_fold_is_stable_and_bounded(self) -> None:
        first = study.stable_fold("retail:42", 8)
        self.assertEqual(first, study.stable_fold("retail:42", 8))
        self.assertIn(first, range(8))

    def test_model_features_exclude_outcome_fields(self) -> None:
        features = set(study.TRACE_FEATURES + study.SEMANTIC_FEATURES)
        forbidden = {"success", "reward", "reward_was_missing", "db_check", "assertions"}
        self.assertFalse(features & forbidden)


class SemanticTraceArtifactTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.features = read_csv("trajectory_features.csv")
        cls.predictions = read_csv("crossfit_predictions.csv")
        cls.splits = read_csv("split_inventory.csv")
        cls.summary = json.loads((HERE / "results" / "summary.json").read_text(encoding="utf-8"))

    def test_real_cohort_shape_and_clusters(self) -> None:
        self.assertEqual(len(self.features), 4448)
        self.assertEqual(len({row["group_id"] for row in self.features}), 278)
        self.assertEqual(set(Counter(row["group_id"] for row in self.features).values()), {16})
        self.assertEqual({row["domain"] for row in self.features}, {"airline", "retail", "telecom"})
        self.assertEqual(len({row["model_id"] for row in self.features}), 4)

    def test_crossfit_prediction_inventory(self) -> None:
        official = [row for row in self.predictions if row["cohort"] == "official"]
        self.assertEqual(len(official), 4448 * 5)
        counts = Counter(row["row_id"] for row in official)
        self.assertEqual(set(counts.values()), {5})
        self.assertEqual(
            {row["method"] for row in official},
            {"prevalence", "trace", "semantic", "toolbag", "full"},
        )
        self.assertTrue(all(0.0 < float(row["probability"]) < 1.0 for row in official))

    def test_split_inventory_holds_out_model_and_task_fold(self) -> None:
        official = [row for row in self.splits if row["cohort"] == "official"]
        self.assertEqual(len(official), 4 * 8)
        for row in official:
            self.assertNotIn(row["heldout_model"], row["training_models"].split(";"))
            self.assertGreater(int(row["training_task_groups"]), 0)
            self.assertGreater(int(row["test_task_groups"]), 0)
        self.assertFalse(self.summary["crossfit"]["task_overlap_between_train_and_test"])
        self.assertFalse(self.summary["crossfit"]["model_overlap_between_train_and_test"])

    def test_negative_result_is_reported_with_uncertainty(self) -> None:
        contrast = self.summary["primary_contrasts"]["semantic_minus_trace"]
        self.assertGreater(contrast["difference"], 0.0)
        self.assertLessEqual(contrast["ci_lower"], 0.0)
        self.assertGreaterEqual(contrast["ci_upper"], 0.0)
        sensitivity = self.summary["numeric_reward_only_sensitivity"]["semantic_minus_trace"]
        self.assertLessEqual(sensitivity["ci_lower"], 0.0)
        self.assertGreaterEqual(sensitivity["ci_upper"], 0.0)

    def test_raw_inputs_and_outputs_are_hashed(self) -> None:
        manifest = json.loads((HERE / "results" / "run_manifest.json").read_text(encoding="utf-8"))
        self.assertEqual(len(manifest["inputs"]["raw_sources"]), 12)
        for relative, expected in manifest["outputs"].items():
            path = HERE / relative
            self.assertTrue(path.exists(), relative)
            self.assertEqual(hashlib.sha256(path.read_bytes()).hexdigest(), expected, relative)
        self.assertEqual(
            hashlib.sha256((HERE / "study_config.json").read_bytes()).hexdigest(),
            manifest["inputs"]["study_config.json"],
        )
        self.assertEqual(
            hashlib.sha256((HERE / "run_realworld_study.py").read_bytes()).hexdigest(),
            manifest["inputs"]["run_realworld_study.py"],
        )

    def test_released_feature_table_contains_no_transcript_text(self) -> None:
        columns = set(self.features[0])
        self.assertFalse({"content", "arguments", "policy", "evaluation_criteria", "reward_info"} & columns)

    def test_package_checksum_inventory(self) -> None:
        inventory = json.loads((HERE / "SHA256SUMS.json").read_text(encoding="utf-8"))
        self.assertEqual(inventory["algorithm"], "SHA-256")
        self.assertEqual(inventory["scope"], ".")
        self.assertEqual(inventory["files"], build_checksums.collect_hashes())
        canonical = json.dumps(
            inventory["files"], sort_keys=True, separators=(",", ":")
        ).encode("utf-8")
        self.assertEqual(
            inventory["inventory_sha256"],
            hashlib.sha256(canonical).hexdigest(),
        )


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