"""ACP stream provider — makes ACP agents look like any other LLM provider.

Implements the :class:`StreamProvider` protocol so that ``agent.py`` can drive
an ACP agent exactly the same way it drives LiteLLM or OAuth providers.

Key design: from agent.py's perspective, the ACP model returns text only (no
tool calls).  The ACP agent executes tools internally; tool status updates are
surfaced via an optional callback for UI display, not for Otto to execute.
"""

from __future__ import annotations

from collections.abc import AsyncIterator, Callable
from types import SimpleNamespace
from typing import Any

from otto.acp.client import ACPClient, ACPPromptDone, ACPTextChunk, ACPToolStatus
from otto.log import get_logger

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


class ACPStreamProvider:
    """StreamProvider shim that delegates to an ACP agent subprocess.

    Parameters
    ----------
    agent_command:
        Command list to spawn the ACP agent (e.g. ``["claude"]``).
    agent_env:
        Optional environment variables for the agent subprocess.
    agent_cwd:
        Working directory for the ACP agent and its session.
    on_tool_status:
        Optional callback ``(name, status, call_id)`` invoked for each
        :class:`ACPToolStatus` update from the agent.  Informational only.
    """

    def __init__(
        self,
        agent_command: list[str],
        agent_env: dict[str, str] | None = None,
        agent_cwd: str | None = None,
        on_tool_status: Callable[[str, str, str], None] | None = None,
        mcp_servers: list[dict] | None = None,
        system_prompt: str | None = None,
        resume_session_id: str | None = None,
    ) -> None:
        self._agent_command = agent_command
        self._agent_env = agent_env
        self._agent_cwd = agent_cwd or "/tmp"
        self.on_tool_status = on_tool_status
        self._mcp_servers = mcp_servers
        self._system_prompt = system_prompt
        self._system_injected = False

        self._client: ACPClient | None = None
        self._session_id: str | None = None
        self._resume_session_id = resume_session_id

    @property
    def session_id(self) -> str | None:
        """The current ACP session ID, if any."""
        return self._session_id

    async def _ensure_client(self) -> ACPClient:
        if self._client is not None and self._client.started:
            return self._client

        _log.info(
            "acp provider starting client",
            command=self._agent_command,
        )
        client = ACPClient(
            command=self._agent_command,
            env=self._agent_env,
            cwd=self._agent_cwd,
        )
        await client.start()
        self._client = client
        return client

    async def _ensure_session(self, client: ACPClient) -> str:
        if self._session_id is not None:
            return self._session_id

        # Try resuming a previous session if we have a persisted ID.
        if self._resume_session_id:
            _log.info(
                "acp provider resuming session",
                resume_id=self._resume_session_id,
            )
            session_id = await client.load_session(
                session_id=self._resume_session_id,
                cwd=self._agent_cwd,
                mcp_servers=self._mcp_servers,
            )
            self._session_id = session_id
            # If we resumed successfully, the system prompt was already
            # injected in the original session.
            if session_id == self._resume_session_id:
                self._system_injected = True
                _log.info("acp provider session resumed", session_id=session_id)
            else:
                _log.info(
                    "acp provider resume failed, new session created",
                    original=self._resume_session_id,
                    new=session_id,
                )
            self._resume_session_id = None
            return session_id

        session_id = await client.new_session(
            cwd=self._agent_cwd,
            mcp_servers=self._mcp_servers,
        )
        self._session_id = session_id
        _log.info("acp provider session created", session_id=session_id)
        return session_id

    async def stream_completion(
        self,
        model: str,
        messages: list[dict],
        *,
        tools: list[dict] | None = None,
        tool_choice: str | dict | None = None,
        max_tokens: int | None = None,
        reasoning_effort: str | None = None,
    ) -> AsyncIterator[Any]:
        client = await self._ensure_client()
        session_id = await self._ensure_session(client)

        # ACP agents manage their own context.  We send only the last user
        # message text rather than the full conversation history.
        user_text = _extract_last_user_text(messages)

        # Inject the system prompt as a preamble on the first message of the session.
        # ACP has no native system_message field, so we prepend it to the user text.
        if self._system_prompt and not self._system_injected:
            user_text = (
                f"[INSTRUCTIONS — follow these for all responses in this session]\n"
                f"{self._system_prompt}\n"
                f"[END INSTRUCTIONS]\n\n"
                f"{user_text}"
            )
            self._system_injected = True

        _log.info(
            "acp provider prompting",
            session_id=session_id,
            text_length=len(user_text),
        )

        return _ACPChunkIterator(
            client=client,
            session_id=session_id,
            user_text=user_text,
            on_tool_status=self.on_tool_status,
        )

    async def close(self) -> None:
        if self._client is not None:
            _log.info("acp provider closing client")
            await self._client.close()
            self._client = None
            self._session_id = None
            self._system_injected = False


def _extract_last_user_text(messages: list[dict]) -> str:
    for msg in reversed(messages):
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return content
            # Handle list-of-blocks format (e.g. vision messages).
            if isinstance(content, list):
                parts = []
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "text":
                        parts.append(block.get("text", ""))
                    elif isinstance(block, str):
                        parts.append(block)
                return "\n".join(parts)
    return ""


class _ACPChunkIterator:
    """Async iterator that translates ACP updates into litellm-compatible chunks.

    Yields ``SimpleNamespace`` objects matching the shape that
    ``agent.py._consume_stream`` expects::

        chunk.choices[0].delta.content   -- text content
        chunk.choices[0].finish_reason   -- "stop" on final chunk
        chunk.usage                      -- None (ACP doesn't report token usage)
    """

    def __init__(
        self,
        client: ACPClient,
        session_id: str,
        user_text: str,
        on_tool_status: Callable[[str, str, str], None] | None = None,
    ) -> None:
        self._client = client
        self._session_id = session_id
        self._user_text = user_text
        self._on_tool_status = on_tool_status
        self._inner: AsyncIterator | None = None

    def __aiter__(self) -> _ACPChunkIterator:
        return self

    async def __anext__(self) -> Any:
        if self._inner is None:
            self._inner = self._stream()
        return await self._inner.__anext__()

    async def _stream(self) -> AsyncIterator[Any]:
        async for update in self._client.prompt(self._session_id, self._user_text):
            if isinstance(update, ACPTextChunk):
                yield _make_chunk(content=update.text)
            elif isinstance(update, ACPToolStatus):
                _log.debug(
                    "acp tool status",
                    name=update.name,
                    status=update.status,
                    call_id=update.call_id,
                )
                if self._on_tool_status is not None:
                    self._on_tool_status(update.name, update.status, update.call_id)
            elif isinstance(update, ACPPromptDone):
                yield _make_chunk(finish_reason=_map_stop_reason(update.stop_reason))


def _make_chunk(
    content: str | None = None,
    finish_reason: str | None = None,
) -> SimpleNamespace:
    delta = SimpleNamespace(content=content, tool_calls=None)
    choice = SimpleNamespace(delta=delta, finish_reason=finish_reason)
    return SimpleNamespace(choices=[choice], usage=None)


def _map_stop_reason(acp_reason: str) -> str:
    mapping = {
        "end_turn": "stop",
        "cancelled": "stop",
        "max_tokens": "max_tokens",
    }
    return mapping.get(acp_reason, "stop")
