"""ACP agent runner — manages any ACP-compatible agent via JSON-RPC over stdio.

Supports any agent that implements the Agent Client Protocol (ACP), such as
Cline (cline --acp), OpenCode (opencode acp), and others.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import shutil
import tempfile
from typing import Any, Awaitable, Callable

logger = logging.getLogger(__name__)


class ACPAgentRunner:
    """Manages an ACP session for the agentirc daemon.

    Works with any ACP-compatible agent by configuring the spawn command
    via the ``acp_command`` parameter (e.g. ``["cline", "--acp"]`` or
    ``["opencode", "acp"]``).
    """

    def __init__(
        self,
        model: str,
        directory: str,
        acp_command: list[str] | None = None,
        system_prompt: str = "",
        on_exit: Callable[[int], Awaitable[None]] | None = None,
        on_message: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
    ) -> None:
        self.model = model
        self.directory = directory
        self.acp_command = acp_command or ["opencode", "acp"]
        self.system_prompt = system_prompt
        self.on_exit = on_exit
        self.on_message = on_message

        self._isolated_home: str | None = None
        self._process: asyncio.subprocess.Process | None = None
        self._session_id: str | None = None
        self._running = False
        self._busy = False
        self._prompt_queue: asyncio.Queue[str] = asyncio.Queue()
        self._task: asyncio.Task | None = None
        self._reader_task: asyncio.Task | None = None
        self._stderr_task: asyncio.Task | None = None
        self._stopping = False
        self._request_id = 0
        self._pending: dict[int, asyncio.Future] = {}
        self._accumulated_text = ""

    def is_running(self) -> bool:
        return self._running

    @property
    def session_id(self) -> str | None:
        return self._session_id

    async def start(self, initial_prompt: str = "") -> None:
        """Start the ACP agent as a subprocess and initialize a session."""
        self._stopping = False

        # Isolate from host config (XDG, etc.)
        self._isolated_home = tempfile.mkdtemp(prefix="agentirc-acp-")
        isolated_env = dict(os.environ)
        isolated_env["HOME"] = self._isolated_home
        isolated_env.pop("XDG_CONFIG_HOME", None)

        cmd_label = " ".join(self.acp_command)
        try:
            # Spawn ACP agent in stdio mode
            # Use a large stdout buffer — ACP messages (especially session/new with
            # all available models) can exceed asyncio's default 64KB line limit.
            self._process = await asyncio.create_subprocess_exec(
                *self.acp_command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                limit=1024 * 1024,  # 1MB line buffer
                env=isolated_env,
            )

            # Start reading responses and stderr
            self._reader_task = asyncio.create_task(self._read_loop())
            self._stderr_task = asyncio.create_task(self._stderr_loop())

            # Initialize with ACP protocol
            resp = await self._send_request("initialize", {
                "protocolVersion": 1,
                "clientCapabilities": {
                    "fs": {"readTextFile": True, "writeTextFile": True},
                    "terminal": True,
                },
                "clientInfo": {"name": "agentirc-acp", "version": "0.1.0"},
            })
            logger.info("ACP initialized (%s): %s", cmd_label, resp)

            # Create a session with model selection
            session_params = {
                "cwd": self.directory,
                "mcpServers": [],
            }
            if self.model:
                session_params["model"] = self.model
            resp = await self._send_request("session/new", session_params)

            logger.info("ACP session/new raw response: %s", json.dumps(resp)[:500])
            result = resp.get("result", {})
            self._session_id = result.get("sessionId")
            self._running = True
            logger.info("ACP session started (%s): %s", cmd_label, self._session_id)

            # Start the prompt processing loop
            self._task = asyncio.create_task(self._prompt_loop())

            # Queue system prompt as the first turn so all subsequent turns
            # are conditioned on it (ACP has no dedicated system instructions field).
            # Queued rather than awaited to avoid blocking start() on LLM completion.
            if self.system_prompt:
                await self.send_prompt(self.system_prompt)

            if initial_prompt:
                await self.send_prompt(initial_prompt)
        except Exception:
            await self.stop()
            raise

    async def stop(self) -> None:
        """Stop the ACP agent process."""
        self._stopping = True
        self._running = False

        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass

        if self._reader_task:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass

        if self._stderr_task:
            self._stderr_task.cancel()
            try:
                await self._stderr_task
            except asyncio.CancelledError:
                pass

        if self._process:
            try:
                self._process.terminate()
                await asyncio.wait_for(self._process.wait(), timeout=5)
            except (asyncio.TimeoutError, ProcessLookupError):
                try:
                    self._process.kill()
                except ProcessLookupError:
                    pass

        # Fail any pending requests
        for future in self._pending.values():
            if not future.done():
                future.set_exception(ConnectionError("Runner stopped"))
        self._pending.clear()

        if self._isolated_home:
            shutil.rmtree(self._isolated_home, ignore_errors=True)
            self._isolated_home = None

    async def send_prompt(self, text: str) -> None:
        """Queue a prompt for the agent."""
        await self._prompt_queue.put(text)

    # ------------------------------------------------------------------
    # Internal: JSON-RPC over stdio
    # ------------------------------------------------------------------

    def _next_id(self) -> int:
        self._request_id += 1
        return self._request_id

    async def _send_request(self, method: str, params: dict, timeout: float = 30) -> dict:
        """Send a JSON-RPC request and wait for the response."""
        if not self._process or not self._process.stdin:
            raise ConnectionError("ACP server not running")

        req_id = self._next_id()
        msg = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params}

        loop = asyncio.get_running_loop()
        future: asyncio.Future[dict] = loop.create_future()
        self._pending[req_id] = future

        line = json.dumps(msg) + "\n"
        self._process.stdin.write(line.encode())
        await self._process.stdin.drain()

        try:
            return await asyncio.wait_for(future, timeout=timeout)
        except (asyncio.TimeoutError, asyncio.CancelledError):
            self._pending.pop(req_id, None)
            if not future.done():
                future.cancel()
            raise

    async def _send_notification(self, method: str, params: dict) -> None:
        """Send a JSON-RPC notification (no response expected)."""
        if not self._process or not self._process.stdin:
            return
        msg = {"jsonrpc": "2.0", "method": method, "params": params}
        line = json.dumps(msg) + "\n"
        self._process.stdin.write(line.encode())
        await self._process.stdin.drain()

    async def _read_loop(self) -> None:
        """Read JSON-RPC messages from stdout."""
        if not self._process or not self._process.stdout:
            return
        try:
            while True:
                line = await self._process.stdout.readline()
                if not line:
                    break
                try:
                    msg = json.loads(line.decode())
                except json.JSONDecodeError:
                    continue

                # Response to a request
                if "id" in msg and ("result" in msg or "error" in msg):
                    req_id = msg["id"]
                    future = self._pending.pop(req_id, None)
                    if future and not future.done():
                        future.set_result(msg)

                # Notification from server
                elif "method" in msg:
                    await self._handle_notification(msg)

        except (asyncio.CancelledError, ConnectionError):
            pass
        except Exception:
            logger.exception("ACP read loop error")
        finally:
            # Wait for process to fully exit and notify daemon for crash recovery
            returncode = -1
            if self._process:
                try:
                    returncode = await asyncio.wait_for(
                        self._process.wait(), timeout=5
                    )
                except asyncio.TimeoutError:
                    try:
                        self._process.kill()
                    except ProcessLookupError:
                        pass

            # Fail any still-pending requests
            for future in self._pending.values():
                if not future.done():
                    future.set_exception(ConnectionError("Process exited"))
            self._pending.clear()

            self._running = False

            # Cancel companion tasks so they don't outlive the process
            for task in (self._task, self._stderr_task):
                if task and not task.done():
                    task.cancel()

            if not self._stopping and self.on_exit:
                await self.on_exit(returncode)

    async def _stderr_loop(self) -> None:
        """Log stderr output from the ACP agent process."""
        if not self._process or not self._process.stderr:
            return
        cmd_name = self.acp_command[0]
        try:
            while True:
                line = await self._process.stderr.readline()
                if not line:
                    break
                text = line.decode("utf-8", errors="replace").rstrip()
                if text:
                    logger.warning("acp[%s] stderr: %s", cmd_name, text)
        except (asyncio.CancelledError, ConnectionError):
            pass

    async def _handle_notification(self, msg: dict) -> None:
        """Handle ACP server notifications."""
        method = msg.get("method", "")
        params = msg.get("params", {})

        if method == "session/update":
            # ACP wraps updates in params.update (not flat in params)
            update = params.get("update", params)
            update_type = update.get("sessionUpdate", "")

            if update_type in ("agent_message_chunk", "agent_thought_chunk"):
                # Accumulate streaming text (both message and thought chunks)
                self._busy = True
                content = update.get("content", {})
                if update_type == "agent_message_chunk" and content.get("type") == "text":
                    self._accumulated_text += content.get("text", "")

            # Check for stopReason — indicates turn is complete
            if "stopReason" in update:
                self._busy = False
                if self.on_message and self._accumulated_text:
                    msg_dict = {
                        "type": "assistant",
                        "model": self.model,
                        "content": [{"type": "text", "text": self._accumulated_text}],
                    }
                    await self.on_message(msg_dict)
                self._accumulated_text = ""

        elif method == "session/request_permission":
            # Auto-approve all permission requests (file changes, commands, etc.)
            req_id = msg.get("id")
            if req_id is not None:
                resp = {"jsonrpc": "2.0", "id": req_id, "result": {"approved": True}}
                if self._process and self._process.stdin:
                    line = json.dumps(resp) + "\n"
                    self._process.stdin.write(line.encode())
                    await self._process.stdin.drain()

        elif method == "error":
            logger.error("ACP error (%s): %s", self.acp_command[0], params)

    async def _prompt_loop(self) -> None:
        """Process queued prompts one at a time."""
        try:
            while self._running:
                text = await self._prompt_queue.get()
                if not self._running or not self._session_id:
                    break

                try:
                    self._busy = True
                    resp = await self._send_request("session/prompt", {
                        "sessionId": self._session_id,
                        "prompt": [{"type": "text", "text": text}],
                    }, timeout=120)

                    # Check if response itself signals turn completion
                    result = resp.get("result", {})
                    if "stopReason" in result:
                        # Turn completed — flush any accumulated text
                        if self.on_message and self._accumulated_text:
                            msg_dict = {
                                "type": "assistant",
                                "model": self.model,
                                "content": [{"type": "text", "text": self._accumulated_text}],
                            }
                            await self.on_message(msg_dict)
                        self._accumulated_text = ""
                        self._busy = False

                    # Wait for turn to complete (via notifications)
                    while self._busy and self._running:
                        await asyncio.sleep(0.1)

                except Exception:
                    logger.exception("ACP turn error")

        except asyncio.CancelledError:
            pass
