#!/usr/bin/env python3
"""RED-first contracts for the shipped receipt challenge and observer seam."""

from __future__ import annotations

import base64
import copy
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_receipt_daemon", DAEMON_PATH)
FRAGMENTS = load_module("lease_broker_normative_fragments", FRAGMENTS_PATH)


class BrokerFixture(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.broker = DAEMON.Broker(DAEMON.StateStore(root / "state.json"))
        registered = self.broker.handle(self.peer, {
            "action": "register_anchor",
            "runtime_generation": 7,
        })
        self.session_id = registered["session_id"]
        self.assertIsInstance(self.session_id, str)

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

    def construction(self) -> tuple[dict[str, object], dict[str, object]]:
        content = b"Constitution\n"
        expected_sha256 = hashlib.sha256(content).hexdigest()
        construction = {
            "manifest_version": 1,
            "generator_version": "wi5-receipt-test",
            "fragments": [{
                "source_id": "authority/constitution",
                "content_base64": base64.b64encode(content).decode("ascii"),
                "expected_sha256": expected_sha256,
            }],
        }
        result = FRAGMENTS.build_payload(
            manifest_version=construction["manifest_version"],
            generator_version=construction["generator_version"],
            fragments=[FRAGMENTS.NormativeFragment("authority/constitution", content, expected_sha256)],
        )
        self.assertEqual(result.injectionDecision, "ACCEPTED")
        self.assertTrue(result.promotion)
        return construction, {
            "compaction_epoch": 3,
            "request_epoch": 8,
            "h_source": result.h_source,
            "h_payload": result.h_payload,
            "schema_version": 1,
        }

    def begin(self, binding: dict[str, object], construction: dict[str, object]) -> dict[str, object]:
        response = self.broker.handle(self.peer, {
            "action": "begin_verification",
            "session_id": self.session_id,
            "runtime_generation": 7,
            "runtime": "pi",
            "binding": binding,
            "construction": construction,
        })
        self.assertEqual(response["state"], DAEMON.LEASE_PENDING)
        self.assertIsInstance(response.get("receipt_challenge"), str)
        self.assertIsInstance(response.get("receipt"), str)
        return response


class BuildPayloadAdmissionTest(BrokerFixture):
    def test_b3_forged_h_source_or_h_payload_is_refused_against_shipped_build_payload(self) -> None:
        construction, trusted = self.construction()
        for field in ("h_source", "h_payload"):
            with self.subTest(field=field):
                forged = dict(trusted)
                forged[field] = "f" * 64
                with self.assertRaisesRegex(DAEMON.BrokerFailure, "PAYLOAD_BINDING_MISMATCH"):
                    self.begin(forged, construction)


class ReceiptObserverTest(BrokerFixture):
    def setUp(self) -> None:
        super().setUp()
        observers = load_module("lease_broker_test_observer", OBSERVER_PATH)
        self.observer = observers.TestReceiptObserver()
        self.broker = DAEMON.Broker(self.broker.store, observer=self.observer)

    def record(self, message: str) -> None:
        self.observer.record_latest_assistant_message(self.session_id, 7, message)

    def observe(self, challenge: str, **untrusted: object) -> dict[str, object]:
        return self.broker.handle(self.peer, {
            "action": "observe_receipt",
            "session_id": self.session_id,
            "runtime_generation": 7,
            "receipt_challenge": challenge,
            **untrusted,
        })

    def promote(self, challenge: str) -> dict[str, object]:
        return self.broker.handle(self.peer, {
            "action": "promote_lease",
            "session_id": self.session_id,
            "runtime_generation": 7,
            "receipt_challenge": challenge,
        })

    def test_b2_echoed_request_observation_is_refused_but_observer_source_promotes(self) -> None:
        construction, binding = self.construction()
        cycle = self.begin(binding, construction)
        challenge = cycle["receipt_challenge"]
        receipt = cycle["receipt"]

        with self.assertRaisesRegex(DAEMON.BrokerFailure, "INVALID_RECEIPT"):
            self.observe(challenge, latest_assistant_message=receipt)
        with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_OBSERVATION_UNAVAILABLE"):
            self.observe(challenge)

        self.record(receipt)
        self.assertEqual(self.observe(challenge)["state"], DAEMON.LEASE_PENDING_PROMOTION)
        self.assertEqual(self.promote(challenge)["state"], DAEMON.LEASE_VERIFIED)
        self.assertEqual(self.broker.handle(self.peer, {
            "action": "authorize_tool",
            "session_id": self.session_id,
            "runtime_generation": 7,
            "runtime": "pi",
            "tool_name": "bash",
        })["decision"], "allow")

    def test_rejected_begin_keeps_revoke_first_fence_for_all_construction_refusals(self) -> None:
        construction, binding = self.construction()
        refusal_cases = {
            "INVALID_CONSTRUCTION": {"bad": "construction"},
            "PAYLOAD_CONSTRUCTION_REFUSED": {
                **construction,
                "fragments": [{
                    **construction["fragments"][0],
                    "expected_sha256": "0" * 64,
                }],
            },
            "PAYLOAD_BINDING_MISMATCH": None,
        }
        for expected_code, rejected_construction in refusal_cases.items():
            with self.subTest(expected_code=expected_code):
                verified = self.begin(binding, construction)
                self.record(verified["receipt"])
                self.observe(verified["receipt_challenge"])
                self.assertEqual(self.promote(verified["receipt_challenge"])["state"], DAEMON.LEASE_VERIFIED)

                rejected_binding = copy.deepcopy(binding)
                if expected_code == "PAYLOAD_BINDING_MISMATCH":
                    rejected_binding["h_payload"] = "f" * 64
                    rejected_construction = construction
                with self.assertRaisesRegex(DAEMON.BrokerFailure, expected_code):
                    self.begin(rejected_binding, rejected_construction)

                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": 7,
                    "runtime": "pi",
                    "tool_name": "bash",
                })
                self.assertEqual(denied["decision"], "deny")
                self.assertEqual(denied["state"], DAEMON.LEASE_UNVERIFIED)

    def test_t26_stale_epoch_receipt_cannot_promote_against_shipped_binding(self) -> None:
        construction, stale_binding = self.construction()
        stale = self.begin(stale_binding, construction)
        current_binding = dict(stale_binding)
        current_binding["compaction_epoch"] = 4
        current_binding["request_epoch"] = 9
        current = self.begin(current_binding, construction)
        self.assertNotEqual(stale["receipt_challenge"], current["receipt_challenge"])

        self.record(stale["receipt"])
        with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_MISMATCH"):
            self.observe(current["receipt_challenge"])
        with self.assertRaisesRegex(DAEMON.BrokerFailure, "INVALID_LEASE_TRANSITION"):
            self.promote(current["receipt_challenge"])

    def test_non_ascii_observation_is_mismatch_and_broker_keeps_serving(self) -> None:
        construction, binding = self.construction()
        cycle = self.begin(binding, construction)
        self.record("I refuse—this is not the receipt")

        with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_MISMATCH"):
            self.observe(cycle["receipt_challenge"])

        denied = self.broker.handle(self.peer, {
            "action": "authorize_tool",
            "session_id": self.session_id,
            "runtime_generation": 7,
            "runtime": "pi",
            "tool_name": "bash",
        })
        self.assertEqual(denied["decision"], "deny")
        self.assertEqual(denied["state"], DAEMON.LEASE_UNVERIFIED)

    def test_t29_altered_model_hash_cannot_promote_against_shipped_binding(self) -> None:
        construction, binding = self.construction()
        cycle = self.begin(binding, construction)
        expected = cycle["receipt"]
        altered_hash = "f" * 64
        self.assertNotEqual(altered_hash, cycle["binding"]["h_payload"])
        altered = expected.replace(cycle["binding"]["h_payload"], altered_hash, 1)
        self.assertNotEqual(altered, expected)

        self.record(altered)
        with self.assertRaisesRegex(DAEMON.BrokerFailure, "RECEIPT_MISMATCH"):
            self.observe(cycle["receipt_challenge"])
        with self.assertRaisesRegex(DAEMON.BrokerFailure, "INVALID_LEASE_TRANSITION"):
            self.promote(cycle["receipt_challenge"])


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