#!/usr/bin/env python3
"""RED-first B1/B2 runtime-boundary contracts for constrained recovery.

Every runtime process in these tests is a fresh child against a private daemon
and Unix sockets. They do not activate a live Mosaic daemon or model stream.
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
import socket
import subprocess
import sys
import tempfile
import time
import unittest
from pathlib import Path


TOOLS = Path(__file__).parents[2] / "framework/tools/lease-broker"
FRAMEWORK = Path(__file__).parents[2] / "framework"
DAEMON = TOOLS / "daemon.py"
GATE = TOOLS / "mutator-gate.py"
RECOVERY = TOOLS / "recover-context.py"
OBSERVER_CLIENT = TOOLS / "receipt-observer-client.py"
CLAUDE_SETTINGS = FRAMEWORK / "runtime/claude/settings.json"
PI_EXTENSION = FRAMEWORK / "runtime/pi/mosaic-extension.ts"


def request(socket_path: Path, value: dict[str, object]) -> dict[str, object]:
    deadline = time.monotonic() + 5.0
    while True:
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
            connection.settimeout(3.0)
            try:
                connection.connect(str(socket_path))
            except ConnectionRefusedError:
                if time.monotonic() >= deadline:
                    raise
                time.sleep(0.02)
                continue
            connection.sendall((json.dumps(value, separators=(",", ":")) + "\n").encode())
            connection.shutdown(socket.SHUT_WR)
            response = bytearray()
            while True:
                chunk = connection.recv(4096)
                if not chunk:
                    break
                response.extend(chunk)
        break
    if not response.endswith(b"\n") or response.count(b"\n") != 1:
        raise AssertionError(f"unframed broker response: {bytes(response)!r}")
    reply = json.loads(response[:-1])
    if not isinstance(reply, dict):
        raise AssertionError("broker response is not an object")
    return reply


def wait_ready(process: subprocess.Popen[str], socket_path: Path) -> None:
    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        try:
            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe:
                probe.settimeout(0.25)
                probe.connect(str(socket_path))
            return
        except (ConnectionRefusedError, FileNotFoundError):
            pass
        if process.poll() is not None:
            output = process.stdout.read() if process.stdout is not None else ""
            raise RuntimeError(f"daemon exited before READY: {output}")
        time.sleep(0.02)
    raise TimeoutError("daemon did not create private broker socket")


class RuntimeBoundaryFixture(unittest.TestCase):
    def setUp(self) -> None:
        self.temporary = tempfile.TemporaryDirectory()
        self.root = Path(self.temporary.name)
        os.chmod(self.root, 0o700)
        self.socket = self.root / "broker.sock"
        self.observer_socket = self.root / "observer.sock"
        self.state = self.root / "state.json"
        self.children: list[subprocess.Popen[str]] = []

    def tearDown(self) -> None:
        for child in self.children:
            if child.poll() is None:
                child.terminate()
                try:
                    child.wait(timeout=3.0)
                except subprocess.TimeoutExpired:
                    child.kill()
                    child.wait()
            if child.stdout is not None:
                child.stdout.close()
        self.temporary.cleanup()

    def start_daemon(self, *, production_observer: bool) -> None:
        arguments = [sys.executable, "-I", "-S", "-B", str(DAEMON), "--socket", str(self.socket), "--state", str(self.state)]
        if production_observer:
            arguments.extend(["--observer-socket", str(self.observer_socket)])
        process = subprocess.Popen(
            arguments,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        self.children.append(process)
        wait_ready(process, self.socket)

    def register(self) -> str:
        reply = request(self.socket, {"action": "register_anchor", "runtime_generation": 1})
        session_id = reply.get("session_id")
        self.assertTrue(reply.get("ok"))
        self.assertIsInstance(session_id, str)
        return session_id

    def environment(self, session_id: str) -> dict[str, str]:
        return {
            **os.environ,
            "MOSAIC_LEASE_BROKER_SOCKET": str(self.socket),
            "MOSAIC_RECEIPT_OBSERVER_SOCKET": str(self.observer_socket),
            "MOSAIC_LEASE_SESSION_ID": session_id,
            "MOSAIC_RUNTIME_GENERATION": "1",
            "MOSAIC_LEASE_RUNTIME": "pi",
        }

    def construction(self) -> Path:
        content = b"WI-6 B2 production observer fixture\n"
        path = self.root / "construction.json"
        path.write_text(json.dumps({
            "manifest_version": 1,
            "generator_version": "wi6-repair-runtime-boundary",
            "fragments": [{
                "source_id": "authority/wi6-repair",
                "content_base64": base64.b64encode(content).decode("ascii"),
                "expected_sha256": hashlib.sha256(content).hexdigest(),
            }],
        }), encoding="utf-8")
        os.chmod(path, 0o600)
        return path


class RecoveryRuntimeBoundaryTest(RuntimeBoundaryFixture):
    def test_b1_claude_exact_recovery_command_is_invocable_unverified_but_bash_is_not(self) -> None:
        self.start_daemon(production_observer=False)
        session_id = self.register()
        environment = self.environment(session_id)
        recovery_command = f"python3 {RECOVERY} complete"
        mapped = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(GATE), "--runtime", "claude", "--recovery-command", str(RECOVERY)],
            input=json.dumps({"tool_name": "Bash", "tool_input": {"command": recovery_command}}),
            text=True,
            capture_output=True,
            env={**environment, "MOSAIC_LEASE_RUNTIME": "claude"},
            check=False,
        )
        self.assertEqual(mapped.returncode, 0, mapped.stderr)
        mapped_begin = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(GATE), "--runtime", "claude", "--recovery-command", str(RECOVERY)],
            input=json.dumps({
                "tool_name": "Bash",
                "tool_input": {
                    "command": (
                        f"python3 {RECOVERY} begin --construction /tmp/construction.json "
                        "--compaction-epoch 1 --request-epoch 1"
                    )
                },
            }),
            text=True,
            capture_output=True,
            env={**environment, "MOSAIC_LEASE_RUNTIME": "claude"},
            check=False,
        )
        self.assertEqual(mapped_begin.returncode, 0, mapped_begin.stderr)
        ordinary_bash = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(GATE), "--runtime", "claude", "--recovery-command", str(RECOVERY)],
            input=json.dumps({"tool_name": "Bash", "tool_input": {"command": "echo not-recovery"}}),
            text=True,
            capture_output=True,
            env={**environment, "MOSAIC_LEASE_RUNTIME": "claude"},
            check=False,
        )
        self.assertEqual(ordinary_bash.returncode, 2)

    def test_b1_pi_registers_only_the_broker_exempt_recovery_tool(self) -> None:
        extension = PI_EXTENSION.read_text(encoding="utf-8")
        self.assertIn("const RECOVERY_TOOL = 'mosaic_context_recover'", extension)
        self.assertIn("name: RECOVERY_TOOL", extension)
        self.assertIn("checkPiMutatorGate(RECOVERY_TOOL)", extension)
        self.assertNotIn("toolName === 'bash' ? RECOVERY_TOOL", extension)

    def test_b2_production_observer_promotes_over_private_transport_and_rejects_broker_supplied_message(self) -> None:
        self.start_daemon(production_observer=True)
        session_id = self.register()
        environment = self.environment(session_id)
        begin = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(RECOVERY), "begin", "--construction", str(self.construction()), "--compaction-epoch", "1", "--request-epoch", "1"],
            text=True,
            capture_output=True,
            env=environment,
            check=False,
        )
        self.assertEqual(begin.returncode, 0, begin.stderr)
        cycle = json.loads(begin.stdout)
        receipt = cycle.get("receipt")
        self.assertIsInstance(receipt, str)

        # S1: production transport is not a broker request field. The public
        # broker endpoint keeps rejecting caller-supplied assistant evidence.
        rejected_begin = request(self.socket, {
            "action": "begin_recovery",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt": receipt,
        })
        self.assertEqual(rejected_begin, {"ok": False, "code": "INVALID_RECOVERY_REQUEST"})
        rejected_observe = request(self.socket, {
            "action": "observe_receipt",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt_challenge": cycle["receipt_challenge"],
            "latest_assistant_message": receipt,
        })
        self.assertEqual(rejected_observe, {"ok": False, "code": "INVALID_RECEIPT"})
        rejected_complete = request(self.socket, {
            "action": "complete_recovery",
            "session_id": session_id,
            "runtime_generation": 1,
            "latest_assistant_message": receipt,
        })
        self.assertEqual(rejected_complete, {"ok": False, "code": "INVALID_RECOVERY_REQUEST"})

        recorded = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(OBSERVER_CLIENT), "--runtime", "pi"],
            input=json.dumps({"latest_assistant_message": receipt}),
            text=True,
            capture_output=True,
            env=environment,
            check=False,
        )
        self.assertEqual(recorded.returncode, 0, recorded.stderr)
        complete = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(RECOVERY), "complete"],
            text=True,
            capture_output=True,
            env=environment,
            check=False,
        )
        self.assertEqual(complete.returncode, 0, complete.stderr)
        self.assertEqual(json.loads(complete.stdout).get("state"), "VERIFIED")

        # The independent Claude transport selects one latest assistant entry
        # from its hook transcript; it is not a Pi/message_end fallback.
        claude_environment = {**environment, "MOSAIC_LEASE_RUNTIME": "claude"}
        claude_begin = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(RECOVERY), "begin", "--construction", str(self.construction()), "--compaction-epoch", "2", "--request-epoch", "2"],
            text=True,
            capture_output=True,
            env=claude_environment,
            check=False,
        )
        self.assertEqual(claude_begin.returncode, 0, claude_begin.stderr)
        claude_receipt = json.loads(claude_begin.stdout)["receipt"]
        transcript = self.root / "claude-transcript.jsonl"
        transcript.write_text(
            json.dumps({"message": {"role": "assistant", "content": claude_receipt}}) + "\n",
            encoding="utf-8",
        )
        claude_recorded = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(OBSERVER_CLIENT), "--runtime", "claude", "--latest-entry"],
            input=json.dumps({"transcript_path": str(transcript)}),
            text=True,
            capture_output=True,
            env=claude_environment,
            check=False,
        )
        self.assertEqual(claude_recorded.returncode, 0, claude_recorded.stderr)
        claude_complete = subprocess.run(
            [sys.executable, "-I", "-S", "-B", str(RECOVERY), "complete"],
            text=True,
            capture_output=True,
            env=claude_environment,
            check=False,
        )
        self.assertEqual(claude_complete.returncode, 0, claude_complete.stderr)
        self.assertEqual(json.loads(claude_complete.stdout).get("state"), "VERIFIED")


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