#!/usr/bin/env python3
"""RED-first WI-6 contracts against the shipped constrained recovery entrypoint."""

from __future__ import annotations

import base64
import hashlib
import importlib.util
import os
import tempfile
import unittest
from pathlib import Path


TOOLS = Path(__file__).parents[2] / "framework/tools/lease-broker"
DAEMON_PATH = TOOLS / "daemon.py"
FRAGMENTS_PATH = TOOLS / "normative_fragments.py"
OBSERVER_PATH = TOOLS / "receipt_observer.py"


def load_module(name: str, path: Path):
    assert path.is_file(), f"shipped module is missing: {path}"
    spec = importlib.util.spec_from_file_location(name, path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"unable to load {name}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


DAEMON = load_module("lease_broker_recovery_daemon", DAEMON_PATH)
FRAGMENTS = load_module("lease_broker_recovery_fragments", FRAGMENTS_PATH)
OBSERVER = load_module("lease_broker_recovery_observer", OBSERVER_PATH)


class RecoveryFixture(unittest.TestCase):
    def setUp(self) -> None:
        self.temporary = tempfile.TemporaryDirectory()
        root = Path(self.temporary.name)
        os.chmod(root, 0o700)
        self.peer = (os.getpid(), os.getuid(), os.getgid())
        self.observer = OBSERVER.TestReceiptObserver()
        self.broker = DAEMON.Broker(DAEMON.StateStore(root / "state.json"), observer=self.observer)
        registered = self.broker.handle(self.peer, {
            "action": "register_anchor",
            "runtime_generation": 11,
        })
        self.session_id = registered["session_id"]

    def tearDown(self) -> None:
        self.temporary.cleanup()

    def construction_and_binding(self) -> tuple[dict[str, object], dict[str, object]]:
        content = b"Mosaic recovery authority\n"
        construction = {
            "manifest_version": 1,
            "generator_version": "wi6-recovery-test",
            "fragments": [{
                "source_id": "authority/recovery",
                "content_base64": base64.b64encode(content).decode("ascii"),
                "expected_sha256": hashlib.sha256(content).hexdigest(),
            }],
        }
        built = FRAGMENTS.build_payload_from_wire(construction)
        self.assertEqual(built.injectionDecision, "ACCEPTED")
        self.assertTrue(built.promotion)
        return construction, {
            "compaction_epoch": 17,
            "request_epoch": 23,
            "h_source": built.h_source,
            "h_payload": built.h_payload,
            "schema_version": 1,
        }

    def begin_normal(self) -> dict[str, object]:
        construction, binding = self.construction_and_binding()
        return self.broker.handle(self.peer, {
            "action": "begin_verification",
            "session_id": self.session_id,
            "runtime_generation": 11,
            "runtime": "pi",
            "binding": binding,
            "construction": construction,
        })

    def begin_recovery(self) -> dict[str, object]:
        construction, binding = self.construction_and_binding()
        return self.broker.handle(self.peer, {
            "action": "begin_recovery",
            "session_id": self.session_id,
            "runtime_generation": 11,
            "runtime": "pi",
            "binding": binding,
            "construction": construction,
        })

    def complete_recovery(self) -> dict[str, object]:
        return self.broker.handle(self.peer, {
            "action": "complete_recovery",
            "session_id": self.session_id,
            "runtime_generation": 11,
        })

    def assert_recovery_refused_unverified(self, code: str) -> None:
        with self.assertRaisesRegex(DAEMON.BrokerFailure, code):
            self.complete_recovery()
        self.assertEqual(self.broker.leases[self.session_id]["state"], DAEMON.LEASE_UNVERIFIED)
        denied = self.broker.handle(self.peer, {
            "action": "authorize_tool",
            "session_id": self.session_id,
            "runtime_generation": 11,
            "runtime": "pi",
            "tool_name": "bash",
        })
        self.assertEqual(denied["decision"], "deny")


class ConstrainedRecoveryContractTest(RecoveryFixture):
    def test_recovery_mints_a_fresh_challenge_distinct_from_normal_path(self) -> None:
        normal = self.begin_normal()
        recovery = self.begin_recovery()

        self.assertEqual(recovery["state"], "PENDING_DELIVERY")
        self.assertNotEqual(normal["receipt_challenge"], recovery["receipt_challenge"])
        self.assertIn(recovery["receipt_challenge"], recovery["receipt"])

    def test_c4_normal_path_receipt_cannot_be_replayed_through_recovery(self) -> None:
        normal = self.begin_normal()
        recovery = self.begin_recovery()
        self.assertNotEqual(normal["receipt_challenge"], recovery["receipt_challenge"])

        self.observer.record_latest_assistant_message(self.session_id, 11, normal["receipt"])
        self.assert_recovery_refused_unverified("RECEIPT_MISMATCH")

    def test_t27_observable_partial_delivery_variants_never_promote(self) -> None:
        variants = {
            "absent": None,
            "malformed": "MOSAIC-RECEIPT{malformed}",
            "prefix-truncated": None,
            "observable-adapter-mutation": None,
            # Tail-only is represented only by this concrete malformed/incomplete
            # delivery. It is not a category-wide tail-only detection claim.
            "tail-only-malformed": "H_payload=tail-only",
        }
        for name, observed in variants.items():
            with self.subTest(name=name):
                recovery = self.begin_recovery()
                if name == "prefix-truncated":
                    observed = recovery["receipt"][:-1]
                elif name == "observable-adapter-mutation":
                    observed = recovery["receipt"].replace("H_payload=", "H_payload=0", 1)
                if observed is not None:
                    self.observer.record_latest_assistant_message(self.session_id, 11, observed)
                self.assert_recovery_refused_unverified(
                    "RECEIPT_OBSERVATION_UNAVAILABLE" if observed is None else "RECEIPT_MISMATCH"
                )

    def test_negative_capability_tail_preserving_middle_drop_is_not_receipt_detectable(self) -> None:
        recovery = self.begin_recovery()

        # The observer seam receives the exact terminal message, not the delivered
        # payload bytes. A middle drop that preserves this tail is therefore T-C
        # and deliberately NOT represented as receipt-detectable; WI-7 server-side
        # evidence owns that residual. This is not an assertion that it is caught.
        self.observer.record_latest_assistant_message(self.session_id, 11, recovery["receipt"])
        promoted = self.complete_recovery()
        self.assertEqual(promoted["state"], DAEMON.LEASE_VERIFIED)


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