from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from turn_detection.provenance import (
    build_freeze_policy_binding,
    file_evidence,
    verify_freeze_manifest,
)

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


class FreezeManifestTest(unittest.TestCase):
    @staticmethod
    def _manifest(checkpoint: Path, config: Path) -> dict:
        checkpoint_evidence = file_evidence(checkpoint, ROOT)
        config_evidence = file_evidence(config, ROOT)
        manifest = {
            "format_version": 1,
            "status": "frozen_for_official_test",
            "checkpoint": dict(checkpoint_evidence),
            "training_config": dict(config_evidence),
            "threshold": 0.7,
            "controller": {"endpoint_threshold": 0.7},
            "official_test": {
                "dataset_id": "pipecat-ai/smart-turn-data-v3.2-test",
                "revision": "0500378e8ed6d38e37b016e24d261e8e6c6a6859",
            },
            "files": [checkpoint_evidence, config_evidence],
        }
        manifest["policy_binding"] = build_freeze_policy_binding(manifest)
        return manifest

    def test_bound_file_tampering_is_rejected(self) -> None:
        with tempfile.TemporaryDirectory(prefix="freeze-test-", dir=ROOT) as directory:
            root = Path(directory)
            checkpoint = root / "best.pt"
            config = root / "config.json"
            manifest_path = root / "frozen_manifest.json"
            checkpoint.write_bytes(b"checkpoint")
            config.write_text("{}", encoding="utf-8")
            manifest = self._manifest(checkpoint, config)
            manifest_path.write_text(json.dumps(manifest), encoding="utf-8")
            loaded = verify_freeze_manifest(
                manifest_path,
                ROOT,
                checkpoint_path=checkpoint,
                threshold=0.7,
            )
            self.assertEqual(loaded["threshold"], 0.7)

            config.write_text('{"changed": true}', encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "frozen file changed"):
                verify_freeze_manifest(manifest_path, ROOT)

    def test_checkpoint_pointer_swap_is_rejected_after_rebinding(self) -> None:
        with tempfile.TemporaryDirectory(prefix="freeze-pointer-test-", dir=ROOT) as directory:
            root = Path(directory)
            checkpoint = root / "best.pt"
            replacement = root / "replacement.pt"
            config = root / "config.json"
            manifest_path = root / "frozen_manifest.json"
            checkpoint.write_bytes(b"checkpoint")
            replacement.write_bytes(b"replacement")
            config.write_text("{}", encoding="utf-8")
            manifest = self._manifest(checkpoint, config)

            # Simulate an attempted bypass which updates the mutable pointer and
            # recomputes the outer policy digest, but does not alter bound files.
            manifest["checkpoint"] = file_evidence(replacement, ROOT)
            manifest["policy_binding"] = build_freeze_policy_binding(manifest)
            manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

            with self.assertRaisesRegex(ValueError, "not identical to hash-bound"):
                verify_freeze_manifest(
                    manifest_path,
                    ROOT,
                    checkpoint_path=replacement,
                    threshold=0.7,
                )

    def test_training_config_pointer_swap_is_rejected_after_rebinding(self) -> None:
        with tempfile.TemporaryDirectory(prefix="freeze-config-test-", dir=ROOT) as directory:
            root = Path(directory)
            checkpoint = root / "best.pt"
            config = root / "config.json"
            replacement = root / "replacement.json"
            manifest_path = root / "frozen_manifest.json"
            checkpoint.write_bytes(b"checkpoint")
            config.write_text("{}", encoding="utf-8")
            replacement.write_text('{"different": true}', encoding="utf-8")
            manifest = self._manifest(checkpoint, config)
            manifest["training_config"] = file_evidence(replacement, ROOT)
            manifest["policy_binding"] = build_freeze_policy_binding(manifest)
            manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

            with self.assertRaisesRegex(ValueError, "not identical to hash-bound"):
                verify_freeze_manifest(manifest_path, ROOT)

    def test_threshold_mutation_breaks_policy_binding(self) -> None:
        with tempfile.TemporaryDirectory(prefix="freeze-threshold-test-", dir=ROOT) as directory:
            root = Path(directory)
            checkpoint = root / "best.pt"
            config = root / "config.json"
            manifest_path = root / "frozen_manifest.json"
            checkpoint.write_bytes(b"checkpoint")
            config.write_text("{}", encoding="utf-8")
            manifest = self._manifest(checkpoint, config)
            manifest["threshold"] = 0.9
            manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

            with self.assertRaisesRegex(ValueError, "policy binding mismatch"):
                verify_freeze_manifest(manifest_path, ROOT, threshold=0.9)

    def test_threshold_must_match_bound_controller_and_evaluated_value(self) -> None:
        with tempfile.TemporaryDirectory(prefix="freeze-controller-test-", dir=ROOT) as directory:
            root = Path(directory)
            checkpoint = root / "best.pt"
            config = root / "config.json"
            manifest_path = root / "frozen_manifest.json"
            checkpoint.write_bytes(b"checkpoint")
            config.write_text("{}", encoding="utf-8")
            manifest = self._manifest(checkpoint, config)
            manifest["threshold"] = 0.9
            manifest["policy_binding"] = build_freeze_policy_binding(manifest)
            manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

            with self.assertRaisesRegex(ValueError, "controller endpoint threshold"):
                verify_freeze_manifest(manifest_path, ROOT, threshold=0.9)

            manifest["controller"]["endpoint_threshold"] = 0.9
            manifest["policy_binding"] = build_freeze_policy_binding(manifest)
            manifest_path.write_text(json.dumps(manifest), encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "evaluated threshold differs"):
                verify_freeze_manifest(manifest_path, ROOT, threshold=0.7)


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