"""Unified async LLM client dispatching to OpenAI or Anthropic."""

from __future__ import annotations

from typing import TYPE_CHECKING

import structlog

from koan.llm.config import LLMConfig, ProviderSettings

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from koan.llm.providers.anthropic import AnthropicProvider
    from koan.llm.providers.openai import OpenAIProvider
    from koan.types import LLMResponse, Message, StreamDelta

log = structlog.get_logger()


class LLMClient:
    """Unified async client over OpenAI and Anthropic.

    Lazily initializes provider clients on first use. Shares a single
    provider instance per API key for connection pooling.

    Usage::

        client = LLMClient()
        config = LLMConfig(provider="openai", model="gpt-4o-mini")
        response = await client.complete(config, [Message(role="user", content="Hello")])
        print(response.text)
    """

    def __init__(self, settings: ProviderSettings | None = None) -> None:
        self._settings = settings or ProviderSettings()
        self._openai: OpenAIProvider | None = None
        self._anthropic: AnthropicProvider | None = None

    async def complete(self, config: LLMConfig, messages: list[Message]) -> LLMResponse:
        """Send messages to the configured provider and return a normalized response."""
        provider = self._get_provider(config)

        log.debug(
            "llm.complete",
            provider=config.provider,
            model=config.model,
            message_count=len(messages),
        )

        response = await provider.complete(config, messages)

        log.info(
            "llm.complete.done",
            provider=config.provider,
            model=response.model,
            input_tokens=response.usage.input_tokens,
            output_tokens=response.usage.output_tokens,
            stop_reason=response.stop_reason,
        )

        return response

    async def stream(self, config: LLMConfig, messages: list[Message]) -> AsyncIterator[StreamDelta]:
        """Stream responses from the configured provider as deltas."""
        provider = self._get_provider(config)

        log.debug(
            "llm.stream",
            provider=config.provider,
            model=config.model,
            message_count=len(messages),
        )

        async for delta in provider.stream(config, messages):
            yield delta

    def _get_provider(self, config: LLMConfig) -> OpenAIProvider | AnthropicProvider:
        """Return (and lazily create) the provider for the given config."""
        if config.provider == "openai":
            if self._openai is None:
                from koan.llm.providers.openai import OpenAIProvider

                key = self._settings.get_api_key("openai")
                self._openai = OpenAIProvider(api_key=key)
            return self._openai

        if self._anthropic is None:
            from koan.llm.providers.anthropic import AnthropicProvider

            key = self._settings.get_api_key("anthropic")
            self._anthropic = AnthropicProvider(api_key=key)
        return self._anthropic

    async def close(self) -> None:
        """Close underlying HTTP clients."""
        if self._openai:
            await self._openai.close()
        if self._anthropic:
            await self._anthropic.close()

    async def __aenter__(self) -> LLMClient:
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.close()
