import unittest
from pathlib import Path

from policyfaultbench.capability_mutations import (
    CAPABILITY_OPERATORS,
    generate_capability_mutants,
)
from policyfaultbench.capability_policy import (
    CapabilityPolicy,
    CapabilityPolicyEngine,
)
from policyfaultbench.generalization_eval import evaluate_generalization
from policyfaultbench.policy import PolicyEngine, PolicySet
from policyfaultbench.scenarios import load_scenarios


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


class PolicyGeneralizationTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.internal = load_scenarios(ROOT / "scenarios/core_40.json")
        cls.external = load_scenarios(
            ROOT / "scenarios/external_agentdojo_40.json"
        )
        cls.reference = PolicySet.load(
            ROOT / "policies/reference_policy.json"
        )
        cls.capability = CapabilityPolicy.load(
            ROOT / "policies/capability_policy_v1.json"
        )

    def test_external_corpus_is_balanced_and_source_grounded(self):
        self.assertEqual(len(self.external), 40)
        self.assertEqual(
            sum(item.expected_decision == "allow" for item in self.external),
            20,
        )
        self.assertEqual(
            sum(item.expected_decision == "deny" for item in self.external),
            20,
        )
        self.assertEqual(
            {item.source_metadata["suite"] for item in self.external},
            {"banking", "workspace", "slack", "travel"},
        )

    def test_both_policies_conform_on_both_frozen_corpora(self):
        engines = (
            PolicyEngine(self.reference),
            CapabilityPolicyEngine(self.capability),
        )
        for scenarios in (self.internal, self.external):
            for engine in engines:
                with self.subTest(
                    corpus=scenarios[0].scenario_id,
                    engine=type(engine).__name__,
                ):
                    self.assertTrue(
                        all(
                            engine.evaluate(item.proposed_call).effect
                            == item.expected_decision
                            for item in scenarios
                        )
                    )

    def test_capability_mutants_are_complete_and_distinct(self):
        mutants = generate_capability_mutants(self.capability)
        self.assertEqual(len(CAPABILITY_OPERATORS), 12)
        self.assertEqual(len(mutants), 12)
        self.assertEqual(
            {item.mutation_id for item in mutants},
            {f"C{index:02d}" for index in range(1, 13)},
        )
        self.assertEqual(len({str(item.as_dict()) for item in mutants}), 12)

    def test_expanded_probes_kill_all_capability_mutants(self):
        report = evaluate_generalization(ROOT)
        capability_reports = [
            item
            for item in report["reports"]
            if item["policy"] == "capability_policy_v1"
        ]
        self.assertEqual(len(capability_reports), 2)
        for item in capability_reports:
            self.assertEqual(item["reference_failures"], [])
            self.assertEqual(item["expanded_killed"], 12)
            self.assertEqual(item["expanded_mutation_score"], 1.0)


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