#!/usr/bin/env python3
"""Fake ``pi --mode rpc`` used by tests so CI never needs node/pi installed.

Speaks a small, deliberately incomplete subset of the protocol documented in
``docs/pi-rpc.md`` (itself summarizing the real
``@earendil-works/pi-coding-agent`` ``docs/rpc.md``):

* Reads JSONL commands from stdin, one per line.
* On ``prompt``: replies with a ``response`` ack, then streams a
  ``message_update`` text-delta event (two when the prompt contains
  "TWODELTA"), optionally a ``tool_execution_start``/``tool_execution_end``
  pair when the prompt contains "TOOL", an ``extension_ui_request`` when the
  prompt contains "PROPOSE", and finally an ``agent_end`` event.
* On ``abort``: replies with a ``response`` ack, then emits ``agent_end``
  with ``aborted: true``.
* On any other/unknown command type: emits an ``{"type": "error", ...}``
  event.

Two knobs exist only so the session daemon's tests can watch this process
from the outside (``tests/test_daemon_lifecycle.py``):

* ``FAKE_PI_PIDDIR`` -- write ``<pid>.pid`` there at start and remove it at
  exit, so a test can count live fake-pi processes without ``pgrep -f``
  (which under ``pytest -n auto`` also sees other tests' fakes).
* ``FAKE_PI_CMDLOG`` -- append every received command object as one JSON
  line, so a test can assert that ``new_session``/``switch_session``
  actually reached pi.

Set ``FAKE_PI_CRASH=1`` to make this process ``exit(1)`` partway through
streaming a prompt's events (after the text delta, before ``agent_end``),
simulating a crashed backend that a real caller must not hang waiting on.

Two knobs model what the *real* pi does around commands (deviation d14):

* ``FAKE_PI_STRICT_ACK=1`` -- behave like pi 0.85.1 when a client
  pipelines: if another command line is already waiting on stdin when this
  process is about to acknowledge one, both are dropped and the process
  goes mute forever. That is the exact shape of d14, so a client that
  writes a prompt without waiting for the previous command's ack hangs
  here too instead of passing the test suite and failing on the Spark.
* ``FAKE_PI_ACK_DELAY=<seconds>`` -- how long to dawdle before an ack,
  widening the window a pipelining client would race into (default 0.2,
  the measured round trip of the real pi).

Like the real pi, this fake owns the name of its session file: it creates
one under ``--session-dir`` on start and on every ``new_session``, and
reports it as ``get_state``'s ``data.sessionFile``.

Two more knobs model a *stuck* pi for the reliable-agent-stop tests:

* ``NVSH_FAKE_IGNORE_CANCEL=1`` -- an ``abort`` command is received and
  silently dropped: no ``response`` ack, no ``agent_end``. This is what a
  harness that has stopped honouring its own cancel protocol looks like
  from the client's side, so a test can prove ``force_stop()`` (kill_tree,
  not the protocol) is what actually ends the turn.
* ``NVSH_FAKE_GRANDCHILD=1`` -- start a ``sleep 600`` child at launch and
  write ``{"harness": <this pid>, "grandchild": <sleep's pid>}`` to
  ``$NVSH_FAKE_PID_FILE``, so a test can assert a whole process *tree* dies
  under ``kill_tree``, not just the direct child nvsh spawned.
* ``NVSH_FAKE_STALL_TURN=1`` -- a ``prompt`` is acked and then left open: no
  ``message_update``, no ``agent_end``, so an end-to-end stop test
  (tests/test_agent_stop.py) has a turn that is genuinely in flight, silent
  and free of approval dialogs when the operator presses Ctrl+C.
"""

from __future__ import annotations

import atexit
import json
import os
import select
import subprocess  # nosec B404 - fixed argv below, test fixture only
import sys
import time
import uuid
from pathlib import Path


def spawn_grandchild_if_requested() -> None:
    """``NVSH_FAKE_GRANDCHILD=1``: start a sleeping child and record both pids."""
    if os.environ.get("NVSH_FAKE_GRANDCHILD") != "1":
        return
    pid_file = os.environ.get("NVSH_FAKE_PID_FILE")
    child = subprocess.Popen(["sleep", "600"])  # nosec B603 B607 - fixed argv
    if pid_file:
        Path(pid_file).write_text(
            json.dumps({"harness": os.getpid(), "grandchild": child.pid}), encoding="utf-8"
        )


def write_pidfile() -> None:
    piddir = os.environ.get("FAKE_PI_PIDDIR")
    if not piddir:
        return
    path = Path(piddir) / f"{os.getpid()}.pid"
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(str(os.getpid()), encoding="utf-8")
    except OSError:
        return
    atexit.register(lambda: path.unlink(missing_ok=True))


def log_command(cmd: dict) -> None:
    cmdlog = os.environ.get("FAKE_PI_CMDLOG")
    if not cmdlog:
        return
    try:
        with open(cmdlog, "a", encoding="utf-8") as handle:
            handle.write(json.dumps(cmd) + "\n")
    except OSError:
        pass


def emit(obj: dict) -> None:
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


STATE = {"session_file": "", "mute": False}


class Reader:
    """Unbuffered line reader over fd 0, with an honest "is more waiting?".

    ``sys.stdin``'s own buffering would hide a pipelined second command
    inside Python's buffer, so the strict-ack check below reads fd 0
    directly and keeps its own remainder.
    """

    def __init__(self) -> None:
        self._buf = b""

    def _fill(self) -> bool:
        chunk = os.read(0, 65536)
        if not chunk:
            return False
        self._buf += chunk
        return True

    def has_pending(self, timeout: float) -> bool:
        """Is a whole further command already available within *timeout*?"""
        if b"\n" in self._buf:
            return True
        if not select.select([0], [], [], timeout)[0]:
            return False
        return self._fill() and b"\n" in self._buf

    def readline(self) -> str:
        while b"\n" not in self._buf:
            if not self._fill():
                line, self._buf = self._buf, b""
                return line.decode("utf-8", errors="replace")
        line, _, self._buf = self._buf.partition(b"\n")
        return line.decode("utf-8", errors="replace")


def session_dir() -> Path | None:
    argv = sys.argv
    if "--session-dir" not in argv:
        return None
    return Path(argv[argv.index("--session-dir") + 1])


def new_session_file() -> str:
    """Create a session file the way pi does: our name, inside --session-dir."""
    directory = session_dir()
    if directory is None:
        return ""
    directory.mkdir(parents=True, exist_ok=True)
    path = directory / f"{uuid.uuid4()}.jsonl"
    path.write_text("", encoding="utf-8")
    STATE["session_file"] = str(path)
    return STATE["session_file"]


def ack(reader: "Reader", cmd: dict, command: str, **extra: object) -> None:
    """Acknowledge one command, modelling pi 0.85.1's pipelining bug.

    With ``FAKE_PI_STRICT_ACK=1``, a client that has already written its
    next command before this ack goes out loses both: the fake stops
    answering anything, forever, exactly as the real pi did in d14.
    """
    delay = float(os.environ.get("FAKE_PI_ACK_DELAY", "0") or 0)
    if os.environ.get("FAKE_PI_STRICT_ACK") == "1":
        if reader.has_pending(delay or 0.05):
            STATE["mute"] = True
            return
    elif delay:
        time.sleep(delay)
    emit({"id": cmd.get("id"), "type": "response", "command": command, "success": True, **extra})


def handle_prompt(reader: "Reader", cmd: dict) -> None:
    message = cmd.get("message", "")
    ack(reader, cmd, "prompt")
    if STATE["mute"] or os.environ.get("NVSH_FAKE_STALL_TURN") == "1":
        return

    emit(
        {
            "type": "message_update",
            "assistantMessageEvent": {"type": "text_delta", "delta": "looking at the failure"},
        }
    )
    if "TWODELTA" in message:
        emit(
            {
                "type": "message_update",
                "assistantMessageEvent": {"type": "text_delta", "delta": " and one more thing"},
            }
        )

    if os.environ.get("FAKE_PI_CRASH") == "1":
        sys.exit(1)

    if "TOOL" in message:
        emit(
            {
                "type": "tool_execution_start",
                "toolCallId": "call_1",
                "toolName": "bash",
                "args": {"command": "nvidia-smi"},
            }
        )
        emit(
            {
                "type": "tool_execution_end",
                "toolCallId": "call_1",
                "toolName": "bash",
                "result": {"content": [{"type": "text", "text": "ok"}]},
                "isError": False,
            }
        )

    if "PROPOSE" in message:
        emit(
            {
                "type": "extension_ui_request",
                "id": "ui-1",
                "method": "confirm",
                "title": "Apply fix?",
                "command": "sudo nvidia-smi -pm 1",
                "message": "Run: sudo nvidia-smi -pm 1",
            }
        )
        # A real pi blocks here for an extension_ui_response before
        # continuing. This fake does not implement that round trip (it is
        # only used to exercise the extension_ui_request -> PROPOSAL
        # mapping); it always finishes the turn regardless.

    emit({"type": "agent_end", "messages": [], "willRetry": False})


def handle_abort(cmd: dict) -> None:
    emit({"id": cmd.get("id"), "type": "response", "command": "abort", "success": True})
    emit({"type": "agent_end", "messages": [], "willRetry": False, "aborted": True})


def handle_unknown(cmd: dict) -> None:
    emit(
        {
            "id": cmd.get("id"),
            "type": "response",
            "command": cmd.get("type", "unknown"),
            "success": False,
            "error": f"unknown command: {cmd.get('type')}",
        }
    )
    emit({"type": "error", "error": f"unknown command: {cmd.get('type')}"})


def main() -> None:
    write_pidfile()
    spawn_grandchild_if_requested()
    new_session_file()
    reader = Reader()
    while True:
        raw_line = reader.readline()
        if not raw_line:
            return
        line = raw_line.strip()
        if not line:
            continue
        try:
            cmd = json.loads(line)
        except json.JSONDecodeError:
            emit({"type": "response", "command": "parse", "success": False, "error": "bad json"})
            continue

        log_command(cmd)
        if STATE["mute"]:
            continue  # pipelined into: this process answers nothing, ever
        cmd_type = cmd.get("type")
        if cmd_type == "prompt":
            handle_prompt(reader, cmd)
        elif cmd_type == "abort":
            if os.environ.get("NVSH_FAKE_IGNORE_CANCEL") == "1":
                continue  # ignore: no ack, no agent_end -- a stuck harness
            handle_abort(cmd)
        elif cmd_type in ("new_session", "switch_session"):
            if cmd_type == "new_session":
                new_session_file()
            else:
                STATE["session_file"] = str(cmd.get("sessionPath") or "")
            ack(reader, cmd, cmd_type, data={"cancelled": False})
        elif cmd_type == "get_state":
            ack(
                reader,
                cmd,
                "get_state",
                data={"isStreaming": False, "sessionFile": STATE["session_file"]},
            )
        else:
            handle_unknown(cmd)


if __name__ == "__main__":
    main()
