import argparse
import json
import tempfile
import unittest
from collections import Counter
from pathlib import Path

from policyfaultbench.cli import command_study_run
from policyfaultbench.study import read_ledger
from policyfaultbench.scenarios import load_scenarios
from policyfaultbench.study import build_schedule, schedule_document


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


class StudyTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.scenarios = load_scenarios(ROOT / "scenarios" / "core_40.json")

    def test_schedule_is_deterministic_balanced_and_randomized(self):
        first = schedule_document(self.scenarios, repetitions=10, seed=20260724)
        second = schedule_document(self.scenarios, repetitions=10, seed=20260724)
        self.assertEqual(first["schedule_hash"], second["schedule_hash"])
        self.assertEqual(first["trial_count"], 400)
        counts = Counter(row["scenario_id"] for row in first["trials"])
        self.assertEqual(set(counts.values()), {10})
        rep1 = [
            row["scenario_id"]
            for row in first["trials"]
            if row["repetition"] == 1
        ]
        rep2 = [
            row["scenario_id"]
            for row in first["trials"]
            if row["repetition"] == 2
        ]
        self.assertNotEqual(rep1, rep2)

    def test_mock_study_resumes_without_duplicate_trials(self):
        with tempfile.TemporaryDirectory() as directory:
            output = Path(directory)
            args = argparse.Namespace(
                root=ROOT,
                provider="mock",
                repetitions=2,
                seed=20260724,
                output_dir=output,
                delay_seconds=0.0,
                max_trials=7,
            )
            self.assertEqual(command_study_run(args), 3)
            args.max_trials = None
            self.assertEqual(command_study_run(args), 0)
            experiment = json.loads(
                (ROOT / "config" / "experiment.json").read_text(
                    encoding="utf-8"
                )
            )
            study_id = experiment["definitive_study"]["experiment_id"]
            event_path = output / f"{study_id}-mock.jsonl"
            before = event_path.read_text(encoding="utf-8")
            self.assertEqual(command_study_run(args), 0)
            self.assertEqual(before, event_path.read_text(encoding="utf-8"))
            status = json.loads(
                (output / "study_status_mock.json").read_text(encoding="utf-8")
            )
            self.assertEqual(status["planned_trials"], 80)
            self.assertEqual(status["primary_complete"], 80)
            self.assertEqual(status["remaining"], 0)

    def test_resume_budget_includes_terminal_abort_cost(self):
        with tempfile.TemporaryDirectory() as directory:
            ledger = Path(directory) / "ledger.jsonl"
            ledger.write_text(
                json.dumps(
                    {
                        "event_type": "run_aborted",
                        "payload": {
                            "trial": {"trial_id": "R01-S01-E"},
                            "terminal": True,
                            "abort_class": "model_conformance",
                            "cost_usd": 0.125,
                        },
                    }
                )
                + "\n",
                encoding="utf-8",
            )
            state = read_ledger(ledger)
            self.assertEqual(state.completed_cost_usd, 0.125)
            self.assertIn("R01-S01-E", state.terminal_trial_ids)


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