"""ACP client — initialize, session management, and prompt lifecycle.

Higher-level wrapper around :class:`ACPTransport` that handles the full ACP
protocol lifecycle: initialize handshake, auth detection, session creation,
prompt/response streaming, cancellation, and shutdown.
"""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any

from acp import PROTOCOL_VERSION, text_block
from acp.client import ClientSideConnection
from acp.schema import (
    AgentCapabilities,
    AgentMessageChunk,
    CreateTerminalResponse,
    PermissionOption,
    ReadTextFileResponse,
    ReleaseTerminalResponse,
    TerminalOutputResponse,
    TextContentBlock,
    ToolCallProgress,
    ToolCallStart,
    ToolCallUpdate,
    WaitForTerminalExitResponse,
    WriteTextFileResponse,
)

from otto.acp import ACPAuthRequiredError
from otto.acp.transport import ACPTransport
from otto.log import get_logger

_log = get_logger("otto.acp.client")

# Sentinel pushed into the update queue to signal that prompt() is done.
_DONE = object()


# ---------------------------------------------------------------------------
# Update types yielded by ACPClient.prompt()
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class ACPTextChunk:
    """A chunk of agent text from an ``agent_message_chunk`` update."""

    text: str


@dataclass(frozen=True, slots=True)
class ACPToolStatus:
    """A tool call status change (start or progress)."""

    name: str
    status: str
    call_id: str


@dataclass(frozen=True, slots=True)
class ACPPromptDone:
    """End of a prompt turn with the stop reason."""

    stop_reason: str


# Union of all update types yielded by prompt().
ACPUpdate = ACPTextChunk | ACPToolStatus | ACPPromptDone


# ---------------------------------------------------------------------------
# OttoACPClientHandler — SDK Client subclass
# ---------------------------------------------------------------------------


class OttoACPClientHandler:
    """Implements the ACP SDK ``Client`` interface.

    Dispatches session updates into an :class:`asyncio.Queue` so that
    :meth:`ACPClient.prompt` can yield them as an async iterator.
    """

    def __init__(self) -> None:
        self._queue: asyncio.Queue[ACPUpdate | object] | None = None

    def bind_queue(self, queue: asyncio.Queue[ACPUpdate | object]) -> None:
        self._queue = queue

    def unbind_queue(self) -> None:
        self._queue = None

    # -- SDK Client interface -----------------------------------------------

    def on_connect(self, conn: Any) -> None:
        pass

    async def session_update(
        self,
        session_id: str,
        update: Any,
        **kwargs: Any,
    ) -> None:
        if isinstance(update, AgentMessageChunk):
            if isinstance(update.content, TextContentBlock):
                item = ACPTextChunk(text=update.content.text)
                if self._queue is not None:
                    self._queue.put_nowait(item)
        elif isinstance(update, ToolCallStart):
            item = ACPToolStatus(
                name=update.title,
                status=update.status or "pending",
                call_id=update.tool_call_id,
            )
            if self._queue is not None:
                self._queue.put_nowait(item)
        elif isinstance(update, ToolCallProgress):
            item = ACPToolStatus(
                name=update.title or "",
                status=update.status or "unknown",
                call_id=update.tool_call_id,
            )
            if self._queue is not None:
                self._queue.put_nowait(item)

    async def request_permission(
        self,
        options: list[PermissionOption],
        session_id: str,
        tool_call: ToolCallUpdate,
        **kwargs: Any,
    ) -> dict[str, Any]:
        # Auto-approve with the first ``allow_once`` option if available,
        # falling back to the first option.
        _log.info(
            "acp permission request (auto-approving)",
            session_id=session_id,
            tool_call_id=getattr(tool_call, "tool_call_id", None),
            options=[o.name for o in options],
        )
        allow_opts = [o for o in options if o.kind == "allow_once"]
        chosen = allow_opts[0] if allow_opts else options[0]
        return {"outcome": {"outcome": "selected", "optionId": chosen.option_id}}

    # -- Stubs for methods required by the Client protocol but not used -----

    async def write_text_file(
        self, content: str, path: str, session_id: str, **kwargs: Any
    ) -> WriteTextFileResponse | None:
        return None

    async def read_text_file(
        self,
        path: str,
        session_id: str,
        limit: int | None = None,
        line: int | None = None,
        **kwargs: Any,
    ) -> ReadTextFileResponse:
        return ReadTextFileResponse(content="")

    async def create_terminal(
        self, command: str, session_id: str, args: list[str] | None = None, **kwargs: Any
    ) -> CreateTerminalResponse:
        return CreateTerminalResponse(terminal_id="unsupported")

    async def terminal_output(
        self, session_id: str, terminal_id: str, **kwargs: Any
    ) -> TerminalOutputResponse:
        return TerminalOutputResponse(output="")

    async def release_terminal(
        self, session_id: str, terminal_id: str, **kwargs: Any
    ) -> ReleaseTerminalResponse | None:
        return None

    async def wait_for_terminal_exit(
        self, session_id: str, terminal_id: str, **kwargs: Any
    ) -> WaitForTerminalExitResponse:
        return WaitForTerminalExitResponse(exit_code=1)

    async def kill_terminal(self, session_id: str, terminal_id: str, **kwargs: Any) -> None:
        return None


# ---------------------------------------------------------------------------
# ACPClient — public interface
# ---------------------------------------------------------------------------


class ACPClient:
    """High-level client for communicating with an ACP agent subprocess.

    Usage::

        async with ACPClient(command=["python3", "agent.py"]) as client:
            session_id = await client.new_session(cwd="/workspace")
            async for update in client.prompt(session_id, "Hello!"):
                if isinstance(update, ACPTextChunk):
                    print(update.text, end="")
    """

    def __init__(
        self,
        command: list[str],
        env: dict[str, str] | None = None,
        cwd: str | None = None,
    ) -> None:
        self._command = command
        self._env = env
        self._cwd = cwd

        self._handler = OttoACPClientHandler()
        self._transport = ACPTransport(
            self._handler,
            command=command,
            env=env,
            cwd=cwd,
        )
        self._capabilities: AgentCapabilities | None = None
        self._connection: ClientSideConnection | None = None
        self._started = False

    @property
    def capabilities(self) -> AgentCapabilities | None:
        return self._capabilities

    @property
    def started(self) -> bool:
        return self._started

    async def start(self) -> None:
        """Spawn the agent subprocess and perform the initialize handshake.

        Raises :class:`ACPAuthRequiredError` if the agent requires
        authentication that the client cannot handle automatically.
        """
        conn, _proc = await self._transport.start()
        self._connection = conn

        init_resp = await conn.initialize(protocol_version=PROTOCOL_VERSION)
        self._capabilities = init_resp.agent_capabilities
        self._started = True

        _log.info(
            "acp client initialized",
            protocol_version=init_resp.protocol_version,
            capabilities=str(self._capabilities),
            auth_methods=len(init_resp.auth_methods or []),
        )

        # Check auth requirements.
        auth_methods = init_resp.auth_methods or []
        if auth_methods:
            self._check_auth(auth_methods)

    def _check_auth(self, auth_methods: list[Any]) -> None:
        """Raise ACPAuthRequiredError if the agent needs auth we cannot handle."""
        # The AuthMethod schema has .id, .name, .description.
        # The *type* of auth (environment_variable, terminal, agent) is
        # communicated via the method id / name.  We inspect the id and name
        # for known keywords to decide the auth type.
        for method in auth_methods:
            method_id = getattr(method, "id", "")
            method_name = getattr(method, "name", "")
            combined = f"{method_id} {method_name}".lower()

            if "environment_variable" in combined or "env" == method_id:
                raise ACPAuthRequiredError(
                    f"Agent requires environment variable authentication. "
                    f"Set the required variables and restart. "
                    f"Method: {method_name} ({method_id})",
                    auth_methods=auth_methods,
                )
            if "terminal" in combined:
                raise ACPAuthRequiredError(
                    f"Agent requires interactive terminal authentication. "
                    f"Run the agent's login command directly. "
                    f"Method: {method_name} ({method_id})",
                    auth_methods=auth_methods,
                )

        # If we reach here the auth type is "agent" (self-handled) or unknown.
        # Log and proceed — the agent manages its own auth.
        _log.info(
            "acp agent has auth methods (agent-managed, proceeding)",
            methods=[getattr(m, "id", str(m)) for m in auth_methods],
        )

    async def new_session(self, cwd: str, mcp_servers: list[Any] | None = None) -> str:
        """Create a new session and return the session ID."""
        if not self._started:
            raise RuntimeError("client not started; call start() first")

        conn = self._connection
        assert conn is not None

        kwargs: dict[str, Any] = {"cwd": cwd}
        if mcp_servers is not None:
            kwargs["mcp_servers"] = mcp_servers

        _log.info(
            "acp new_session request",
            cwd=cwd,
            mcp_server_count=len(mcp_servers) if mcp_servers else 0,
            mcp_server_names=[
                s.get("name", "?") if isinstance(s, dict) else "?" for s in (mcp_servers or [])
            ],
        )

        session = await conn.new_session(**kwargs)
        _log.info("acp session created", session_id=session.session_id)
        return session.session_id

    async def load_session(
        self, session_id: str, cwd: str, mcp_servers: list[Any] | None = None
    ) -> str:
        """Resume an existing session by ID.  Returns the session ID.

        Falls back to :meth:`new_session` if the agent rejects the resume.
        """
        if not self._started:
            raise RuntimeError("client not started; call start() first")

        conn = self._connection
        assert conn is not None

        kwargs: dict[str, Any] = {"cwd": cwd, "session_id": session_id}
        if mcp_servers is not None:
            kwargs["mcp_servers"] = mcp_servers

        _log.info("acp load_session request", session_id=session_id, cwd=cwd)

        try:
            await conn.load_session(**kwargs)
            _log.info("acp session resumed", session_id=session_id)
            return session_id
        except Exception:
            _log.warning(
                "acp load_session failed, falling back to new session",
                session_id=session_id,
                exc_info=True,
            )
            return await self.new_session(cwd=cwd, mcp_servers=mcp_servers)

    async def prompt(self, session_id: str, text: str) -> AsyncIterator[ACPUpdate]:
        """Send a prompt and yield streaming updates.

        Yields :class:`ACPTextChunk`, :class:`ACPToolStatus`, and finally
        :class:`ACPPromptDone` when the turn is complete.
        """
        if not self._started:
            raise RuntimeError("client not started; call start() first")

        conn = self._connection
        assert conn is not None

        queue: asyncio.Queue[ACPUpdate | object] = asyncio.Queue()
        self._handler.bind_queue(queue)

        try:
            # Run prompt() in a task so we can yield updates from the queue
            # as they arrive rather than waiting for the full response.
            prompt_task = asyncio.create_task(
                conn.prompt(
                    session_id=session_id,
                    prompt=[text_block(text)],
                )
            )

            # Yield updates from the queue until prompt() finishes.
            while not prompt_task.done():
                try:
                    item = await asyncio.wait_for(queue.get(), timeout=0.05)
                    yield item  # type: ignore[misc]
                except TimeoutError:
                    continue

            # Drain any remaining items in the queue.
            while not queue.empty():
                item = queue.get_nowait()
                yield item  # type: ignore[misc]

            # Get the prompt response and yield the done marker.
            response = prompt_task.result()
            yield ACPPromptDone(stop_reason=response.stop_reason)

        finally:
            self._handler.unbind_queue()

    async def cancel(self, session_id: str) -> None:
        """Send a cancel notification for the given session."""
        if not self._started:
            return

        conn = self._connection
        if conn is None:
            return

        _log.info("acp cancel sent", session_id=session_id)
        await conn.cancel(session_id=session_id)

    async def close(self) -> None:
        """Shut down the agent subprocess gracefully."""
        if not self._started:
            return

        self._started = False
        self._capabilities = None
        self._connection = None
        await self._transport.stop()
        _log.info("acp client closed")

    async def __aenter__(self) -> ACPClient:
        await self.start()
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.close()
