#!/usr/bin/env python3
"""P6 constrained-recovery probe; BUILT ONLY and Mos-gated.

DO NOT self-fire. Under Mos authorization only:
  python3 -I -S -B docs/compaction-refresh/probes/p6_constrained_recovery.py

The default three isolated runs launch the shipped daemon plus its production
observer transport on private sockets. The driver invokes the shipped recovery
command and adapter gate identity; it never resets broker state, mocks promote,
or taps a live model-output stream.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import importlib.util
import json
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path


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


def load_shipped_fragments():
    spec = importlib.util.spec_from_file_location("p6_shipped_fragments", FRAGMENTS)
    if spec is None or spec.loader is None:
        raise RuntimeError("shipped normative construction unavailable")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def request(socket_path: Path, value: dict[str, object]) -> dict[str, object]:
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(3.0)
        connection.connect(str(socket_path))
        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)
    if not response.endswith(b"\n") or response.count(b"\n") != 1:
        raise AssertionError(f"unframed broker reply: {bytes(response)!r}")
    reply = json.loads(response[:-1])
    if not isinstance(reply, dict):
        raise AssertionError("broker reply 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:
        if socket_path.exists():
            return
        if process.poll() is not None:
            output = process.stdout.read() if process.stdout is not None else ""
            raise RuntimeError(f"shipped daemon exited before READY: {output}")
        time.sleep(0.02)
    raise TimeoutError("shipped daemon did not create private probe socket")


def run_json(command: list[str], environment: dict[str, str], input_value: object | None = None) -> dict[str, object]:
    completed = subprocess.run(
        command,
        input=None if input_value is None else json.dumps(input_value),
        text=True,
        capture_output=True,
        env=environment,
        check=False,
    )
    if not completed.stdout.endswith("\n"):
        raise AssertionError(f"command omitted framed result: {completed.stderr!r}")
    reply = json.loads(completed.stdout)
    if not isinstance(reply, dict):
        raise AssertionError("command result is not an object")
    return reply


def gate_recovery(runtime: str, phase: str, environment: dict[str, str]) -> None:
    command = [sys.executable, "-I", "-S", "-B", str(GATE), "--runtime", runtime]
    if runtime == "claude":
        command.extend(["--recovery-command", str(RECOVERY_COMMAND)])
        recovery_invocation = (
            f"python3 {RECOVERY_COMMAND} begin --construction /tmp/p6.json "
            "--compaction-epoch 1 --request-epoch 1"
            if phase == "begin"
            else f"python3 {RECOVERY_COMMAND} complete"
        )
        value = {"tool_name": "Bash", "tool_input": {"command": recovery_invocation}}
    else:
        value = {"tool_name": "mosaic_context_recover"}
    completed = subprocess.run(command, input=json.dumps(value), text=True, capture_output=True, env=environment, check=False)
    if completed.returncode != 0:
        raise AssertionError(f"{runtime} recovery invocation remained gated: {completed.stderr!r}")


def record_production_observation(runtime: str, message: str, root: Path, environment: dict[str, str]) -> None:
    command = [sys.executable, "-I", "-S", "-B", str(OBSERVER_CLIENT), "--runtime", runtime]
    if runtime == "claude":
        transcript = root / "claude-transcript.jsonl"
        transcript.write_text(json.dumps({"message": {"role": "assistant", "content": message}}) + "\n", encoding="utf-8")
        payload = {"transcript_path": str(transcript)}
        command.append("--latest-entry")
    else:
        payload = {"latest_assistant_message": message}
    completed = subprocess.run(command, input=json.dumps(payload), text=True, capture_output=True, env=environment, check=False)
    if completed.returncode != 0:
        raise AssertionError(f"{runtime} production observer transport refused: {completed.stderr!r}")


def run_once(index: int, runtime: str) -> None:
    # Parity guard: drive the shipped command and the repaired adapter/observer
    # bytes, not a shadow receipt or promotion implementation.
    recovery_source = RECOVERY_COMMAND.read_text(encoding="utf-8")
    if '"action": "begin_recovery"' not in recovery_source or '"action": "complete_recovery"' not in recovery_source:
        raise AssertionError("P6 parity guard: recovery command no longer drives shipped broker entrypoints")
    gate_source = GATE.read_text(encoding="utf-8")
    if "--recovery-command" not in CLAUDE_SETTINGS.read_text(encoding="utf-8"):
        raise AssertionError("P6 parity guard: Claude recovery mapping is missing")
    if "_SHELL_ACTIVE" not in gate_source or "argv[1] != str(recovery_command)" not in gate_source:
        raise AssertionError("P6 parity guard: Claude mapping is not literal-only")
    if "const RECOVERY_TOOL = 'mosaic_context_recover'" not in PI_EXTENSION.read_text(encoding="utf-8"):
        raise AssertionError("P6 parity guard: Pi recovery tool mapping is missing")

    fragments = load_shipped_fragments()
    root = Path(tempfile.mkdtemp(prefix=f"mosaic-p6-recovery-{index}-"))
    os.chmod(root, 0o700)
    socket_path = root / "broker.sock"
    observer_socket = root / "observer.sock"
    state_path = root / "state.json"
    construction_path = root / "construction.json"
    content = b"P6 constrained recovery fixture\n"
    construction = {
        "manifest_version": 1,
        "generator_version": "p6-constrained-recovery",
        "fragments": [{
            "source_id": "authority/p6",
            "content_base64": base64.b64encode(content).decode("ascii"),
            "expected_sha256": hashlib.sha256(content).hexdigest(),
        }],
    }
    construction_path.write_text(json.dumps(construction), encoding="utf-8")
    os.chmod(construction_path, 0o600)
    process = subprocess.Popen(
        [sys.executable, "-I", "-S", "-B", str(DAEMON), "--socket", str(socket_path),
         "--state", str(state_path), "--observer-socket", str(observer_socket)],
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    try:
        wait_ready(process, socket_path)
        registered = request(socket_path, {"action": "register_anchor", "runtime_generation": 1})
        session_id = registered.get("session_id")
        if registered.get("ok") is not True or not isinstance(session_id, str):
            raise AssertionError(f"broker anchor registration failed: {registered!r}")
        built = fragments.build_payload_from_wire(construction)
        normal = request(socket_path, {
            "action": "begin_verification", "session_id": session_id, "runtime_generation": 1,
            "runtime": runtime, "construction": construction,
            "binding": {"compaction_epoch": index, "request_epoch": index + 100,
                        "h_source": built.h_source, "h_payload": built.h_payload, "schema_version": 1},
        })
        normal_challenge = normal.get("receipt_challenge")
        normal_receipt = normal.get("receipt")
        if not isinstance(normal_challenge, str) or not isinstance(normal_receipt, str):
            raise AssertionError("normal path did not mint a receipt challenge")
        environment = {
            **os.environ,
            "MOSAIC_LEASE_BROKER_SOCKET": str(socket_path),
            "MOSAIC_RECEIPT_OBSERVER_SOCKET": str(observer_socket),
            "MOSAIC_LEASE_SESSION_ID": session_id,
            "MOSAIC_RUNTIME_GENERATION": "1",
            "MOSAIC_LEASE_RUNTIME": runtime,
        }
        gate_recovery(runtime, "begin", environment)
        recovery = run_json([
            sys.executable, "-I", "-S", "-B", str(RECOVERY_COMMAND), "begin", "--construction", str(construction_path),
            "--compaction-epoch", str(index + 10), "--request-epoch", str(index + 110),
        ], environment)
        challenge = recovery.get("receipt_challenge")
        receipt = recovery.get("receipt")
        if recovery.get("state") != "PENDING_DELIVERY" or not isinstance(challenge, str) or not isinstance(receipt, str):
            raise AssertionError(f"recovery command did not drive pending delivery: {recovery!r}")
        if challenge == normal_challenge:
            raise AssertionError("recovery reused a normal-path challenge")

        # C4: production observer content is still exact-current-cycle only.
        record_production_observation(runtime, normal_receipt, root, environment)
        refused = run_json([sys.executable, "-I", "-S", "-B", str(RECOVERY_COMMAND), "complete"], environment)
        if refused.get("ok") is not False or refused.get("code") != "RECEIPT_MISMATCH":
            raise AssertionError(f"normal-path receipt replay was not refused: {refused!r}")

        gate_recovery(runtime, "begin", environment)
        recovery = run_json([
            sys.executable, "-I", "-S", "-B", str(RECOVERY_COMMAND), "begin", "--construction", str(construction_path),
            "--compaction-epoch", str(index + 20), "--request-epoch", str(index + 120),
        ], environment)
        receipt = recovery.get("receipt")
        if recovery.get("state") != "PENDING_DELIVERY" or not isinstance(receipt, str):
            raise AssertionError(f"fresh recovery retry did not pend: {recovery!r}")
        record_production_observation(runtime, receipt, root, environment)
        gate_recovery(runtime, "complete", environment)
        promoted = run_json([sys.executable, "-I", "-S", "-B", str(RECOVERY_COMMAND), "complete"], environment)
        if promoted.get("ok") is not True or promoted.get("state") != "VERIFIED":
            raise AssertionError(f"recovery consume-before-promote failed: {promoted!r}")
        replay = run_json([sys.executable, "-I", "-S", "-B", str(RECOVERY_COMMAND), "complete"], environment)
        if replay.get("ok") is not False or replay.get("code") != "INVALID_LEASE_TRANSITION":
            raise AssertionError(f"consumed recovery challenge re-promoted: {replay!r}")
    finally:
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=3.0)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()
        shutil.rmtree(root, ignore_errors=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs", type=int, default=3)
    arguments = parser.parse_args()
    if arguments.runs != 3:
        raise SystemExit("P6 requires exactly three isolated runs")
    for index, runtime in enumerate(("pi", "claude", "pi")):
        run_once(index, runtime)
    print("P6 constrained recovery probe PASS: 3 isolated shipped recovery-command runs")


if __name__ == "__main__":
    main()
