"""LLM client with multi-provider streaming support.

Provides a unified async streaming interface for OpenAI, Anthropic,
Ollama, and vLLM providers.  Used by ``Action.call_llm()`` to stream
text chunks from the configured LLM.
"""

from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional

if TYPE_CHECKING:
    from oedix.core.config import AgentConfig

logger = logging.getLogger("oedix.llm")


class LLMClient:
    """Unified LLM client with async streaming across providers.

    Usage::

        client = LLMClient(config)
        async for chunk in client.stream(messages):
            print(chunk, end="")
    """

    def __init__(self, config: AgentConfig) -> None:
        self.config = config
        self._openai_client: Any = None
        self._anthropic_client: Any = None

    async def stream(
        self,
        messages: list[dict[str, str]],
        *,
        provider: Optional[str] = None,
        model: Optional[str] = None,
        **kwargs: Any,
    ) -> AsyncGenerator[str, None]:
        """Stream text chunks from the LLM.

        Args:
            messages: Chat messages in OpenAI format
                (``[{"role": "system"|"user"|"assistant", "content": "..."}]``).
            provider: Override the default provider key (e.g. "openai").
            model: Override the default model key (e.g. "gpt-5.4").
            **kwargs: Extra parameters forwarded to the provider API.

        Yields:
            Text chunks from the LLM response.
        """
        prov_name = provider or self.config.default_provider
        prov_config = self.config.get_provider(prov_name)

        # Resolve model_id
        model_key = model or self.config.default_model
        extra: dict[str, Any] = {}
        try:
            model_config = self.config.get_model(model_key, prov_name)
            model_id = model_config.model_id
            max_tokens = kwargs.pop(
                "max_tokens", model_config.max_output_tokens
            )
            extra = model_config.extra
        except KeyError:
            # Model not in config — use key directly as model_id
            model_id = model_key
            max_tokens = kwargs.pop("max_tokens", 8192)

        api_key = prov_config.api_key
        base_url = prov_config.base_url

        if prov_name in ("openai", "vllm"):
            use_max_completion = (
                extra.get("use_max_completion_tokens", False)
                or prov_name == "openai"
            )
            async for chunk in self._stream_openai(
                messages, model_id, api_key, base_url, max_tokens,
                use_max_completion_tokens=use_max_completion, **kwargs
            ):
                yield chunk
        elif prov_name == "anthropic":
            async for chunk in self._stream_anthropic(
                messages, model_id, api_key, max_tokens, **kwargs
            ):
                yield chunk
        elif prov_name == "ollama":
            async for chunk in self._stream_ollama(
                messages, model_id, base_url, **kwargs
            ):
                yield chunk
        else:
            raise ValueError(f"Unsupported LLM provider: {prov_name}")

    # ------------------------------------------------------------------
    # Provider-specific streaming
    # ------------------------------------------------------------------

    async def _stream_openai(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str,
        max_tokens: int,
        *,
        use_max_completion_tokens: bool = False,
        **kwargs: Any,
    ) -> AsyncGenerator[str, None]:
        """Stream from OpenAI-compatible APIs (OpenAI, vLLM)."""
        try:
            from openai import AsyncOpenAI
        except ImportError:
            raise ImportError(
                "openai package required: pip install lib[openai]"
            )

        client = self._get_openai_client(api_key, base_url)
        system_msg, chat_messages = self._split_system(messages)

        # Some models (e.g. gpt-5.x) require max_completion_tokens
        token_param = (
            "max_completion_tokens" if use_max_completion_tokens
            else "max_tokens"
        )
        params: dict[str, Any] = {
            "model": model,
            "messages": (
                [{"role": "system", "content": system_msg}]
                if system_msg
                else []
            )
            + chat_messages,
            "stream": True,
            token_param: max_tokens,
            **kwargs,
        }

        response = await client.chat.completions.create(**params)
        async for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    async def _stream_anthropic(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        max_tokens: int,
        **kwargs: Any,
    ) -> AsyncGenerator[str, None]:
        """Stream from Anthropic Claude API."""
        try:
            from anthropic import AsyncAnthropic
        except ImportError:
            raise ImportError(
                "anthropic package required: pip install lib[anthropic]"
            )

        client = self._get_anthropic_client(api_key)
        system_msg, chat_messages = self._split_system(messages)

        params: dict[str, Any] = {
            "model": model,
            "messages": chat_messages,
            "stream": True,
            "max_tokens": max_tokens,
            **kwargs,
        }
        if system_msg:
            params["system"] = system_msg

        async with client.messages.stream(**params) as stream:
            async for text in stream.text_stream:
                yield text

    async def _stream_ollama(
        self,
        messages: list[dict],
        model: str,
        base_url: str,
        **kwargs: Any,
    ) -> AsyncGenerator[str, None]:
        """Stream from local Ollama server."""
        url_base = base_url or "http://localhost:11434"
        url = f"{url_base}/api/chat"
        payload = {
            "model": model,
            "messages": messages,
            "stream": True,
            **kwargs,
        }

        try:
            import aiohttp

            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as resp:
                    async for line in resp.content:
                        line_str = line.decode("utf-8").strip()
                        if not line_str:
                            continue
                        try:
                            data = json.loads(line_str)
                            content = data.get("message", {}).get(
                                "content", ""
                            )
                            if content:
                                yield content
                            if data.get("done"):
                                break
                        except json.JSONDecodeError:
                            continue
        except ImportError:
            # Fallback: sync requests in thread
            import asyncio

            chunks = await asyncio.to_thread(
                self._ollama_sync, url, payload
            )
            for chunk in chunks:
                yield chunk

    @staticmethod
    def _ollama_sync(url: str, payload: dict) -> list[str]:
        """Synchronous Ollama streaming (fallback when aiohttp unavailable)."""
        import requests

        chunks: list[str] = []
        resp = requests.post(url, json=payload, stream=True)
        for line in resp.iter_lines():
            if not line:
                continue
            try:
                data = json.loads(line)
                content = data.get("message", {}).get("content", "")
                if content:
                    chunks.append(content)
                if data.get("done"):
                    break
            except json.JSONDecodeError:
                continue
        return chunks

    # ------------------------------------------------------------------
    # Client management
    # ------------------------------------------------------------------

    def _get_openai_client(self, api_key: str, base_url: str) -> Any:
        """Get or create cached OpenAI async client."""
        cache_key = (api_key, base_url)
        if self._openai_client is None or getattr(
            self, "_openai_cache_key", None
        ) != cache_key:
            from openai import AsyncOpenAI

            params: dict[str, Any] = {
                "api_key": api_key or "EMPTY"
            }
            if base_url:
                params["base_url"] = base_url
            self._openai_client = AsyncOpenAI(**params)
            self._openai_cache_key = cache_key
        return self._openai_client

    def _get_anthropic_client(self, api_key: str) -> Any:
        """Get or create cached Anthropic async client."""
        if self._anthropic_client is None or getattr(
            self, "_anthropic_cache_key", None
        ) != api_key:
            from anthropic import AsyncAnthropic

            self._anthropic_client = AsyncAnthropic(api_key=api_key)
            self._anthropic_cache_key = api_key
        return self._anthropic_client

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _split_system(
        messages: list[dict],
    ) -> tuple[str, list[dict]]:
        """Split system message from chat messages.

        Returns:
            (system_content, remaining_messages)
        """
        system = ""
        chat: list[dict] = []
        for msg in messages:
            if msg.get("role") == "system":
                system = (
                    system + "\n" + msg["content"]
                    if system
                    else msg["content"]
                )
            else:
                chat.append(msg)
        return system, chat
