#!/usr/bin/env python3
"""Atomic-write, resume, failure-retention, and validation tests."""

from pathlib import Path
import sys
import tempfile
import unittest

import numpy as np


CODE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(CODE_ROOT))

from experiments.io import TrialPayload  # noqa: E402
from experiments.plan import build_trial_plan  # noqa: E402
from experiments.runner import run_trial_plan  # noqa: E402
from experiments.validate import validate_batch  # noqa: E402


def runner_specification():
    return {
        "study_id": "runner_contract",
        "split": "pilot",
        "root_seed": 17,
        "replicates": 2,
        "methods": ["a", "b"],
        "trajectories": ["trajectory_001"],
        "factors": {"delay_ms": [0]},
    }


def successful_executor(trial):
    method_offset = 0.0 if trial["method"]["method_id"] == "a" else 1.0
    time = np.arange(5, dtype=float) * 0.002
    return TrialPayload(
        samples={
            "time": time,
            "signal": np.full((5, 2), method_offset, dtype=float),
        },
        events=[{"sample_index": 2, "event": "marker"}],
        metadata={"executor": "unit-test"},
    )


class ExperimentRunnerContractTest(unittest.TestCase):
    def test_dry_run_writes_nothing(self):
        plan = build_trial_plan(runner_specification())
        with tempfile.TemporaryDirectory() as temporary:
            batch = Path(temporary) / "batch"
            report = run_trial_plan(plan, batch, dry_run=True)
            self.assertEqual(report["trial_count"], 4)
            self.assertFalse(batch.exists())

    def test_atomic_run_validate_and_resume(self):
        plan = build_trial_plan(runner_specification())
        with tempfile.TemporaryDirectory() as temporary:
            batch = Path(temporary) / "batch"
            first = run_trial_plan(plan, batch, successful_executor)
            self.assertEqual(first["completed_trials"], 4)
            self.assertEqual(first["failed_trials"], 0)
            validation = validate_batch(batch)
            self.assertTrue(validation["valid"], validation["errors"])
            self.assertEqual(validation["completed_trials"], 4)
            self.assertFalse(
                any(path.name.startswith(".") for path in (batch / "raw").iterdir())
            )

            second = run_trial_plan(plan, batch, successful_executor, resume=True)
            self.assertEqual(second["completed_trials"], 0)
            self.assertEqual(second["skipped_trials"], 4)

    def test_failure_is_retained_without_committing_trial(self):
        plan = build_trial_plan(runner_specification())

        def failing_executor(_trial):
            raise RuntimeError("intentional failure")

        with tempfile.TemporaryDirectory() as temporary:
            batch = Path(temporary) / "batch"
            report = run_trial_plan(
                plan,
                batch,
                failing_executor,
                max_trials=1,
            )
            self.assertEqual(report["failed_trials"], 1)
            self.assertFalse((batch / "raw" / plan["trials"][0]["trial_id"]).exists())
            failure_files = list((batch / "failures").rglob("*.json"))
            self.assertEqual(len(failure_files), 1)


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