"""LLM client classes and factory for kokoro-chat.

Extracted from chat.py to keep the ChatEngine class focused on orchestration.
Provides:
  - Channel adapters for streaming / no-op modes
  - Usage-tracking AI client wrapper
  - Streaming motosan-chat client adapter
  - Factory functions for building LLM clients from ChatConfig
"""
from __future__ import annotations

import asyncio
import logging
from typing import Any, AsyncIterator

from motosan_ai import Client as AiClient
from motosan_chat import StreamChunk

from kokoro_chat.config import ChatConfig

logger = logging.getLogger(__name__)

# Sentinel value that signals the end of a streaming queue
STREAM_SENTINEL = object()


# ── Channel adapters ─────────────────────────────────────────────────────────

class NullChannel:
    """No-op channel for non-streaming API mode."""

    @property
    def name(self) -> str:
        return "api"

    async def listen(self) -> AsyncIterator:
        return
        yield  # pragma: no cover

    async def stream_chunk(self, user_id: str, chunk: str) -> None:
        pass

    async def stream_done(self, user_id: str) -> None:
        pass

    async def send(self, user_id: str, text: str) -> None:
        pass


class StreamingChannel:
    """Channel that pushes streaming chunks into an asyncio.Queue."""

    def __init__(self, queue: asyncio.Queue) -> None:
        self._queue = queue

    @property
    def name(self) -> str:
        return "api"

    async def listen(self) -> AsyncIterator:
        return
        yield  # pragma: no cover

    async def stream_chunk(self, user_id: str, chunk: str) -> None:
        if chunk.startswith("\U0001f527"):  # filter tool-name annotations
            return
        await self._queue.put(chunk)

    async def stream_done(self, user_id: str) -> None:
        pass

    async def send(self, user_id: str, text: str) -> None:
        pass


# ── LLM client wrappers ──────────────────────────────────────────────────────

class UsageTrackingAiClient:
    """Wrap motosan-ai client and accumulate usage."""

    def __init__(self, client: AiClient, model: str, max_tokens: int) -> None:
        self._client = client
        self._model = model
        self._max_tokens = max_tokens
        self.total_input_tokens = 0
        self.total_output_tokens = 0

    async def chat(self, messages: list[Any], **kwargs: Any) -> Any:
        provider_options = dict(kwargs.pop("provider_options", {}) or {})
        provider_options.setdefault("model", self._model)
        response = await self._client.chat(
            messages,
            max_tokens=self._max_tokens,
            provider_options=provider_options,
            **kwargs,
        )
        usage = getattr(response, "usage", None)
        if usage:
            self.total_input_tokens += int(getattr(usage, "input_tokens", 0))
            self.total_output_tokens += int(getattr(usage, "output_tokens", 0))
        return response

    async def stream(self, messages: list[Any], **kwargs: Any) -> AsyncIterator[Any]:
        provider_options = dict(kwargs.pop("provider_options", {}) or {})
        provider_options.setdefault("model", self._model)
        async for event in self._client.stream(
            messages,
            max_tokens=self._max_tokens,
            provider_options=provider_options,
            **kwargs,
        ):
            yield event

    async def stream_with_tools(
        self, messages: list[Any], **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        provider_options = dict(kwargs.pop("provider_options", {}) or {})
        provider_options.setdefault("model", self._model)
        async for event in self._client.stream(
            messages,
            max_tokens=self._max_tokens,
            provider_options=provider_options,
            **kwargs,
        ):
            etype = getattr(event, "event_type", "text")
            if etype == "text":
                yield StreamChunk(kind="text", content=getattr(event, "content", ""))
            elif etype == "tool_call_start":
                yield StreamChunk(
                    kind="tool_call_start",
                    tool_id=getattr(event, "tool_call_id", "") or "",
                    tool_name=getattr(event, "tool_call_name", "") or "",
                )
            elif etype == "tool_call_args":
                yield StreamChunk(
                    kind="tool_call_args",
                    tool_id=getattr(event, "tool_call_id", "") or "",
                    content=getattr(event, "tool_call_args_delta", "") or "",
                )
            elif etype == "tool_call_end":
                yield StreamChunk(
                    kind="tool_call_end",
                    tool_id=getattr(event, "tool_call_id", "") or "",
                )


class StreamingMotosanAiClient:
    """Custom LLM client with usage tracking and prompt-cache support."""

    def __init__(
        self,
        tracking_client: UsageTrackingAiClient,
        system: str | None = None,
        system_prompt_blocks: list[dict] | None = None,
        prompt_cache_enabled: bool = False,
    ) -> None:
        self._tracking = tracking_client
        self._system = system
        self._system_prompt_blocks = system_prompt_blocks
        self._prompt_cache_enabled = prompt_cache_enabled

    def _resolve_system(self) -> str | list[dict] | None:
        if self._prompt_cache_enabled and self._system_prompt_blocks:
            return self._system_prompt_blocks
        return self._system

    async def chat(self, messages: list, tools: list) -> Any:
        from motosan_chat.ai import _to_ai_message, _to_ai_tool
        ai_messages = [_to_ai_message(m) for m in messages]
        ai_tools = [_to_ai_tool(t) for t in tools]
        response = await self._tracking.chat(
            ai_messages,
            tools=ai_tools,
            system=self._resolve_system(),
        )
        from motosan_chat import LlmResponse
        from motosan_chat.agent import ToolCallItem
        tool_calls = getattr(response, "tool_calls", [])
        if len(tool_calls) == 1:
            tc = tool_calls[0]
            return LlmResponse.tool_call(
                id=tc.id, name=tc.name, args=getattr(tc, "input", {}),
            )
        elif len(tool_calls) > 1:
            return LlmResponse.tool_calls([
                ToolCallItem(id=tc.id, name=tc.name, args=getattr(tc, "input", {}))
                for tc in tool_calls
            ])
        return LlmResponse.message(getattr(response, "content", ""))

    async def stream(self, messages: list) -> AsyncIterator[str]:
        from motosan_chat.ai import _to_ai_message
        ai_messages = [_to_ai_message(m) for m in messages]
        async for event in self._tracking.stream(
            ai_messages, system=self._resolve_system(),
        ):
            yield getattr(event, "content", "")

    async def stream_with_tools(
        self, messages: list, tools: list,
    ) -> AsyncIterator[StreamChunk]:
        from motosan_chat.ai import _to_ai_message, _to_ai_tool
        ai_messages = [_to_ai_message(m) for m in messages]
        ai_tools = [_to_ai_tool(t) for t in tools]
        async for chunk in self._tracking.stream_with_tools(
            ai_messages,
            tools=ai_tools,
            system=self._resolve_system(),
        ):
            yield chunk


# ── Factory functions ────────────────────────────────────────────────────────

def build_ai_client(
    provider: str, model: str, config: ChatConfig,
) -> AiClient | None:
    """Build a motosan-ai AiClient for the given provider/model.

    Returns None when credentials are missing or the provider is unknown.
    """
    provider = provider.lower()
    if provider == "anthropic":
        api_key = config.anthropic_api_key or config.anthropic_oauth_token
        if not api_key:
            logger.warning("No Anthropic credential for provider '%s'", provider)
            return None
        return AiClient.anthropic(api_key=api_key, model=model)
    if provider == "openai":
        if not config.openai_api_key:
            logger.warning("No OpenAI key for provider '%s'", provider)
            return None
        return AiClient.openai(api_key=config.openai_api_key, model=model)
    if provider == "minimax":
        if not config.minimax_api_key:
            logger.warning("No MiniMax key for provider '%s'", provider)
            return None
        return AiClient.minimax(api_key=config.minimax_api_key, model=model)
    logger.warning("Unknown provider '%s'", provider)
    return None
