"""
PiAdapter — drives the `pi` coding agent via `--mode rpc` (JSONL stdin/stdout).

Pi (github.com/earendil-works/pi) is a multi-LLM coding agent. It runs as a
persistent subprocess in RPC mode, speaking strict JSONL (one JSON object per
line, LF-delimited) on stdin/stdout.

Install: npm install -g @earendil-works/pi-coding-agent
     or: curl -fsSL https://pi.dev/install.sh | sh

This adapter targets the pi 0.79.x RPC protocol. The authoritative spec ships
with the package at `docs/rpc.md`. Key shapes we rely on:

  Send to stdin (commands):
    {"type": "prompt",  "message": "<text>", "images": [{"type":"image","data":"<b64>","mimeType":"image/png"}]}
    {"type": "abort"}
    {"type": "extension_ui_response", "id": "<gate_id>", "confirmed": true|false}   # confirm dialogs
    {"type": "extension_ui_response", "id": "<gate_id>", "value": "<option>"}        # select/input/editor
    {"type": "extension_ui_response", "id": "<gate_id>", "cancelled": true}          # dismiss

  Receive from stdout (events — never carry an "id", only responses do):
    {"type": "agent_start"}
    {"type": "message_update", "assistantMessageEvent": {"type": "text_delta",     "delta": "..."}}
    {"type": "message_update", "assistantMessageEvent": {"type": "thinking_delta", "delta": "..."}}
    {"type": "message_update", "assistantMessageEvent": {"type": "...", "partial": {"usage": {...}}}}
    {"type": "tool_execution_start", "toolCallId": "...", "toolName": "...", "args": {...}}
    {"type": "tool_execution_end",   "toolCallId": "...", "toolName": "...",
                                     "result": {"content": [{"type":"text","text":"..."}]}, "isError": false}
    {"type": "turn_end", "message": {"usage": {"input","output","cacheRead","cost":{"total"}}}, "toolResults": [...]}
    {"type": "extension_ui_request", "id": "...", "method": "confirm"|"select"|"input"|"editor"|"notify"|..., ...}
    {"type": "agent_end", "messages": [...]}        # ← turn terminator (no id, no success field)

Session continuity:
  Sessions are JSONL files. `--session <path|id>` resumes/creates a session tree.

LLM: multi-provider. We default to OpenRouter, pulling the key from
     $OPENROUTER_API_KEY or, failing that, the Hermes credential pool
     (~/.hermes/auth.json). Override the model via the new-session picker.

Tools: read, write, edit, bash (built-in). MCP servers / skills add more.
"""
from __future__ import annotations
import asyncio
import json
import os
import re
import shutil
import logging
import uuid
from pathlib import Path
from typing import AsyncIterator, Callable

logger = logging.getLogger(__name__)

from ..attachments import Attachment, build_text_with_inline
from ..blocks import (
    Block, AgentInfoBlock, ToolCallBlock, ThinkingBlock, TextBlock,
    AudioHintBlock, SessionEndBlock, ErrorBlock, PermissionRequestBlock,
    SessionStatusHintBlock, tool_result_or_large,
)
from .base import (
    BaseAdapter, ExternalSession, ORPolicy, split_markdown, agent_env,
    tool_label_subject, context_window_for_model,
)

from ..adapter_logging import setup_adapter_logger
from ..platform_paths import user_config_dir as _config_dir

_SESSIONS_DIR = _config_dir("pi") / "sessions"

# A transcript written within this window of a scan ⇒ activity "turn" (running),
# else "open". Matches claude's _TURN_RECENT_SECS; the old 2.5 s was too tight — a
# busy session rarely lands inside it at scan time and mis-showed as merely "open".
_TURN_RECENT_SECS = 20.0


def _pi_session_cwd(path: Path) -> str:
    """The cwd recorded on line 1 of a pi session file ({type:session, cwd:…})."""
    try:
        with path.open("r", errors="replace") as f:
            obj = json.loads(f.readline())
        if obj.get("type") == "session":
            return obj.get("cwd", "") or ""
    except Exception:
        pass
    return ""


def _active_pi_cwds() -> set[str]:
    """Realpath working dirs of running *external* pi processes (exe basename ``pi``).

    pi (≥0.79) no longer holds its session ``.jsonl`` open for the process lifetime —
    it append-and-closes per event — so ``lsof`` on the session file (the old ownership
    signal) returns nothing and every external pi session looked idle and was dropped
    from the drawer. The robust live signal is the process cwd, mirroring claude's
    ``_scan_claude_procs``: a pi session file whose recorded cwd has a live pi process
    is the one in a terminal right now. Beacon's own spawned pi subprocesses are
    excluded (process subtree + cmdline signature)."""
    import subprocess
    from .base import own_subtree_pids, is_beacon_spawned_cmdline
    out: set[str] = set()
    try:
        # Full command line (not just comm) so a Beacon-spawned pi — including orphans
        # from a prior bridge instance, not in our subtree — can be excluded by
        # signature. The executable basename must be `pi` to count.
        ps = subprocess.check_output(
            ["ps", "-axo", "pid=,command="], text=True, stderr=subprocess.DEVNULL)
    except Exception:
        return out
    own = own_subtree_pids()  # exclude Beacon's own spawned pi subprocesses
    for line in ps.splitlines():
        line = line.strip()
        parts = line.split(None, 1)
        if len(parts) != 2:
            continue
        pid, cmd = parts
        if not pid.isdigit():
            continue
        exe = cmd.split()[0] if cmd else ""
        if os.path.basename(exe) != "pi":
            continue
        if int(pid) in own or is_beacon_spawned_cmdline(cmd):
            continue
        try:
            # -a ANDs -p and -d → exactly this pid's cwd as an `n<path>` line.
            cwd_out = subprocess.check_output(
                ["lsof", "-p", pid, "-a", "-d", "cwd", "-Fn"],
                text=True, stderr=subprocess.DEVNULL)
        except Exception:
            continue
        for ln in cwd_out.splitlines():
            if ln.startswith("n/") and os.path.isdir(ln[1:]):
                out.add(os.path.realpath(ln[1:]))
    return out


def _pi_session_meta(path: Path) -> tuple[str, str]:
    """(started_at, first_user_prompt) from the head of a pi session file."""
    started, prompt = "", ""
    try:
        with path.open("r", errors="replace") as f:
            for _ in range(60):
                line = f.readline()
                if not line:
                    break
                try:
                    o = json.loads(line)
                except Exception:
                    continue
                if not started and o.get("timestamp"):
                    started = o["timestamp"]
                if not prompt and o.get("type") == "message":
                    m = o.get("message") or {}
                    if m.get("role") == "user":
                        for c in m.get("content") or []:
                            if isinstance(c, dict) and c.get("type") == "text" and c.get("text"):
                                prompt = c["text"][:200]
                                break
                if started and prompt:
                    break
    except Exception:
        pass
    return started, prompt


def _scan_pi_sessions(limit: int = 20) -> list[ExternalSession]:
    """Discover pi sessions on disk (~/.config/pi/sessions/*.jsonl). Beacon-created
    files (named by relay session_id) are filtered centrally in headless."""
    if not _SESSIONS_DIR.exists():
        return []
    import time as _t
    active_cwds = _active_pi_cwds()
    cands: list[tuple[float, Path]] = []
    for f in _SESSIONS_DIR.glob("*.jsonl"):
        try:
            cands.append((f.stat().st_mtime, f))
        except OSError:
            continue
    cands.sort(key=lambda x: x[0], reverse=True)
    now = _t.time()
    claimed: set[str] = set()  # one live session per active cwd — the newest transcript
    out: list[ExternalSession] = []
    for mtime, f in cands[:limit]:
        started, prompt = _pi_session_meta(f)
        cwd = _pi_session_cwd(f)
        real = os.path.realpath(cwd) if cwd else ""
        # Live iff a running external pi owns this session's cwd AND this is the newest
        # transcript in that cwd (pi no longer holds the file open, so cwd — not an fd —
        # is the ownership signal; see _active_pi_cwds).
        is_live = bool(real) and real in active_cwds and real not in claimed
        if is_live:
            claimed.add(real)
        activity = "idle"
        if is_live:
            activity = "turn" if (now - mtime) <= _TURN_RECENT_SECS else "open"
        out.append(ExternalSession(
            session_id=f.stem, agent_id="pi", cwd=cwd,
            last_prompt=prompt, started_at=started, is_active=is_live, activity=activity))
    return out

# pi emits one JSON object per line; some events/responses (notably
# get_available_models — hundreds of models) exceed asyncio's default 64 KB
# StreamReader limit and would raise LimitOverrunError on `async for`.
_RPC_STREAM_LIMIT = 16 * 1024 * 1024

_BUILTIN_TOOLS = [
    {"name": "read",  "description": "Read files"},
    {"name": "write", "description": "Write / create files"},
    {"name": "edit",  "description": "Edit files (search-and-replace)"},
    {"name": "bash",  "description": "Run shell commands"},
]


# Transient provider failures pi retries in-process (mirrors the herdr pi-extension
# `retryableErrorPattern`). Matched against the last assistant message's errorMessage.
_RETRYABLE_ERR_RE = re.compile(
    r"overload|provider.?returned.?error|rate.?limit|too many requests|\b429\b|"
    # Explicit transient status codes only — a bare `\b5\d\d\b` also matched incidental
    # numbers in non-retryable text ("context window 512 tokens exceeded") → false hold.
    r"\b50[0234]\b|"
    r"service.?unavailable|server.?error|internal.?error|network.?error|"
    r"connection.?(error|refused|lost)|websocket.?(closed|error)|other side closed|"
    r"fetch failed|upstream.?connect|reset before headers|socket hang up|"
    r"ended without|http2 request did not get a response|timed? out|timeout|"
    # undici's transport failure is the bare string "terminated"; an unqualified
    # `terminated` also matched deliberate stops ("Process terminated by user").
    r"^terminated\.?$|(connection|stream|socket|request).?terminated|retry delay",
    re.I,
)


# While a pi tool is executing, the per-line read timeout must NOT be the provider-hang
# timeout: a `bash` tool doing a build/install/test run emits no events for many minutes,
# and firing there desynchronizes the adapter from a pi run that is still very much alive
# (the turn dies mid-output; the next prompt is then read against the OLD run's leftover
# events). Generous upper bound only so a truly wedged tool can't pin the turn forever —
# the user can always /stop (cap_interrupt).
_TOOL_TIMEOUT_SECS = float(os.environ.get("TERMINULL_PI_TOOL_TIMEOUT_SECS", "1800"))


def _agent_end_errored(event: dict) -> bool:
    """True if pi's `agent_end` reports a FAILED attempt (last assistant message
    `stopReason == "error"`). A successful end must never be held for a retry — pi has
    nothing left to retry — no matter what earlier `auto_retry` events said."""
    msgs = event.get("messages")
    if not isinstance(msgs, list):
        return False
    for m in reversed(msgs):
        if isinstance(m, dict) and m.get("role") == "assistant":
            return m.get("stopReason") == "error"
    return False


def _retryable_agent_end_error(event: dict) -> str | None:
    """If pi's `agent_end` reports a *retryable* provider error on its last assistant
    message, return that error text, else None. Fields (`stopReason`/`errorMessage`) are
    pi's own, observed by the herdr pi extension. Unknown/renamed fields ⇒ None ⇒ the
    caller falls back to the normal turn-end (so a schema drift can't wedge a turn)."""
    msgs = event.get("messages")
    if not isinstance(msgs, list):
        return None
    last_asst = None
    for m in reversed(msgs):
        if isinstance(m, dict) and m.get("role") == "assistant":
            last_asst = m
            break
    if not last_asst or last_asst.get("stopReason") != "error":
        return None
    err = str(last_asst.get("errorMessage") or "")
    return err if _RETRYABLE_ERR_RE.search(err) else None


class PiAdapter(BaseAdapter):
    cap_images = True            # per-model; overridden at make_agent_info time (self._cap_images)
    cap_documents = False        # PDFs inlined as a path hint; pi has no native doc input
    # Pi ignores ~/.config/pi/mcp.json in --mode rpc and has no --mcp flag, so the global
    # MCP script never reaches it. Instead we deliver the beacon tools via a pi *extension*
    # (~/.terminull/pi_terminull.ts) loaded with `--extension` at spawn — pi's first-class tool
    # path, which DOES reach the agent. The extension forwards to the same SentinelServer
    # over TCP, routing by TERMINULL_SESSION_ID. (Pi's native extension_ui_request dialogs are
    # still mapped in stream() for pi's own confirm/select prompts.)
    cap_sentinel = True
    cap_thinking = True
    cap_permissions = True
    cap_subagents = True
    cap_interrupt = True
    cap_input_during_turn = False  # serializes sends in its own _send_queue (no true
                                   # mid-turn interleave) — route mid-turn msgs through
                                   # the bridge's droppable queue, not an invisible defer
    provides_models = True
    provides_usage = True
    or_policy = ORPolicy.PRIMARY  # OpenRouter is pi's default provider
    agent_id = "pi"
    display_name = "Pi"
    beta = False              # live-verified (in daily use); no Beta badge
    cli_binary = "pi"

    slash_command_transforms: dict[str, object] = {
        "/plan": lambda args: (
            f"Before writing any code, create a detailed step-by-step plan for:\n"
            f"{args or 'the requested task'}\n\n"
            f"After I confirm the plan, proceed with implementation."
        ),
    }

    def __init__(
        self,
        binary: str | None = None,
        cwd: str | None = None,
        model: str | None = None,
        thinking: str | None = None,
        resume_session_id: str | None = None,
        on_session_id: "Callable[[str], None] | None" = None,
        isolated_home: bool = False,
    ) -> None:
        self._binary = binary or shutil.which("pi") or "pi"
        self._cwd = cwd
        # Multi-tenant on a shared runner (cloud gateway): pin HOME under cwd so anything pi keeps
        # under $HOME is per-room, not co-mingled across tenants. Off for the normal bridge (the
        # user's real HOME holds pi's config/auth). Handoff isn't used in cloud mode, so the
        # module-level session-dir lookup being HOME-independent is harmless there.
        self._isolated_home = isolated_home
        # Resume an existing pi session file by stem (e.g. one started in a terminal and
        # adopted via the sidebar). None ⇒ name the session file by the Beacon session_id.
        self._resume_session_id = resume_session_id
        # Persists pi's session-file id as the ext_sid (handoff dedup + durable /clear).
        # pi has no server-assigned id, so the file stem IS the identity; on clear we mint
        # a fresh stem and report it here so a reconnect resumes the NEW file, not the old.
        self._on_session_id = on_session_id
        # None → let pi resolve its own default (depends on which provider has
        # credentials); the real model is detected from get_state in start().
        self._model = model
        self._thinking = thinking          # "off"|"minimal"|"low"|"medium"|"high"|"xhigh"|None
        # Track if model was explicitly set by user (vs detected from pi's default)
        self._model_explicitly_set = bool(model)
        self._proc: asyncio.subprocess.Process | None = None
        self._session_file: Path | None = None
        self._session_id: str = ""
        
        # Session-specific adapter log file for debugging
        self._adapter_logger: logging.Logger | None = None
        # Placeholder log methods before setup
        self._log = self._log_debug = self._log_warn = self._log_err = lambda _: None
        # Pi has no live set-model RPC — a model switch is applied by relaunching the
        # subprocess with the new --model. set_model() raises this; stream() returns when
        # it's set so the session loop calls start() again (the generic restart path).
        self._pending_restart = False
        # queue items: (message_text, images_payload)
        self._send_queue: asyncio.Queue[tuple[str, list[dict]]] = asyncio.Queue()
        self._fetch_store: dict[str, str] = {}
        self._ui_requests: dict[str, dict] = {}   # gate_id → {method, options}

        # Usage accounting (mirrors KiloCodeAdapter conventions)
        self._context_window = context_window_for_model(model) if model else 0
        self._cap_images = True
        self._turn_input_tokens = 0
        self._turn_cache_read = 0
        self._total_output_tokens = 0
        self._total_cost_usd = 0.0
        self._total_turns = 0
        # Retryable-error hold (see agent_end handler): pi auto-retries a transient
        # provider error (overload/rate-limit/5xx) in-process, emitting agent_end for the
        # FAILED attempt first. Without this we'd flash 'completed'/'Finished' and then the
        # retry streams into an already-ended turn. Reset per user turn in _read_turn.
        self._retry_holds = 0
        self._saw_auto_retry = False

    # ── Session handoff (docs/session-handoff.md) ──────────────────────────────

    @classmethod
    def scan_external_sessions(cls) -> list[ExternalSession]:
        return _scan_pi_sessions()

    @classmethod
    def is_cwd_externally_owned(cls, cwd: str, sid: str = "") -> bool:
        if not cwd:
            return False
        real = os.path.realpath(os.path.expanduser(cwd))
        if real not in _active_pi_cwds():
            return False
        # Cwd is active — but is *this* session the one that's live?
        # Only the newest session file in the cwd is considered owned.
        if not sid:
            return True  # no sid to check against, assume owned
        # Find the newest session file whose cwd matches
        newest_mtime = 0.0
        newest_sid = ""
        for f in _SESSIONS_DIR.glob("*.jsonl"):
            try:
                if _pi_session_cwd(f) and os.path.realpath(_pi_session_cwd(f)) == real:
                    mtime = f.stat().st_mtime
                    if mtime > newest_mtime:
                        newest_mtime = mtime
                        newest_sid = f.stem
            except OSError:
                continue
        # This session is owned iff it's the newest one in the active cwd
        return newest_sid == sid

    @classmethod
    def make_transcript_mirror(cls, cwd: str, sid: str, *, from_end: bool = False):
        from ..transcript_mirror import PiTranscriptMirror
        f = _SESSIONS_DIR / f"{sid}.jsonl"
        return PiTranscriptMirror(f, from_end=from_end) if (sid and f.exists()) else None

    @classmethod
    def terminal_resume_command(cls, sid: str) -> str | None:
        import shlex
        if not sid:
            return "pi"
        return f"pi --session {shlex.quote(str(_SESSIONS_DIR / f'{sid}.jsonl'))}"

    @classmethod
    def readiness(cls) -> tuple[bool, str]:
        if cls.shared_provider_key():
            return True, ""
        # Pi may still have its own provider configured (not cheap to detect), so
        # don't block — just nudge the high-leverage fix.
        return True, "Add an OpenRouter key for the full model list (or configure a pi provider)."

    def _setup_adapter_logger(self, session_id: str) -> None:
        """Create a per-session log file for detailed adapter debugging."""
        self._adapter_logger = setup_adapter_logger(
            "pi",
            session_id,
            extra={"model": self._model, "cwd": self._cwd}
        )
        # Keep a reference to the actual logger for _log calls
        self._log = self._adapter_logger.info
        self._log_debug = self._adapter_logger.debug
        self._log_warn = self._adapter_logger.warning
        self._log_err = self._adapter_logger.error

    def _log(self, msg: str, level: str = "info") -> None:
        """Log to adapter-specific log file."""
        if not self._adapter_logger:
            return
        # Use the bound methods for performance (set up in _setup_adapter_logger)
        if level == "debug":
            self._log_debug(msg)
        elif level == "warning" or level == "warn":
            self._log_warn(msg)
        elif level == "error" or level == "err":
            self._log_err(msg)
        else:
            self._log(msg)  # info is bound to self._log

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    async def start(self, session_id: str) -> AgentInfoBlock:
        self._session_id = session_id
        
        # Setup session-specific adapter log
        self._setup_adapter_logger(session_id)
        
        # A restart (model switch) re-enters start() while the old process is still
        # alive — stop it first so we don't leak an orphaned pi holding the session.
        if self._proc and self._proc.returncode is None:
            await self.stop()
        _SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
        # Resume an adopted external session file if given, else key by the Beacon sid.
        self._session_file = _SESSIONS_DIR / f"{self._resume_session_id or session_id}.jsonl"

        cmd = [self._binary, "--mode", "rpc", "--session", str(self._session_file)]
        inj = self.shared_provider_injection(self._model)
        # When the provider is injected separately via --provider, the --model
        # arg should be the raw model slug (e.g. anthropic/claude-sonnet-4.5),
        # not our prefixed bridge identifier (openrouter/anthropic/...).
        model_arg = self._model
        if model_arg and inj.active and model_arg.startswith(inj.provider_id + "/"):
            model_arg = model_arg[len(inj.provider_id) + 1:]
        if model_arg:
            cmd += ["--model", model_arg]
            logger.info("pi start[%s]: model selected=%r → --model %r (provider injected=%s)",
                        session_id[:12], self._model, model_arg, inj.active)
        else:
            # No model resolved. pi will fall back to whatever default the active provider
            # exposes — typically a free model, which surfaces to the user as the "free
            # models" error. This should not happen for a session where the user picked a
            # model: the bridge now persists the selection (_session_model) and re-applies
            # it on every respawn. Log loudly so a recurrence is obvious in the bridge log.
            logger.warning("pi start[%s]: NO model selected (self._model=%r) — pi will use a "
                           "provider default. If you picked a model this is a bug in the model "
                           "hand-off (set_model / _session_model fallback).",
                           session_id[:12], self._model)
        if self._thinking:
            cmd += ["--thinking", self._thinking]

        # Beacon sentinel tools: deliver them as a pi extension (pi ignores mcp.json in
        # rpc mode). The extension reads the sentinel port from ~/.terminull/sentinel_port
        # and routes by TERMINULL_SESSION_ID (stamped into env by agent_env below).
        if type(self).cap_sentinel:
            try:
                from ..sentinel_tools import write_pi_extension
                cmd += ["--extension", str(write_pi_extension())]
            except Exception as e:
                logger.warning(f"pi: failed to write beacon extension, sentinel tools off: {e}")

        cmd += inj.argv
        extra_env = dict(inj.env)
        if self._isolated_home and self._cwd:
            home = os.path.join(self._cwd, ".home")
            os.makedirs(home, exist_ok=True)
            extra_env["HOME"] = home

        logger.info(f"Starting pi with command: {cmd}")
        logger.info(f"Pi environment variables: {extra_env}")
        
        # Log effective OpenRouter context
        if inj.active:
            key_preview = f"{inj.env.get('OPENROUTER_API_KEY', '')[:8]}..." if inj.env.get('OPENROUTER_API_KEY') else "MISSING"
            logger.info(f"Pi Adapter: Using OpenRouter key preview: {key_preview}")
        else:
            logger.info("Pi Adapter: Shared provider injection NOT active")
        _dbg = os.environ.get("TERMINULL_PI_STDERR_FILE")
        _stderr = open(_dbg, "ab") if _dbg else asyncio.subprocess.DEVNULL
        self._proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=_stderr,
            cwd=self._cwd,
            env=agent_env(extra_env, session_id=self._session_id),
            limit=_RPC_STREAM_LIMIT,
        )
        logger.info(f"Pi process started with PID: {self._proc.pid}")

        # Detect the model pi actually resolved (may be its own default when we
        # passed no --model), plus its real context window / image support.
        await self._detect_model()

        return self.make_agent_info(
            name="Pi",
            tools=_BUILTIN_TOOLS,
            slash_commands=["/plan"],
            working_dirs=[self._cwd] if self._cwd else [],
            cap_images=self._cap_images,
        )

    async def _detect_model(self, timeout: float = 15) -> None:
        """Query get_state to learn the model pi resolved (and its real context
        window / image support). Safe to call before the first prompt OR after a
        turn ends — no events interleave with the response while pi is idle on the
        send queue. A short `timeout` is used on the post-turn refresh so a wedged
        pi can't stall turn completion."""
        # If a model switch is pending, pi is about to restart with the user-selected
        # model. Don't let the old process's get_state response overwrite that selection.
        if self._pending_restart:
            return
        try:
            data = await asyncio.wait_for(self._request("get_state"), timeout=timeout)
        except Exception:
            return
        model = (data or {}).get("model") or {}
        mid = model.get("id")
        if not mid:
            return
        prov = model.get("provider", "")
        resolved = f"{prov}/{mid}" if prov else mid

        # Don't overwrite an explicitly-set model. Pi may fall back to its default
        # (from ~/.pi/agent/settings.json) if the requested model isn't available,
        # but we want to preserve the user's selection so it's retried on restart.
        # Only accept pi's resolved model when we had no explicit selection.
        if self._model_explicitly_set:
            if resolved != self._model:
                logger.warning("pi _detect_model[%s]: pi resolved %r but user selected %r — keeping user selection",
                                self._session_id[:12], resolved, self._model)
            # Still update context window and image caps from the resolved model.
        else:
            if resolved != self._model:
                logger.info("pi _detect_model[%s]: pi resolved model %r (was %r)",
                            self._session_id[:12], resolved, self._model)
            self._model = resolved

        if model.get("contextWindow"):
            self._context_window = model["contextWindow"]
        self._cap_images = "image" in (model.get("input") or ["text"])

    async def _request(self, cmd_type: str, **kw) -> dict:
        """Send an RPC command and return its response `data`. Reads stdout until
        the matching response id; ignores any interleaved events."""
        assert self._proc is not None and self._proc.stdout is not None
        req_id = uuid.uuid4().hex
        await self._write({"type": cmd_type, "id": req_id, **kw})
        while True:
            raw = await self._proc.stdout.readline()
            if not raw:
                return {}
            try:
                ev = json.loads(raw.decode(errors="replace").rstrip("\r\n"))
                # Log raw events to detect silent failures / 401s from the pi side
                logger.debug(f"Pi RPC In: {json.dumps(ev)}")
            except Exception:
                continue
            if ev.get("type") == "response" and ev.get("id") == req_id:
                return ev.get("data") or {}

    def clear_history(self) -> None:
        # pi has no server-assigned id — the session-file stem IS the identity, and a fresh
        # session would otherwise be re-named by the Beacon session_id (so nulling resume
        # reuses the SAME file → not cleared). Mint a fresh stem, persist it as the ext_sid,
        # and restart the held process so start() opens the new (empty) session file.
        new_id = uuid.uuid4().hex
        self._resume_session_id = new_id
        self._session_file = _SESSIONS_DIR / f"{new_id}.jsonl"
        if self._on_session_id:
            self._on_session_id(new_id)
        self._pending_restart = True
        if self._proc and self._proc.returncode is None:
            try:
                if self._proc.stdin:
                    self._proc.stdin.close()
            except Exception:
                pass
            self._proc.terminate()

    async def stop(self) -> None:
        if self._proc and self._proc.returncode is None:
            logger.info(f"Stopping pi process with PID: {self._proc.pid}")
            try:
                if self._proc.stdin:
                    self._proc.stdin.close()
            except Exception as e:
                logger.warning(f"Error closing pi stdin: {e}")
            # SIGINT (Ctrl+C) lets Pi flush its session before exiting.
            # Fall back to SIGTERM, then SIGKILL if still stuck.
            import signal as _signal
            try:
                self._proc.send_signal(_signal.SIGINT)
            except ProcessLookupError:
                pass
            try:
                await asyncio.wait_for(self._proc.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                logger.info("Pi process didn't exit on SIGINT, sending SIGTERM...")
                try:
                    self._proc.terminate()
                except ProcessLookupError:
                    pass
                try:
                    await asyncio.wait_for(self._proc.wait(), timeout=3.0)
                except asyncio.TimeoutError:
                    logger.warning("Pi process didn't exit on SIGTERM, sending SIGKILL...")
                    try:
                        self._proc.kill()
                    except ProcessLookupError:
                        pass
                    await self._proc.wait()
            logger.info(f"Pi process with PID: {self._proc.pid} stopped with return code: {self._proc.returncode}")

    # ── I/O ───────────────────────────────────────────────────────────────────

    async def send(self, text: str, attachments: list[Attachment] | None = None) -> None:
        atts = attachments or []
        images = [
            {"type": "image", "data": a.b64, "mimeType": a.media_type}
            for a in atts if a.is_image
        ]
        # Non-image attachments (text/pdf/binary) are inlined into the prompt text.
        others = [a for a in atts if not a.is_image]
        first = not self._sentinel_preamble_sent
        self._sentinel_preamble_sent = True
        full_text = self.sentinel_client_preamble(is_first_turn=first) + build_text_with_inline(text, others)
        await self._send_queue.put((full_text, images))

    async def respond_permission(self, permission_id: str, approved: bool) -> None:
        req = self._ui_requests.pop(permission_id, {})
        method = req.get("method", "confirm")
        if method == "confirm":
            payload = {"confirmed": approved}
        elif approved:
            # select/input/editor: approve maps to the first ("allow") option.
            options = req.get("options") or []
            payload = {"value": options[0] if options else ""}
        else:
            payload = {"cancelled": True}
        await self._write({"type": "extension_ui_response", "id": permission_id, **payload})

    async def interrupt(self) -> None:
        await self._write({"type": "abort"})

    def set_model(self, model: str) -> None:
        """Select a new model. Pi can only honor --model at launch, so flag a restart;
        stream() returns on the flag and the session loop relaunches us with it. Without
        this the picked model was silently dropped and pi kept running its default."""
        if not model or model == self._model:
            return
        self._model = model
        self._model_explicitly_set = True  # user explicitly chose this model
        self._context_window = context_window_for_model(model) or 0
        self._pending_restart = True
        logger.info(f"pi set_model: {model!r} — relaunching on next turn")

    async def stream(self) -> AsyncIterator[Block]:
        self._log("Stream started, waiting for prompt")
        while True:
            # A model switch is pending — return so the session loop relaunches pi via
            # start() with the new --model (pi has no live set-model RPC). Reached after
            # a turn completes (loop top) for a mid-turn switch.
            if self._pending_restart:
                self._log("Stream returning: pending restart")
                return
            text, images = await self._send_queue.get()
            self._log(f"Got prompt from queue: {len(text)} chars, {len(images)} images")
            
            if self._pending_restart:
                # Switch arrived while idle on the queue — requeue this prompt and bail
                # so the relaunched process handles it with the new model.
                await self._send_queue.put((text, images))
                self._log("Requeued prompt for restart")
                return
            # Auto-reconnect: pi is persistent, but it can exit between turns (the user
            # hit stop → SIGINT/SIGTERM, a crash, or an OS kill). Respawn before the
            # prompt — requeue it and flag a restart so the session loop relaunches us
            # (same proven path as a model switch), resuming the same session file —
            # instead of looping forever on "Pi is not running".
            if self._proc is None or self._proc.returncode is not None:
                rc = self._proc.returncode if self._proc else None
                logger.info(f"pi not running (rc={rc}) — reconnecting before prompt")
                self._pending_restart = True
                await self._send_queue.put((text, images))
                return
            cmd: dict = {"type": "prompt", "message": text}
            if images:
                cmd["images"] = images
            try:
                await self._write(cmd)
            except (BrokenPipeError, ConnectionResetError, AssertionError):
                # pi exited before/at the prompt (e.g. crashed at startup, or a
                # resumed session file recorded a cwd that no longer exists — pi
                # aborts in that case). Degrade to a clean error turn instead of
                # letting the write exception crash the stream.
                rc = self._proc.returncode if self._proc else None
                logger.error(f"Pi adapter process exited unexpectedly. Return code: {rc}", exc_info=True)
                yield AudioHintBlock(state="agent_silent")
                yield ErrorBlock(
                    message=f"Pi is not running (exited{f' with code {rc}' if rc else ''}); "
                            "the turn could not be started.",
                    recoverable=False,
                )
                yield SessionEndBlock(reason="error")
                continue
            async for block in self._read_turn():
                yield block

    # ── Turn loop ─────────────────────────────────────────────────────────────

    async def _read_turn(self, timeout: float = 120.0,
                         tool_timeout: float | None = None) -> AsyncIterator[Block]:
        """Read one agent run: from the prompt we just sent until `agent_end`.

        Args:
            timeout: Max seconds to wait for each event *while waiting on the provider*.
                     Provider APIs can be slow (OpenRouter free tier often hangs), so
                     the default is 2 min.
            tool_timeout: The same, but while one of pi's tools is still executing (a
                     `bash` build/install legitimately runs silent for many minutes).
                     Defaults to _TOOL_TIMEOUT_SECS."""
        assert self._proc is not None and self._proc.stdout is not None
        yield AudioHintBlock(state="agent_speaking")

        text_buf = ""            # assistant text accumulated between tool boundaries
        thinking_active = False
        thinking_buf = ""
        # Per-run usage (one prompt → possibly several turns)
        run_output = 0
        run_cost = 0.0
        turn_ended = False
        emitted_content = False   # any text / tool / thinking shown this turn?
        tools_running: set[str] = set()   # toolCallIds pi has started but not ended
        self._retry_holds = 0     # fresh retryable-error budget per user turn
        self._saw_auto_retry = False

        self._log("Turn read started, waiting for events")

        def flush_text():
            nonlocal text_buf, emitted_content
            blocks = []
            if text_buf.strip():
                blocks = list(split_markdown(text_buf))
                emitted_content = True
            text_buf = ""
            return blocks

        # Read with per-line timeout to detect frozen provider APIs
        reader = self._proc.stdout
        event_count = 0
        tool_tmo = _TOOL_TIMEOUT_SECS if tool_timeout is None else tool_timeout
        while True:
            # A tool in flight means pi is working, not hung — wait on the tool budget.
            waiting_on_tool = bool(tools_running)
            wait = tool_tmo if waiting_on_tool else timeout
            try:
                raw = await asyncio.wait_for(reader.readline(), timeout=wait)
            except asyncio.TimeoutError:
                what = (f"tool{'s' if len(tools_running) > 1 else ''} "
                        f"{', '.join(sorted(tools_running))} still running"
                        if waiting_on_tool else "provider API hung?")
                logger.warning(f"pi adapter: no event for {wait}s — {what}")
                self._log(f"TIMEOUT: no event for {wait}s ({what}), "
                          f"{event_count} events received so far")
                yield AudioHintBlock(state="agent_silent")
                yield ErrorBlock(
                    message=(
                        f"A tool has been running for {wait / 60:.0f} min with no output — "
                        "giving up on this turn. Try /stop or retry."
                        if waiting_on_tool else
                        f"Model stopped responding for {wait:.0f}s. The provider API may be hung — try /stop or retry."
                    ),
                    recoverable=True,
                )
                yield SessionEndBlock(reason="timeout")
                return
            
            if not raw:
                # stdout EOF
                self._log("stdout EOF — pi process closed")
                break
                
            line = raw.decode(errors="replace").rstrip("\r\n")
            logger.debug(f"Received from pi: {line}")
            if not line:
                continue
            try:
                event = json.loads(line)
            except Exception:
                continue

            etype = event.get("type", "")
            event_count += 1
            
            # Log every event to adapter log
            self._log(f"Event #{event_count}: {etype}")

            # Responses to commands (have an "id"); not turn events — ignore here.
            if etype == "response":
                continue

            if etype == "message_update":
                ame = event.get("assistantMessageEvent") or {}
                sub = ame.get("type", "")
                if sub == "text_delta":
                    text_buf += ame.get("delta", "")
                elif sub == "thinking_start":
                    thinking_active = True
                    thinking_buf = ""
                    emitted_content = True
                    yield ThinkingBlock(thinking="", done=False)
                elif sub == "thinking_delta":
                    chunk = ame.get("delta", "")
                    thinking_buf += chunk
                    yield ThinkingBlock(thinking=chunk, done=False)
                elif sub == "thinking_end":
                    if thinking_active:
                        yield ThinkingBlock(thinking=thinking_buf, done=True)
                    thinking_active = False
                    thinking_buf = ""

            elif etype == "tool_execution_start":
                exec_id = event.get("toolCallId", str(uuid.uuid4()))
                tool_name = event.get("toolName", "unknown")
                args = event.get("args", {}) or {}
                # Flush any text the assistant emitted before calling the tool.
                for b in flush_text():
                    yield b
                label, subject = tool_label_subject(tool_name, args)
                tools_running.add(tool_name or exec_id)
                emitted_content = True
                yield AudioHintBlock(state="agent_silent")
                yield ToolCallBlock(id=exec_id, name=tool_name, args=args,
                                    status="running", label=label, subject=subject)

            elif etype == "tool_execution_end":
                exec_id = event.get("toolCallId", str(uuid.uuid4()))
                tool_name = event.get("toolName", "")
                is_error = bool(event.get("isError"))
                tools_running.discard(tool_name or exec_id)
                content = _join_tool_content(event.get("result"))
                label, subject = tool_label_subject(tool_name, {})
                yield ToolCallBlock(id=exec_id, name=tool_name, args={},
                                    status="error" if is_error else "done",
                                    label=label, subject=subject)
                if content:
                    yield tool_result_or_large(exec_id, content, self._fetch_store,
                                               is_error=is_error)
                yield AudioHintBlock(state="agent_speaking")

            elif etype == "turn_end":
                # Authoritative per-turn usage/cost.
                usage = (event.get("message") or {}).get("usage") or {}
                self._turn_input_tokens = usage.get("input", self._turn_input_tokens)
                self._turn_cache_read = usage.get("cacheRead", 0)
                run_output += usage.get("output", 0)
                run_cost += (usage.get("cost") or {}).get("total", 0.0)
                self._total_turns += 1

            elif etype == "extension_ui_request":
                method = event.get("method", "")
                # Fire-and-forget UI methods carry no response; surface dialogs only.
                if method in ("confirm", "select", "input", "editor"):
                    gate_id = event.get("id", str(uuid.uuid4()))
                    self._ui_requests[gate_id] = {
                        "method": method, "options": event.get("options"),
                    }
                    action = event.get("title") or event.get("message") or method
                    details = event.get("message", "") if method == "confirm" else ""
                    yield PermissionRequestBlock(id=gate_id, action=action,
                                                 risk="medium", details=details)

            elif etype == "extension_error":
                yield ErrorBlock(message=event.get("error", "extension error"),
                                 recoverable=True)

            elif etype == "agent_end":
                # Retryable provider error → HOLD, don't end. pi auto-retries a transient
                # failure (overload/rate-limit/5xx) in-process and emits agent_end for the
                # FAILED attempt first; ending here would flash 'completed'/'Finished' + an
                # 'empty response' notice, then the retry would stream into an already-ended
                # turn (the herdr pi-extension `holdForRetry`). Keep reading — the retry
                # emits a fresh agent_start → content → agent_end. Bounded: if pi does NOT
                # retry, the per-line read timeout above ends the turn; the per-turn hold
                # cap stops an unbounded loop. On unknown/renamed fields the helper returns
                # None ⇒ normal end (no wedge).
                # ⚠️ Hold ONLY on a FAILED attempt. `emitted_content` is NOT a usable
                # guard here: it's set when text is *flushed* (at a tool boundary), so a
                # plain text-only turn still reads False at agent_end. And
                # `_saw_auto_retry` is sticky for the rest of the turn, so a retry that
                # SUCCEEDED used to hold its own successful end → the turn ran until the
                # 120s read timeout and its text was discarded below.
                errored = _agent_end_errored(event)
                retry_err = _retryable_agent_end_error(event) if errored else None
                if errored and (retry_err or self._saw_auto_retry) and self._retry_holds < 3:
                    self._retry_holds += 1
                    self._saw_auto_retry = False
                    self._log(f"agent_end: retryable error, holding for pi retry "
                              f"({self._retry_holds}/3): {(retry_err or 'auto_retry')[:120]}")
                    text_buf = ""
                    thinking_buf = ""
                    thinking_active = False
                    continue
                for b in flush_text():
                    yield b
                if thinking_active:
                    yield ThinkingBlock(thinking=thinking_buf, done=True)
                # Empty completion (no text/tool/thinking) — common with rate-limited
                # OpenRouter ":free" models, which often return nothing. Surface a notice
                # so the turn isn't a bare, unexplained end card.
                if not emitted_content:
                    yield TextBlock(
                        text="_The model returned an empty response. Free-tier "
                             "(`:free`) models are frequently rate-limited — try again "
                             "or switch to a paid model._",
                        done=True,
                    )
                self._total_output_tokens += run_output
                self._total_cost_usd += run_cost
                yield AudioHintBlock(state="agent_silent")
                # pi can silently fall back to a different model mid-session (e.g. a
                # provider error drops it to deepseek). The turn is over and pi is idle
                # on the send queue now — the same no-interleave state as before the
                # first prompt — so re-query get_state and surface the model actually in
                # effect. Without this the header/selector keep showing the stale picked
                # model until the user manually re-selects.
                await self._detect_model(timeout=4)
                hint = await self.get_status_hint()
                if hint is not None:
                    yield hint
                yield SessionEndBlock(
                    reason="completed",
                    cost_usd=run_cost,
                    input_tokens=self._turn_input_tokens,
                    output_tokens=run_output,
                    cache_read_tokens=self._turn_cache_read,
                )
                turn_ended = True
                return

            elif etype.startswith("auto_retry"):
                # pi announced an in-process retry — corroborates (and covers the ordering
                # where this arrives before the failed agent_end) the retryable-error hold.
                self._saw_auto_retry = True

            # agent_start, turn_start, message_start/_end, tool_execution_update,
            # compaction_*, queue_update → no-ops for now.

        # stdout closed before agent_end — process crashed or was stopped (e.g. /stop →
        # SIGINT/SIGTERM). Fold this run's partial usage into the totals and surface it on
        # the end block so an interrupted turn still reports tokens spent. A signal death
        # (rc < 0) is a deliberate/forced stop, not an error.
        if not turn_ended:
            for b in flush_text():
                yield b
            yield AudioHintBlock(state="agent_silent")
            rc = self._proc.returncode if self._proc else None
            self._total_output_tokens += run_output
            self._total_cost_usd += run_cost
            reason = "completed" if rc in (None, 0) else ("stopped" if rc < 0 else "error")
            yield SessionEndBlock(
                reason=reason,
                cost_usd=run_cost,
                input_tokens=self._turn_input_tokens,
                output_tokens=run_output,
                cache_read_tokens=self._turn_cache_read,
            )

    # ── Usage reporting ─────────────────────────────────────────────────────────

    async def get_usage(self) -> dict:
        return {
            "model":                self._model,
            "input_tokens":         self._turn_input_tokens,
            "output_tokens":        self._total_output_tokens,
            "cost_usd":             self._total_cost_usd,
            "turns":                self._total_turns,
            # Real context occupancy = fresh input + cached portion (cache read alone is
            # most of the prompt once caching kicks in; fresh input understates "% full").
            "last_input_tokens":    self._turn_input_tokens + self._turn_cache_read,
            "cache_read_tokens":    self._turn_cache_read,
            "cache_creation_tokens": 0,
        }

    async def get_status_hint(self) -> SessionStatusHintBlock | None:
        if not self._model:
            return None
        return SessionStatusHintBlock(
            model=self._model,
            turns=self._total_turns,
            cost_usd=self._total_cost_usd,
            output_tokens=self._total_output_tokens,
            cache_creation_tokens=0,
            last_input_tokens=self._turn_input_tokens + self._turn_cache_read,
            context_window=self._context_window,
        )

    # ── Model picker ──────────────────────────────────────────────────────────

    # Popular coding models surfaced first; the rest follow alphabetically.
    _PRIORITY = (
        "claude-opus", "claude-sonnet", "claude-haiku", "gpt-5", "gpt-4.1",
        "gpt-4o", "o4-", "o3", "gemini-2.5-pro", "gemini-2.5-flash", "grok-4",
        "grok-code", "deepseek", "kimi", "qwen3-coder", "glm",
    )
    _MODEL_LIMIT = None  # was 60 — let the full provider list through; iOS can search

    @classmethod
    async def list_models(cls, extra: dict | None = None) -> list[dict]:
        """Live model list from pi's get_available_models (depends on which
        provider has credentials), prioritized and capped. Falls back to the
        shared OpenRouter model list (cached per process) when the pi probe
        fails, then to a hardcoded list as a last resort."""
        raw = await cls._probe_available_models()
        if raw:
            return cls._format_models(raw)
        return await cls.shared_catalog_models()

    @classmethod
    async def _probe_available_models(cls) -> list[dict]:
        binary = shutil.which("pi") or "pi"
        inj = cls.shared_provider_injection(None)
        cmd = [binary, "--mode", "rpc", "--no-session", *inj.argv]
        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.DEVNULL, env=agent_env(inj.env),
                limit=_RPC_STREAM_LIMIT,
            )
        except Exception:
            return []
        out: list[dict] = []
        try:
            req = uuid.uuid4().hex
            assert proc.stdin and proc.stdout
            proc.stdin.write((json.dumps({"type": "get_available_models", "id": req}) + "\n").encode())
            await proc.stdin.drain()
            while True:
                raw = await asyncio.wait_for(proc.stdout.readline(), timeout=20)
                if not raw:
                    break
                try:
                    ev = json.loads(raw.decode(errors="replace").rstrip("\r\n"))
                except Exception:
                    continue
                if ev.get("type") == "response" and ev.get("id") == req:
                    out = (ev.get("data") or {}).get("models", []) or []
                    break
        except Exception:
            out = []
        finally:
            try:
                proc.terminate()
                await proc.wait()
            except Exception:
                pass
        return out

    @classmethod
    def _format_models(cls, raw: list[dict]) -> list[dict]:
        def rank(m: dict) -> int:
            mid = m.get("id", "").lower()
            for i, p in enumerate(cls._PRIORITY):
                if p in mid:
                    return i
            return len(cls._PRIORITY) + 1

        def entry(m: dict) -> dict:
            prov, mid = m.get("provider", ""), m.get("id", "")
            full = f"{prov}/{mid}" if prov else mid
            bits: list[str] = [prov] if prov else []
            if m.get("reasoning"):
                bits.append("reasoning")
            cw = m.get("contextWindow") or 0
            if cw:
                bits.append(f"{cw // 1_000_000}M ctx" if cw >= 1_000_000 else f"{cw // 1000}k ctx")
            cost_in = (m.get("cost") or {}).get("input")
            if cost_in:
                bits.append(f"${cost_in:.4g}/1M in")
            return {"id": full, "label": mid.rsplit("/", 1)[-1], "detail": " · ".join(bits)}

        ordered = sorted(raw, key=lambda m: (rank(m), m.get("id", "")))
        if cls._MODEL_LIMIT:
            ordered = ordered[:cls._MODEL_LIMIT]
        return [entry(m) for m in ordered]

    # ── Helpers ───────────────────────────────────────────────────────────────

    async def _write(self, obj: dict) -> None:
        assert self._proc is not None and self._proc.stdin is not None
        self._proc.stdin.write((json.dumps(obj) + "\n").encode())
        logger.debug(f"Sending to pi: {obj}")
        await self._proc.stdin.drain()


def _join_tool_content(result: object) -> str:
    """Flatten a pi tool result into plain text.

    Shape: {"content": [{"type": "text", "text": "..."}, ...], "details": {...}}.
    """
    if not isinstance(result, dict):
        return str(result) if result else ""
    parts: list[str] = []
    for item in result.get("content", []) or []:
        if isinstance(item, dict) and item.get("type") == "text":
            parts.append(item.get("text", ""))
        elif isinstance(item, str):
            parts.append(item)
    return "\n".join(p for p in parts if p)
