import json
import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

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


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


class GeneralizationStudyTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.protocol = json.loads(
            (
                ROOT
                / "config"
                / "policy_generalization_protocol_v1.json"
            ).read_text(encoding="utf-8")
        )
        cls.scenarios = load_scenarios(
            ROOT / "scenarios" / "external_agentdojo_40.json"
        )

    def test_schedule_is_deterministic_and_gate_is_embedded(self):
        settings = self.protocol["design"]
        gate = self.protocol["preflight_gate"]["scenario_ids"]
        first = extension_schedule_document(
            self.scenarios,
            settings["repetitions_per_scenario"],
            settings["randomization_seed"],
            gate,
        )
        second = extension_schedule_document(
            self.scenarios,
            settings["repetitions_per_scenario"],
            settings["randomization_seed"],
            gate,
        )
        self.assertEqual(first, second)
        self.assertEqual(first["trial_count"], 400)
        self.assertEqual(first["gate_trial_count"], 12)
        self.assertEqual(
            [item["trial_id"] for item in first["trials"][:12]],
            first["gate_trial_ids"],
        )

    def test_mock_extension_resumes_after_gate_without_duplicates(self):
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory)
            env = os.environ.copy()
            env["PYTHONPATH"] = str(ROOT / "src")
            command = [
                sys.executable,
                str(
                    ROOT
                    / "scripts"
                    / "run_policy_generalization_extension.py"
                ),
                "--root",
                str(ROOT),
                "--provider",
                "mock",
                "--output-dir",
                str(output),
            ]
            gate = subprocess.run(
                command, env=env, text=True, capture_output=True
            )
            self.assertEqual(gate.returncode, 3, gate.stderr)
            status = json.loads(
                (
                    output
                    / "policy_generalization_status_mock.json"
                ).read_text(encoding="utf-8")
            )
            self.assertTrue(status["embedded_gate_passed"])
            self.assertEqual(status["primary_complete"], 12)

            completion = subprocess.run(
                command, env=env, text=True, capture_output=True
            )
            self.assertEqual(completion.returncode, 0, completion.stderr)
            status = json.loads(
                (
                    output
                    / "policy_generalization_status_mock.json"
                ).read_text(encoding="utf-8")
            )
            self.assertEqual(status["primary_complete"], 400)
            ledger = (
                output
                / "pfb-policy-generalization-v1-mock.jsonl"
            )
            before = ledger.read_text(encoding="utf-8")
            again = subprocess.run(
                command, env=env, text=True, capture_output=True
            )
            self.assertEqual(again.returncode, 0, again.stderr)
            self.assertEqual(before, ledger.read_text(encoding="utf-8"))


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