"""Pi-RPC subprocess client for communicating with the pi coding agent."""

from __future__ import annotations

import json
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path

from abreka.errors import PiRpcError

# Pi-rpc streams tool-call args as `message_update` events whose nested
# `assistantMessageEvent.type` is one of these. Each such event is a
# cumulative snapshot of the full args-so-far, not an incremental diff,
# producing O(N²) bytes per tool call. The final args already live in the
# `ame:toolcall_end` event (and the top-level `tool_execution_end`), and
# nothing in abreka reads these deltas back — the webui parser only consumes
# `thinking_delta` and `text_delta` subtypes of `message_update`.
DROPPED_AME_SUBTYPES: frozenset[str] = frozenset({"toolcall_delta"})

SESSION_CAPTURE_MODES: frozenset[str] = frozenset({"full", "slim", "minimal"})

# Minimal mode keeps only assistant text deltas (no thinking trace) plus
# toolcall start/end metadata from message_update.
_MINIMAL_AME_SUBTYPES: frozenset[str] = frozenset(
    {"text_start", "text_delta", "text_end", "toolcall_start", "toolcall_end"}
)


def normalize_session_capture_mode(mode: str) -> str:
    """Normalize/validate session capture mode."""
    normalized = (mode or "slim").strip().lower()
    if normalized not in SESSION_CAPTURE_MODES:
        allowed = ", ".join(sorted(SESSION_CAPTURE_MODES))
        raise ValueError(
            f"Invalid session capture mode: {mode!r}. Expected one of: {allowed}."
        )
    return normalized


def is_dropped_event(event: dict) -> bool:
    """Return True if *event* is a redundant cumulative-snapshot subtype.

    Top-level `message_update` events whose nested assistantMessageEvent
    type is in DROPPED_AME_SUBTYPES are pure streaming bloat. Everything
    else passes through.
    """
    if event.get("type") != "message_update":
        return False
    ame = event.get("assistantMessageEvent")
    if not isinstance(ame, dict):
        return False
    return ame.get("type") in DROPPED_AME_SUBTYPES


def _strip_message_update_snapshots(event: dict) -> dict:
    """Return a copy of a message_update event without cumulative snapshots."""
    compacted = dict(event)
    compacted.pop("message", None)

    ame = compacted.get("assistantMessageEvent")
    if isinstance(ame, dict) and "partial" in ame:
        ame_compacted = dict(ame)
        ame_compacted.pop("partial", None)
        compacted["assistantMessageEvent"] = ame_compacted

    return compacted


def compact_event_for_session(event: dict, mode: str = "slim") -> dict | None:
    """Transform *event* for persisted session storage.

    Returns:
        - None: drop the event from persisted session data.
        - dict: event to persist (possibly compacted).
    """
    capture_mode = normalize_session_capture_mode(mode)

    if is_dropped_event(event):
        return None

    if capture_mode == "full":
        return event

    if event.get("type") != "message_update":
        return event

    ame = event.get("assistantMessageEvent")
    if capture_mode == "minimal":
        if not isinstance(ame, dict):
            return None
        if ame.get("type") not in _MINIMAL_AME_SUBTYPES:
            return None

    return _strip_message_update_snapshots(event)


@dataclass
class PiResult:
    """Result from a pi-rpc prompt."""
    text: str
    events: list[dict]


@dataclass
class PiRpcClient:
    """Client that spawns and communicates with pi in RPC mode."""

    provider: str
    model: str
    cwd: Path
    thinking_level: str = "xhigh"
    extra_args: list[str] = field(default_factory=list)
    debug: bool = False
    session_path: Path | None = None
    session_capture_mode: str = "slim"
    collect_result_events: bool = False
    _process: subprocess.Popen | None = field(default=None, init=False, repr=False)
    _all_events: list[dict] = field(default_factory=list, init=False, repr=False)
    _eof: bool = field(default=False, init=False, repr=False)

    def __post_init__(self) -> None:
        self.session_capture_mode = normalize_session_capture_mode(self.session_capture_mode)

    def start(self) -> None:
        """Spawn the pi subprocess in RPC mode."""
        cmd = [
            "pi", "--mode", "rpc",
            "--provider", self.provider,
            "--model", self.model,
            "--thinking", self.thinking_level,
            *self.extra_args,
        ]
        if self.session_path is not None:
            self.session_path.parent.mkdir(parents=True, exist_ok=True)
            cmd.extend(["--session", str(self.session_path)])
        else:
            cmd.append("--no-session")
        if self.debug:
            print(f"[pi-rpc] spawning: {' '.join(cmd)}", file=sys.stderr)
            print(f"[pi-rpc] cwd: {self.cwd}", file=sys.stderr)
        try:
            self._process = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                cwd=self.cwd,
            )
        except FileNotFoundError:
            raise PiRpcError(
                "pi command not found. Install pi-rpc: https://github.com/badlogic/pi-mono"
            )

    def _send(self, command: dict) -> None:
        """Send a JSON command to pi's stdin."""
        if self._process is None or self._process.stdin is None:
            raise PiRpcError("Pi process not started. Call start() first.")
        try:
            line = json.dumps(command) + "\n"
            if self.debug:
                print(f"[pi-rpc] >>> {line.strip()}", file=sys.stderr)
            self._process.stdin.write(line)
            self._process.stdin.flush()
        except BrokenPipeError:
            stderr = ""
            if self._process.stderr:
                stderr = self._process.stderr.read()
            raise PiRpcError(f"Pi process terminated unexpectedly. stderr: {stderr}")

    def _read_event(self) -> dict | None:
        """Read one JSONL event from pi's stdout. Returns None on EOF or skip."""
        if self._process is None or self._process.stdout is None:
            self._eof = True
            return None
        line = self._process.stdout.readline()
        if not line:
            # Empty string from readline() means EOF — pipe is closed
            self._eof = True
            return None
        line = line.strip()
        if not line:
            return None
        # Only parse lines that look like JSON
        if not line.startswith("{"):
            if self.debug:
                print(f"[pi-rpc] non-json: {line[:200]}", file=sys.stderr)
            return None
        try:
            event = json.loads(line)
            if self.debug:
                etype = event.get("type", "?")
                # Compact debug: don't dump full text deltas
                if etype == "message_update":
                    print(f"[pi-rpc] <<< {etype} (text delta)", file=sys.stderr)
                else:
                    compact = json.dumps(event)
                    if len(compact) > 300:
                        compact = compact[:300] + "..."
                    print(f"[pi-rpc] <<< {compact}", file=sys.stderr)
            return event
        except json.JSONDecodeError:
            if self.debug:
                print(f"[pi-rpc] bad json: {line[:200]}", file=sys.stderr)
            return None

    def prompt(self, message: str, on_event: callable | None = None) -> PiResult:
        """Send a prompt and collect events until agent_end.

        Args:
            message: The prompt text to send.
            on_event: Optional callback for each event (for live streaming).

        Returns:
            PiResult with the full assistant text and optionally collected raw
            events (empty unless ``collect_result_events=True``).
        """
        self._send({"type": "prompt", "message": message})

        events: list[dict] | None = [] if self.collect_result_events else None
        text_parts: list[str] = []

        while True:
            event = self._read_event()
            if event is None:
                if self._eof:
                    # Pipe closed — agent is done (may not have sent agent_end)
                    if self.debug:
                        print("[pi-rpc] EOF on stdout, ending", file=sys.stderr)
                    break
                # Non-JSON line or empty — skip
                continue

            if events is not None:
                events.append(event)
            self._record_event(event)

            if on_event:
                on_event(event)

            # Extract text deltas
            event_type = event.get("type")
            if event_type == "message_update":
                ame = event.get("assistantMessageEvent", {})
                if ame.get("type") == "text_delta":
                    text_parts.append(ame.get("delta", ""))

            # Check for completion
            if event_type == "agent_end":
                break
            # Break on turn_end only if it's an error (no agent_end will follow)
            if event_type == "turn_end":
                msg = event.get("message", {})
                if msg.get("stopReason") == "error":
                    break

        return PiResult(text="".join(text_parts), events=events or [])

    def _record_event(self, event: dict) -> None:
        compacted = compact_event_for_session(event, self.session_capture_mode)
        if compacted is None:
            return
        self._all_events.append(compacted)

    def save_session(self, path: Path) -> None:
        """Save all collected events to a JSON file."""
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(self._all_events, f, indent=2)
            f.write("\n")

    def stop(self) -> None:
        """Terminate the pi subprocess."""
        if self._process is not None:
            try:
                if self._process.stdin:
                    self._process.stdin.close()
                self._process.terminate()
                self._process.wait(timeout=5)
            except Exception:
                self._process.kill()
            finally:
                self._process = None

    def __enter__(self) -> PiRpcClient:
        self.start()
        return self

    def __exit__(self, *args) -> None:
        self.stop()
