"""
Synchronous ACP Client for Augment CLI

A clean, easy-to-use wrapper around the Agent Client Protocol for communicating
with the Augment CLI agent.
"""

import asyncio
import os
from abc import ABC, abstractmethod
from queue import Empty, Queue
from threading import Thread
from typing import Optional, Any, List

from acp import (
    Client,
    ClientSideConnection,
    InitializeRequest,
    NewSessionRequest,
    PromptRequest,
    RequestError,
    SessionNotification,
    text_block,
    PROTOCOL_VERSION,
    spawn_agent_process,
)
from acp.schema import (
    RequestPermissionRequest,
    RequestPermissionResponse,
    AllowedOutcome,
)


class AgentEventListener(ABC):
    """
    Interface for listening to agent events.

    Implement this interface to receive notifications about what the agent is doing.
    """

    @abstractmethod
    def on_agent_message_chunk(self, text: str) -> None:
        """
        Called when the agent sends a message chunk (streaming).

        The agent streams its response in real-time. This method is called
        multiple times with small chunks of text that together form the
        complete message.

        Args:
            text: A chunk of text from the agent's response
        """
        pass

    @abstractmethod
    def on_tool_call(
        self,
        tool_call_id: str,
        title: str,
        kind: Optional[str] = None,
        status: Optional[str] = None,
    ) -> None:
        """
        Called when the agent makes a tool call.

        Args:
            tool_call_id: Unique identifier for this tool call
            title: Human-readable description of what the tool is doing
            kind: Category of tool (read, edit, delete, execute, etc.)
            status: Current status (pending, in_progress, completed, failed)
        """
        pass

    @abstractmethod
    def on_tool_response(
        self,
        tool_call_id: str,
        status: Optional[str] = None,
        content: Optional[Any] = None,
    ) -> None:
        """
        Called when a tool responds with results.

        Args:
            tool_call_id: Unique identifier for this tool call
            status: Response status (completed, failed, etc.)
            content: Response content/results from the tool
        """
        pass

    def on_agent_thought(self, text: str) -> None:
        """
        Called when the agent shares its internal reasoning.

        Args:
            text: The thought content from the agent
        """
        pass

    def on_agent_message(self, message: str) -> None:
        """
        Called when the agent finishes sending a complete message.

        This is called once after all message chunks have been sent,
        with the complete assembled message.

        Args:
            message: The complete message from the agent
        """
        pass


class _InternalACPClient(Client):
    """Internal ACP client implementation."""

    def __init__(self, listener: Optional[AgentEventListener] = None):
        self.listener = listener
        self.last_response = ""

    async def requestPermission(
        self, params: RequestPermissionRequest
    ) -> RequestPermissionResponse:
        """
        Handle permission requests from the CLI (e.g., indexing permission).
        Auto-approves all requests by selecting the first "allow" option.
        """
        # Find the first "allow" option (allow_once or allow_always)
        allow_option = None
        for option in params.options:
            if option.kind.startswith("allow"):
                allow_option = option
                break

        if not allow_option and params.options:
            # If no allow option found, just select the first option
            allow_option = params.options[0]

        if not allow_option:
            # No options available, return a default
            return RequestPermissionResponse(
                outcome=AllowedOutcome(optionId="default", outcome="selected")
            )

        # Return approval response using proper ACP types
        return RequestPermissionResponse(
            outcome=AllowedOutcome(optionId=allow_option.optionId, outcome="selected")
        )

    async def writeTextFile(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("fs/write_text_file")

    async def readTextFile(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("fs/read_text_file")

    async def createTerminal(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("terminal/create")

    async def terminalOutput(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("terminal/output")

    async def releaseTerminal(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("terminal/release")

    async def waitForTerminalExit(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("terminal/wait_for_exit")

    async def killTerminal(self, params):  # type: ignore[override]
        raise RequestError.method_not_found("terminal/kill")

    async def sessionUpdate(self, params: SessionNotification) -> None:
        update = params.update
        if isinstance(update, dict):
            kind = update.get("sessionUpdate")
            content = update.get("content")
        else:
            kind = getattr(update, "sessionUpdate", None)
            content = getattr(update, "content", None)

        # Handle agent message chunks
        if kind == "agent_message_chunk" and content is not None:
            if isinstance(content, dict):
                text = content.get("text", "")
            else:
                text = getattr(content, "text", "")
            self.last_response += text
            if self.listener:
                self.listener.on_agent_message_chunk(text)

        # Handle agent thoughts
        elif kind == "agent_thought_chunk" and content is not None:
            if isinstance(content, dict):
                text = content.get("text", "")
            else:
                text = getattr(content, "text", "")
            if self.listener:
                self.listener.on_agent_thought(text)

        # Handle tool calls
        elif kind == "tool_call":
            tool_call_id = getattr(update, "toolCallId", "unknown")
            title = getattr(update, "title", "")
            tool_kind = getattr(update, "kind", None)
            status = getattr(update, "status", None)
            if self.listener:
                self.listener.on_tool_call(tool_call_id, title, tool_kind, status)

        # Handle tool responses
        elif kind == "tool_call_update":
            tool_call_id = getattr(update, "toolCallId", "unknown")
            status = getattr(update, "status", None)
            content = getattr(update, "content", None)
            if self.listener:
                self.listener.on_tool_response(tool_call_id, status, content)

    async def extMethod(self, method: str, params: dict) -> dict:  # noqa: ARG002
        raise RequestError.method_not_found(method)

    async def extNotification(self, method: str, params: dict) -> None:  # noqa: ARG002
        raise RequestError.method_not_found(method)

    def get_last_response(self) -> str:
        return self.last_response.strip()


class ACPClient:
    """ACP client interface."""

    def start(self) -> None:
        """
        Start the agent process and establish ACP connection.

        Raises:
            RuntimeError: If the agent is already started
            Exception: If initialization fails
        """
        raise NotImplementedError()

    def stop(self) -> None:
        """Stop the agent process and cleanup resources."""
        raise NotImplementedError()

    def send_message(self, message: str, timeout: float = 30.0) -> str:
        """
        Send a message to the agent and get the response.

        Args:
            message: The message to send
            timeout: Maximum time to wait for response (seconds)

        Returns:
            The agent's response as a string

        Raises:
            RuntimeError: If the agent is not started
            TimeoutError: If the response takes too long
        """
        raise NotImplementedError()

    def clear_context(self) -> None:
        """
        Clear the session context by restarting the agent.

        This stops the current agent and starts a new one with a fresh session.
        """
        raise NotImplementedError()

    @property
    def is_running(self) -> bool:
        """Check if the agent is currently running."""
        raise NotImplementedError()


class AuggieACPClient(ACPClient):
    """
    Synchronous ACP client for the Augment CLI agent.

    This client provides a simple interface for:
    - Starting/stopping the agent
    - Sending messages and getting responses
    - Listening to agent events (messages, tool calls, etc.)
    - Clearing session context

    Example:
        ```python
        # Create a client with an event listener
        client = AuggieACPClient(listener=MyListener())

        # Start the agent
        client.start()

        # Send a message
        response = client.send_message("What is 2 + 2?")
        print(response)

        # Clear context and start fresh
        client.clear_context()

        # Stop the agent
        client.stop()
        ```
    """

    def __init__(
        self,
        cli_path: Optional[str] = None,
        listener: Optional[AgentEventListener] = None,
        model: Optional[str] = None,
        workspace_root: Optional[str] = None,
        acp_max_tool_result_bytes: int = 35 * 1024,
        removed_tools: Optional[List[str]] = None,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
    ):
        """
        Initialize the ACP client.

        Args:
            cli_path: Path to the Augment CLI. If None, uses default location.
            listener: Optional event listener to receive agent events.
            model: AI model to use (e.g., "claude-3-5-sonnet-latest", "gpt-4o").
                   If None, uses the CLI's default model.
            workspace_root: Workspace root directory. If None, uses current directory.
            acp_max_tool_result_bytes: Maximum bytes for tool results sent through ACP.
                   Default: 35KB (35840 bytes). This prevents Python's asyncio.StreamReader
                   from hitting its 64KB line buffer limit. Set higher if needed, but
                   keep under 64KB to avoid LimitOverrunError.
            removed_tools: List of tool names to remove/disable (e.g., ["github-api", "linear"]).
                   These tools will not be available to the agent.
            api_key: Optional API key for authentication. If provided, sets AUGMENT_API_TOKEN
                   environment variable for the agent process.
            api_url: Optional API URL. If not provided, uses AUGMENT_API_URL environment variable,
                   or defaults to "https://api.augmentcode.com". Sets AUGMENT_API_URL environment
                   variable for the agent process.
        """
        if cli_path is None:
            # Default to 'auggie' in PATH
            cli_path = "auggie"

        self.cli_path = cli_path

        self.listener = listener
        self.model = model
        self.workspace_root = workspace_root
        self.acp_max_tool_result_bytes = acp_max_tool_result_bytes
        self.removed_tools = removed_tools or []
        self.api_key = api_key
        self.api_url = (
            api_url
            if api_url is not None
            else os.getenv("AUGMENT_API_URL", "https://api.augmentcode.com")
        )
        self._client: Optional[_InternalACPClient] = None
        self._conn: Optional[ClientSideConnection] = None
        self._session_id: Optional[str] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._thread: Optional[Thread] = None
        self._context = None
        self._ready_queue: Optional[Queue] = None

    def start(self, timeout: float = 30.0) -> None:
        """
        Start the agent process and establish ACP connection.

        Args:
            timeout: Maximum time to wait for the agent to start (seconds)

        Raises:
            RuntimeError: If the agent is already started
            TimeoutError: If the agent fails to start within the timeout
            Exception: If initialization fails
        """
        if self._thread is not None:
            raise RuntimeError("Agent already started")

        self._ready_queue = Queue()
        self._thread = Thread(target=self._run_async_loop, daemon=True)
        self._thread.start()

        # Wait for initialization with timeout
        try:
            result = self._ready_queue.get(timeout=timeout)
            if isinstance(result, Exception):
                raise result
        except Empty:
            # Queue.get() timed out - no result received within timeout
            raise TimeoutError(
                f"Agent failed to start within {timeout} seconds. "
                f"Check that the CLI path is correct and the agent process can start. "
                f"CLI path: {self.cli_path}"
            )

    def stop(self) -> None:
        """Stop the agent process and cleanup resources."""
        if self._loop is not None:
            asyncio.run_coroutine_threadsafe(self._async_stop(), self._loop)
            self._loop.call_soon_threadsafe(self._loop.stop)
            if self._thread is not None:
                self._thread.join(timeout=2.0)

        self._client = None
        self._conn = None
        self._session_id = None
        self._loop = None
        self._thread = None
        self._context = None

    def send_message(self, message: str, timeout: float = 30.0) -> str:
        """
        Send a message to the agent and get the response.

        Args:
            message: The message to send
            timeout: Maximum time to wait for response (seconds)

        Returns:
            The agent's response as a string

        Raises:
            RuntimeError: If the agent is not started
            TimeoutError: If the response takes too long
        """
        if self._loop is None or self._conn is None or self._client is None:
            raise RuntimeError("Agent not started. Call start() first.")

        # Reset the response
        self._client.last_response = ""

        # Schedule the async query
        future = asyncio.run_coroutine_threadsafe(
            self._async_send_message(message), self._loop
        )

        # Wait for completion - when prompt() completes, the message is done
        # (per ACP spec, there is no agent_message_end event)
        future.result(timeout=timeout)

        response = self._client.get_last_response()

        # Call listener with complete message now that prompt() has completed
        if self.listener and response:
            self.listener.on_agent_message(response)

        return response

    def clear_context(self) -> None:
        """
        Clear the session context by restarting the agent.

        This stops the current agent and starts a new one with a fresh session.
        """
        self.stop()
        self.start()

    @property
    def session_id(self) -> Optional[str]:
        """Get the current session ID."""
        return self._session_id

    @property
    def is_running(self) -> bool:
        """Check if the agent is currently running."""
        return self._thread is not None and self._loop is not None

    def _run_async_loop(self):
        """Run the asyncio event loop in a background thread."""
        try:
            self._loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self._loop)
            self._loop.run_until_complete(self._async_start())
            self._ready_queue.put(True)
            self._loop.run_forever()
        except Exception as e:
            self._ready_queue.put(e)

    async def _async_start(self):
        """Async initialization."""
        self._client = _InternalACPClient(self.listener)

        # Build CLI arguments
        cli_path_str = str(self.cli_path)
        if cli_path_str.endswith(".mjs") or cli_path_str.endswith(".js"):
            # It's a JS file, run with node
            cli_args = ["node", cli_path_str]
        else:
            # It's a binary or script (like 'auggie'), run directly
            cli_args = [cli_path_str]

        cli_args.extend(
            [
                "--acp",
                "--log-file",
                "./auggie-acp.log",
            ]
        )

        # Add model if specified
        if self.model:
            cli_args.extend(["--model", self.model])

        # Add workspace root if specified
        if self.workspace_root:
            cli_args.extend(["--workspace-root", self.workspace_root])

        # Add removed tools if specified and enabled
        if self.removed_tools:
            for tool in self.removed_tools:
                cli_args.extend(["--remove-tool", tool])

        # Add ACP max tool result bytes
        # TODO: Re-enable once --acp-max-tool-result-bytes is in pre-release
        # cli_args.extend(
        #     ["--acp-max-tool-result-bytes", str(self.acp_max_tool_result_bytes)]
        # )

        # Set environment variables for API authentication
        # Build environment dict to pass to spawn_agent_process
        env = os.environ.copy()
        if self.api_key:
            env["AUGMENT_API_TOKEN"] = self.api_key
        if self.api_url:
            env["AUGMENT_API_URL"] = self.api_url

        # Spawn the agent process with environment variables
        self._context = spawn_agent_process(
            lambda _agent: self._client, *cli_args, env=env
        )

        # Start the process and get connection
        conn_proc = await self._context.__aenter__()
        self._conn, self._proc = conn_proc

        # Create a task to monitor if the process exits early
        async def wait_for_process_exit():
            """Wait for the process to exit and raise an error if it does."""
            await self._proc.wait()
            stderr = ""
            if self._proc.stderr:
                try:
                    stderr_bytes = await asyncio.wait_for(
                        self._proc.stderr.read(), timeout=1.0
                    )
                    stderr = stderr_bytes.decode("utf-8", errors="replace")
                except Exception:
                    pass
            raise RuntimeError(
                f"Agent process exited with code {self._proc.returncode}. "
                f"CLI path: {self.cli_path}\n"
                f"Stderr: {stderr}"
            )

        # Check if process has already exited
        if self._proc.returncode is not None:
            # Process already exited
            stderr = ""
            if self._proc.stderr:
                stderr_bytes = await self._proc.stderr.read()
                stderr = stderr_bytes.decode("utf-8", errors="replace")
            raise RuntimeError(
                f"Agent process exited immediately with code {self._proc.returncode}. "
                f"CLI path: {self.cli_path}\n"
                f"Stderr: {stderr}"
            )

        # Create process monitor task
        monitor_task = asyncio.create_task(wait_for_process_exit())

        try:
            # Race between initialization and process exit
            init_task = asyncio.create_task(
                self._conn.initialize(
                    InitializeRequest(
                        protocolVersion=PROTOCOL_VERSION, clientCapabilities=None
                    )
                )
            )
            done, pending = await asyncio.wait(
                [init_task, monitor_task], return_when=asyncio.FIRST_COMPLETED
            )

            # If monitor task completed first, it means process exited
            if monitor_task in done:
                # Cancel the init task
                init_task.cancel()
                # Re-raise the exception from monitor_task
                await monitor_task

            # Otherwise, initialization succeeded
            await init_task

            # Use workspace_root as cwd if provided, otherwise use current directory
            cwd = self.workspace_root if self.workspace_root else os.getcwd()

            # Race between session creation and process exit
            session_task = asyncio.create_task(
                self._conn.newSession(NewSessionRequest(mcpServers=[], cwd=cwd))
            )
            done, pending = await asyncio.wait(
                [session_task, monitor_task], return_when=asyncio.FIRST_COMPLETED
            )

            # If monitor task completed first, it means process exited
            if monitor_task in done:
                # Cancel the session task
                session_task.cancel()
                # Re-raise the exception from monitor_task
                await monitor_task

            # Otherwise, session creation succeeded
            session = await session_task
            self._session_id = session.sessionId

            # Keep the monitor task running in the background
            # (don't cancel it, it will keep monitoring the process)
        except Exception:
            # If anything fails, cancel the monitor task
            monitor_task.cancel()
            raise

    async def _async_send_message(self, message: str):
        """Async message sending."""
        await self._conn.prompt(
            PromptRequest(
                sessionId=self._session_id,
                prompt=[text_block(message)],
            )
        )

    async def _async_stop(self):
        """Async cleanup."""
        if self._context is not None:
            await self._context.__aexit__(None, None, None)

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()
