#!/usr/bin/env python3
"""P5 Gate0 replay probe; BUILT ONLY, execution is Mos-gated.

Run only under fresh-executor authorization:
  python3 -I -S -B docs/compaction-refresh/probes/p5_receipt_replay.py

Each of the default three isolated runs launches the shipped lease-broker daemon
in a distinct private temporary directory.  This driver never changes broker
state directly and does not replace the promote gate: every transition is sent
over the daemon's real Unix socket.  It proves the shipped order is
PENDING_DELIVERY -> observe/evidence commit -> consume -> VERIFIED and that a
consumed challenge cannot be replayed or reopen/renew its lease.
"""

from __future__ import annotations

import argparse
import base64
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"
FRAGMENTS = TOOLS / "normative_fragments.py"


def load_shipped_fragments():
    if not FRAGMENTS.is_file():
        raise RuntimeError(f"shipped normative construction missing: {FRAGMENTS}")
    spec = importlib.util.spec_from_file_location("p5_shipped_normative_fragments", FRAGMENTS)
    if spec is None or spec.loader is None:
        raise RuntimeError("unable to load shipped normative construction")
    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}")
    parsed = json.loads(response[:-1])
    if not isinstance(parsed, dict):
        raise AssertionError(f"non-object broker reply: {parsed!r}")
    return parsed


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 expect_refused(reply: dict[str, object], code: str) -> None:
    if reply != {"ok": False, "code": code}:
        raise AssertionError(f"expected refusal {code}, got {reply!r}")


def run_once(index: int) -> str:
    fragments = load_shipped_fragments()
    root = Path(tempfile.mkdtemp(prefix=f"mosaic-p5-replay-{index}-"))
    os.chmod(root, 0o700)
    socket_path = root / "broker.sock"
    state_path = root / "state.json"
    observer_path = root / "test-observer.json"
    process = subprocess.Popen(
        [
            sys.executable, "-I", "-S", "-B", str(DAEMON), "--socket", str(socket_path),
            "--state", str(state_path), "--test-observer-file", str(observer_path),
        ],
        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})
        if registered.get("ok") is not True or not isinstance(registered.get("session_id"), str):
            raise AssertionError(f"registration failed: {registered!r}")
        session_id = registered["session_id"]
        construction = fragments.build_payload(
            manifest_version=1,
            generator_version="p5-replay-probe",
            fragments=[
                fragments.NormativeFragment(
                    "authority/probe",
                    b"P5 shipped transition driver\n",
                    "63537df1a6cb0d80195a96757ab11d629e5b5e1f23be167218b84cb195b1c1d6",
                ),
            ],
        )
        if construction.injectionDecision != "ACCEPTED" or not construction.promotion:
            raise AssertionError("shipped normative construction refused P5 fixture")
        binding = {
            "compaction_epoch": index,
            "request_epoch": index + 100,
            "h_source": construction.h_source,
            "h_payload": construction.h_payload,
            "schema_version": 1,
        }
        construction_request = {
            "manifest_version": 1,
            "generator_version": "p5-replay-probe",
            "fragments": [{
                "source_id": "authority/probe",
                "content_base64": base64.b64encode(b"P5 shipped transition driver\n").decode("ascii"),
                "expected_sha256": "63537df1a6cb0d80195a96757ab11d629e5b5e1f23be167218b84cb195b1c1d6",
            }],
        }
        pending = request(socket_path, {
            "action": "begin_verification",
            "session_id": session_id,
            "runtime_generation": 1,
            "runtime": "pi",
            "binding": binding,
            "construction": construction_request,
        })
        if pending.get("ok") is not True or pending.get("state") != "PENDING_VERIFICATION":
            raise AssertionError(f"shipped pending-delivery transition failed: {pending!r}")
        challenge = pending.get("receipt_challenge")
        receipt = pending.get("receipt")
        if not isinstance(challenge, str) or not isinstance(receipt, str):
            raise AssertionError(f"shipped broker did not mint a receipt challenge: {pending!r}")

        # Promotion before observation/evidence/consumption is forbidden.
        expect_refused(request(socket_path, {
            "action": "promote_lease",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt_challenge": challenge,
        }), "INVALID_LEASE_TRANSITION")

        observer_path.write_text(json.dumps({
            "session_id": session_id,
            "runtime_generation": 1,
            "latest_assistant_message": receipt,
        }), encoding="utf-8")
        os.chmod(observer_path, 0o600)
        observed = request(socket_path, {
            "action": "observe_receipt",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt_challenge": challenge,
        })
        if observed.get("ok") is not True or observed.get("state") != "PENDING_PROMOTION":
            raise AssertionError(f"shipped evidence transition failed: {observed!r}")
        durable = json.loads(state_path.read_text(encoding="utf-8"))
        evidence = durable["tokens"][challenge].get("evidence")
        if not isinstance(evidence, dict) or not isinstance(evidence.get("h_latest_assistant"), str):
            raise AssertionError("shipped receipt evidence was not committed before consume/promote")

        promoted = request(socket_path, {
            "action": "promote_lease",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt_challenge": challenge,
        })
        if promoted.get("ok") is not True or promoted.get("state") != "VERIFIED":
            raise AssertionError(f"shipped consume-before-promote transition failed: {promoted!r}")

        # T25/T28: the actual consumed challenge, re-presented through the
        # shipped daemon, can neither be observed again nor re-promote/reopen.
        expect_refused(request(socket_path, {
            "action": "observe_receipt",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt_challenge": challenge,
        }), "RECEIPT_REPLAY")
        expect_refused(request(socket_path, {
            "action": "promote_lease",
            "session_id": session_id,
            "runtime_generation": 1,
            "receipt_challenge": challenge,
        }), "RECEIPT_REPLAY")
        return challenge
    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("P5 requires exactly three isolated runs")
    challenges = [run_once(index) for index in range(arguments.runs)]
    if len(set(challenges)) != arguments.runs:
        raise AssertionError("separate shipped cycles did not mint unique challenges")
    print("P5 receipt replay probe PASS: 3 isolated shipped-daemon runs")


if __name__ == "__main__":
    main()
