#!/usr/bin/env python3
"""Reproducibility provenance and locked-split safeguards."""

import hashlib
from pathlib import Path
import sys
import tempfile
import unittest
from unittest.mock import patch


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

from experiments.manifest import build_batch_manifest, ensure_batch  # noqa: E402
from experiments.plan import build_trial_plan  # noqa: E402
from experiments.schema import ContractError  # noqa: E402


def _plan(split: str = "pilot"):
    return build_trial_plan(
        {
            "study_id": "manifest_contract",
            "split": split,
            "root_seed": 7,
            "replicates": 1,
            "methods": ["method"],
            "trajectories": ["trajectory"],
        }
    )


class ExperimentManifestTest(unittest.TestCase):
    def test_reproducibility_files_are_hashed_when_present(self):
        contents = {
            "pyproject.toml": b"[project]\nname='example'\n",
            "uv.lock": b"version = 1\n",
            "code/config/master_7dof.urdf": b"<robot name='master'/>\n",
            "code/config/real_slave_7dof.urdf": b"<robot name='slave'/>\n",
        }
        with tempfile.TemporaryDirectory() as temporary:
            source_root = Path(temporary)
            for relative_path, data in contents.items():
                path = source_root / relative_path
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_bytes(data)

            manifest = build_batch_manifest(_plan(), source_root=source_root)

        recorded = manifest["source"]["files"]
        self.assertEqual(set(recorded), set(contents))
        for relative_path, data in contents.items():
            self.assertEqual(
                recorded[relative_path]["sha256"],
                hashlib.sha256(data).hexdigest(),
            )

    def test_runtime_records_optional_scientific_versions_when_importable(self):
        manifest = build_batch_manifest(_plan())
        for module_name in ("scipy", "pinocchio"):
            try:
                module = __import__(module_name)
            except Exception:
                self.assertNotIn(module_name, manifest["runtime"])
            else:
                self.assertIn(module_name, manifest["runtime"])
                expected = getattr(module, "__version__", None)
                if expected is not None:
                    self.assertEqual(
                        manifest["runtime"][module_name],
                        str(expected),
                    )

    def test_locked_split_rejects_missing_git_metadata(self):
        with tempfile.TemporaryDirectory() as temporary:
            with self.assertRaisesRegex(
                ContractError, "locked split requires available Git metadata"
            ):
                build_batch_manifest(
                    _plan("locked"),
                    source_root=Path(temporary),
                )

    def test_locked_split_rejects_dirty_git_worktree(self):
        dirty = {"available": True, "commit": "abc123", "dirty": True}
        with patch("experiments.manifest._git_provenance", return_value=dirty):
            with self.assertRaisesRegex(
                ContractError, "locked split requires a clean Git worktree"
            ):
                build_batch_manifest(_plan("locked"))

    def test_pilot_split_allows_missing_git_metadata(self):
        unavailable = {"available": False, "commit": None, "dirty": None}
        with patch("experiments.manifest._git_provenance", return_value=unavailable):
            manifest = build_batch_manifest(_plan("pilot"))
        self.assertFalse(manifest["source"]["available"])

    def test_existing_locked_batch_rechecks_current_git_state(self):
        clean = {"available": True, "commit": "abc123", "dirty": False}
        unavailable = {"available": False, "commit": None, "dirty": None}
        with tempfile.TemporaryDirectory() as temporary:
            source_root = Path(temporary) / "source"
            batch_dir = Path(temporary) / "batch"
            source_root.mkdir()
            plan = _plan("locked")
            with patch("experiments.manifest._git_provenance", return_value=clean):
                ensure_batch(batch_dir, plan, source_root=source_root)

            with patch(
                "experiments.manifest._git_provenance",
                return_value=unavailable,
            ):
                with self.assertRaisesRegex(
                    ContractError, "locked split requires available Git metadata"
                ):
                    ensure_batch(batch_dir, plan, source_root=source_root)

    def test_existing_locked_batch_rejects_commit_drift(self):
        initial = {"available": True, "commit": "abc123", "dirty": False}
        changed = {"available": True, "commit": "def456", "dirty": False}
        with tempfile.TemporaryDirectory() as temporary:
            source_root = Path(temporary) / "source"
            batch_dir = Path(temporary) / "batch"
            source_root.mkdir()
            plan = _plan("locked")
            with patch("experiments.manifest._git_provenance", return_value=initial):
                ensure_batch(batch_dir, plan, source_root=source_root)

            with patch("experiments.manifest._git_provenance", return_value=changed):
                with self.assertRaisesRegex(
                    ContractError, "locked batch Git commit mismatch"
                ):
                    ensure_batch(batch_dir, plan, source_root=source_root)

    def test_existing_locked_batch_rejects_source_file_drift(self):
        clean = {"available": True, "commit": "abc123", "dirty": False}
        with tempfile.TemporaryDirectory() as temporary:
            source_root = Path(temporary) / "source"
            batch_dir = Path(temporary) / "batch"
            source_root.mkdir()
            pyproject = source_root / "pyproject.toml"
            pyproject.write_text("initial\n", encoding="utf-8")
            plan = _plan("locked")
            with patch(
                "experiments.manifest._git_provenance",
                return_value=dict(clean),
            ):
                ensure_batch(batch_dir, plan, source_root=source_root)

            pyproject.write_text("changed\n", encoding="utf-8")
            with patch(
                "experiments.manifest._git_provenance",
                return_value=dict(clean),
            ):
                with self.assertRaisesRegex(
                    ContractError,
                    "locked batch source file hash mismatch for: pyproject.toml",
                ):
                    ensure_batch(batch_dir, plan, source_root=source_root)


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