#!/usr/bin/env python3
"""D4-only same-PID runtime-generation revocation harness.

AUTHORING NOTE: this file is intentionally not executed until the separately
ratified FIRE authorization.  When run later, every invocation creates its own
/tmp fixture and launches the real Pi RPC runtime with only the D4 extension
and ``p3_generation_broker.py``.  It does not use the broader Gate0 runner.
"""

from __future__ import annotations

import argparse
import json
import os
import queue
import signal
import socket
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path
from typing import Any, Callable

HERE = Path(__file__).resolve().parent
# WI-2 is deliberately kept in its reviewed worktree until the release package
# contains the gated launcher.  This harness pins the source it will execute;
# falling back to the released `mosaic` binary is forbidden.
GATED_WI_ROOT = HERE.parents[3].parent / "stack-827-probe3"
GATED_WI_HEAD = "abd2791f59b3f06f46dd08e55298ced72f6aa7c2"
GATED_LAUNCHER = GATED_WI_ROOT / "packages/mosaic/framework/tools/lease-broker/launch-runtime.py"


class PiRpc:
    """Small JSON-RPC client for an isolated real Pi process."""

    def __init__(self, command: list[str], cwd: Path, env: dict[str, str]) -> None:
        self.process = subprocess.Popen(
            command,
            cwd=cwd,
            env=env,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
            start_new_session=True,
        )
        self.events: queue.Queue[dict[str, Any]] = queue.Queue()
        self.stderr_lines: list[str] = []
        threading.Thread(target=self._read_stdout, daemon=True).start()
        threading.Thread(target=self._read_stderr, daemon=True).start()

    def _read_stdout(self) -> None:
        assert self.process.stdout is not None
        for line in self.process.stdout:
            try:
                self.events.put(json.loads(line))
            except json.JSONDecodeError:
                continue

    def _read_stderr(self) -> None:
        assert self.process.stderr is not None
        for line in self.process.stderr:
            self.stderr_lines.append(line.rstrip("\n"))

    def send(self, payload: dict[str, object]) -> None:
        assert self.process.stdin is not None
        self.process.stdin.write(json.dumps(payload) + "\n")
        self.process.stdin.flush()

    def wait(
        self,
        predicate: Callable[[dict[str, Any]], bool],
        description: str,
        timeout: float = 180,
    ) -> dict[str, Any]:
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            if self.process.poll() is not None and self.events.empty():
                detail = " | ".join(self.stderr_lines[-5:])
                raise RuntimeError(
                    f"Pi exited {self.process.returncode} while waiting for {description}: {detail}"
                )
            try:
                event = self.events.get(timeout=0.2)
            except queue.Empty:
                continue
            if predicate(event):
                return event
        raise TimeoutError(f"timed out waiting for {description}")

    def response(self, request_id: str, timeout: float = 180) -> dict[str, Any]:
        return self.wait(
            lambda event: event.get("type") == "response" and event.get("id") == request_id,
            f"response {request_id}",
            timeout,
        )

    def prompt_and_settle(self, request_id: str, message: str) -> None:
        self.send({"id": request_id, "type": "prompt", "message": message})
        response = self.response(request_id)
        if not response.get("success"):
            raise RuntimeError(f"prompt rejected: {response}")
        self.wait(
            lambda event: event.get("type") == "agent_settled",
            f"agent_settled {request_id}",
        )

    def close(self) -> None:
        if self.process.poll() is None:
            try:
                os.killpg(self.process.pid, signal.SIGTERM)
            except ProcessLookupError:
                pass
            try:
                self.process.wait(timeout=8)
            except subprocess.TimeoutExpired:
                os.killpg(self.process.pid, signal.SIGKILL)
                self.process.wait(timeout=5)


def wait_path(path: Path, timeout: float = 20) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if path.exists():
            return
        time.sleep(0.05)
    raise TimeoutError(f"timed out waiting for {path}")


def request(path: Path, payload: dict[str, object]) -> dict[str, Any]:
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn:
        conn.connect(str(path))
        conn.sendall((json.dumps(payload) + "\n").encode())
        reply = conn.makefile("r", encoding="utf-8").readline()
    return json.loads(reply)


def jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text().splitlines() if line]


def write_extension(path: Path) -> None:
    """Write the minimal Pi lifecycle bridge into the isolated fixture only."""

    path.write_text(
        """import type { ExtensionAPI } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import { appendFileSync, readFileSync } from 'node:fs';
import net from 'node:net';

const socketPath = process.env['D4_GENERATION_SOCKET'];
const logPath = process.env['D4_PI_LOG'];

function starttime(): number {
  const text = readFileSync(`/proc/${process.pid}/stat`, 'utf8');
  const close = text.lastIndexOf(')');
  return Number(text.slice(close + 2).trim().split(/\\s+/)[19]);
}

function log(event: string, details: Record<string, unknown> = {}): void {
  if (!logPath) return;
  appendFileSync(logPath, `${JSON.stringify({ event, pid: process.pid, starttime_ticks: starttime(), ...details })}\\n`);
}

function broker(payload: Record<string, unknown>): Promise<Record<string, unknown>> {
  if (!socketPath) return Promise.reject(new Error('D4_GENERATION_SOCKET is required'));
  return new Promise((resolve, reject) => {
    const connection = net.createConnection(socketPath);
    let buffer = '';
    connection.setEncoding('utf8');
    connection.on('connect', () => connection.write(`${JSON.stringify(payload)}\\n`));
    connection.on('data', (chunk) => {
      buffer += chunk;
      const newline = buffer.indexOf('\\n');
      if (newline < 0) return;
      connection.end();
      resolve(JSON.parse(buffer.slice(0, newline)) as Record<string, unknown>);
    });
    connection.on('error', reject);
  });
}

export default function register(pi: ExtensionAPI): void {
  pi.on('session_start', async (event) => {
    const lifecycle = await broker({ action: 'lifecycle', phase: 'start', reason: event.reason });
    log('session_start', { reason: event.reason, lifecycle });
    if (event.reason === 'reload') {
      const generation = lifecycle['new_generation'];
      if (typeof generation !== 'number') throw new Error('broker did not return new_generation');
      const current = await broker({ action: 'authorize-probe', generation });
      const superseded = await broker({ action: 'authorize-probe', generation: generation - 1 });
      log('d4_generation_authorization', { generation, current, superseded });
    }
  });

  pi.on('session_shutdown', async (event) => {
    const lifecycle = await broker({ action: 'lifecycle', phase: 'shutdown', reason: event.reason });
    log('session_shutdown', { reason: event.reason, lifecycle });
  });

  pi.registerTool({
    name: 'd4_fixture_promote',
    label: 'D4 Fixture Promotion',
    description: 'Promotes only the fixture lease needed for the D4 revocation check.',
    parameters: Type.Object({}),
    async execute() {
      // the promotion step is a D4 test fixture, not a P2 evidence-gathering authorization.
      const promotion = await broker({ action: 'promote-probe' });
      log('fixture_promotion', { promotion });
      return { content: [{ type: 'text', text: 'D4 fixture promotion complete' }] };
    },
  });

  pi.registerCommand('d4-reload', {
    description: 'D4-only same-PID reload boundary.',
    handler: async (_args, context) => {
      await context.reload();
    },
  });
}
""",
        encoding="utf-8",
    )


def gated_launcher_precondition(root: Path, socket_path: Path, environment: dict[str, str]) -> None:
    """Refuse to launch unless the exact WI launcher binds this run's fixture."""

    if environment.get("MOSAIC_LEASE_BROKER_SOCKET") != str(socket_path):
        raise RuntimeError("D4 precondition: lease broker socket is not this fixture")
    if socket_path.parent != root or root.parent != Path(tempfile.gettempdir()):
        raise RuntimeError("D4 precondition: fixture socket is outside this run's temporary root")
    try:
        head = subprocess.check_output(
            ["git", "-C", str(GATED_WI_ROOT), "rev-parse", "HEAD"], text=True
        ).strip()
        launcher_source = GATED_LAUNCHER.read_text(encoding="utf-8")
    except (OSError, subprocess.CalledProcessError) as error:
        raise RuntimeError("D4 precondition: gated WI launcher is unavailable") from error
    if head != GATED_WI_HEAD:
        raise RuntimeError(f"D4 precondition: gated WI head mismatch: {head}")

    register = launcher_source.find('"action": "register_anchor"')
    execute = launcher_source.find("execute(command[0], command, environment)")
    fixture_socket = launcher_source.find('source_environment["MOSAIC_LEASE_BROKER_SOCKET"]')
    if register < 0 or execute < 0 or register >= execute or fixture_socket < 0:
        raise RuntimeError("D4 precondition: launcher is not register-before-exec fixture-bound")


def assert_d4(records: list[dict[str, Any]]) -> dict[str, object]:
    lifecycle = [record for record in records if record.get("event") == "runtime_generation_bump"]
    promotion = next(record for record in records if record.get("event") == "probe_lease_promoted")
    launcher_registration = next(
        record for record in records if record.get("event") == "lease_anchor_registered"
    )
    reload_revoke = next(
        record
        for record in lifecycle
        if record.get("reason") == "reload" and record.get("phase") == "shutdown"
    )
    reload_start = next(
        record
        for record in lifecycle
        if record.get("reason") == "reload" and record.get("phase") == "start"
    )
    authorization = [
        record for record in records if record.get("event") == "generation_authorization"
    ]
    current_generation = reload_start["new_generation"]
    current_authorization = next(
        record
        for record in authorization
        if record.get("requested_generation") == current_generation
    )
    superseded_authorization = next(
        record
        for record in authorization
        if record.get("requested_generation") == current_generation - 1
    )

    identities = {
        (record["peercred"]["pid"], record["starttime_ticks"])
        for record in [*lifecycle, promotion, launcher_registration, *authorization]
    }
    assert len(identities) == 1, f"same Pi PID/starttime did not persist: {identities}"

    generations = [record["new_generation"] for record in lifecycle]
    assert all(
        previous < current for previous, current in zip(generations, generations[1:])
    ), f"generation did not strictly increase: {generations}"
    observed_reasons = {record.get("reason") for record in lifecycle}
    assert {"startup", "reload", "fork", "new", "resume"} <= observed_reasons, (
        f"missing lifecycle boundary: {observed_reasons}"
    )
    assert launcher_registration.get("session_id_shape") == "hex-256"
    assert reload_revoke.get("prior_lease") == "VERIFIED"
    assert reload_revoke.get("prior_lease_revoked") is True
    assert current_authorization.get("code") == "MUTATOR_UNVERIFIED"
    assert superseded_authorization.get("code") == "STALE_GENERATION"

    return {
        "same_pid_starttime": next(iter(identities)),
        "generations": generations,
        "reload_revoke_verified": True,
        "lease_anchor_fixture": True,
        "new_generation_code": "MUTATOR_UNVERIFIED",
        "superseded_generation_code": "STALE_GENERATION",
    }


def run_once(index: int) -> Path:
    root = Path(tempfile.mkdtemp(prefix=f"gate0-d4-{index}-"))
    workspace = root / "workspace"
    sessions = root / "sessions"
    workspace.mkdir()
    sessions.mkdir()
    socket_path = root / "generation.sock"
    generation_log = root / "generation.jsonl"
    pi_log = root / "pi.jsonl"
    extension = root / "d4_extension.ts"
    write_extension(extension)

    environment = os.environ.copy()
    # Do not let the gated launcher inherit or fall back to its live/default
    # broker.  Both its registration handshake and D4 lifecycle use this one
    # per-run p3 fixture socket; there is no second broker.
    environment.pop("MOSAIC_LEASE_BROKER_SOCKET", None)
    environment.update(
        {
            "D4_GENERATION_SOCKET": str(socket_path),
            "MOSAIC_LEASE_BROKER_SOCKET": str(socket_path),
            "D4_PI_LOG": str(pi_log),
            "MOSAIC_PI_FORCE_SKILLS": "",
            "PI_SKIP_VERSION_CHECK": "1",
        }
    )
    # Must run before the fixture broker or Pi process is launched.  It proves
    # the launcher registers before exec and can only read this fixture socket.
    gated_launcher_precondition(root, socket_path, environment)
    broker = subprocess.Popen(
        [
            sys.executable,
            str(HERE / "p3_generation_broker.py"),
            "--socket",
            str(socket_path),
            "--log",
            str(generation_log),
        ],
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    wait_path(socket_path)
    pi = PiRpc(
        [
            sys.executable,
            str(GATED_LAUNCHER),
            "--runtime",
            "pi",
            "--",
            "pi",
            "--mode",
            "rpc",
            "--session-dir",
            str(sessions),
            "--no-extensions",
            "--no-context-files",
            "--no-prompt-templates",
            "--model",
            "openai-codex/gpt-5.6-sol",
            "--thinking",
            "medium",
            "--extension",
            str(extension),
        ],
        workspace,
        environment,
    )
    try:
        pi.send({"id": "state", "type": "get_state"})
        state = pi.response("state")
        original_session = state["data"]["sessionFile"]
        pi.prompt_and_settle(
            "fixture-promote",
            "Call d4_fixture_promote exactly once, then stop.",
        )
        pi.send({"id": "reload", "type": "prompt", "message": "/d4-reload"})
        reload_response = pi.response("reload")
        if not reload_response.get("success"):
            raise RuntimeError(f"reload failed: {reload_response}")
        for request_id, request_payload in [
            ("clone", {"id": "clone", "type": "clone"}),
            ("new", {"id": "new", "type": "new_session"}),
            (
                "resume",
                {"id": "resume", "type": "switch_session", "sessionPath": original_session},
            ),
        ]:
            pi.send(request_payload)
            response = pi.response(request_id)
            if not response.get("success") or response.get("data", {}).get("cancelled"):
                raise RuntimeError(f"{request_id} failed: {response}")
        assertions = assert_d4(jsonl(generation_log))
        (root / "machine-assertions.json").write_text(
            json.dumps(assertions, sort_keys=True, indent=2) + "\n",
            encoding="utf-8",
        )
        print(f"run={index} evidence_dir={root}")
        print("machine_assertions=PASS")
        print(json.dumps(assertions, sort_keys=True))
    except Exception as error:
        (root / "machine-assertions.json").write_text(
            json.dumps({"error": f"{type(error).__name__}: {error}"}, sort_keys=True, indent=2)
            + "\n",
            encoding="utf-8",
        )
        print(f"run={index} evidence_dir={root}")
        print("machine_assertions=FAIL")
        print(f"error={type(error).__name__}: {error}")
        raise
    finally:
        pi.close()
        try:
            request(socket_path, {"action": "shutdown-broker"})
        except OSError:
            pass
        try:
            broker.wait(timeout=5)
        except subprocess.TimeoutExpired:
            broker.kill()
            broker.wait()
    return root


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--runs", type=int, default=3)
    args = parser.parse_args()
    if args.runs < 1:
        raise SystemExit("--runs must be at least 1")
    roots: list[Path] = []
    for index in range(1, args.runs + 1):
        roots.append(run_once(index))
    print("d4_isolation_runs=" + ",".join(str(root) for root in roots))


if __name__ == "__main__":
    main()