"""LLM Gateway - 统一 LLM 调用入口"""

import logging
import time

from agentkit.core.exceptions import LLMProviderError, ModelNotFoundError
from agentkit.llm.config import LLMConfig
from agentkit.llm.protocol import LLMProvider, LLMRequest, LLMResponse, StreamChunk, TokenUsage
from agentkit.llm.providers.tracker import UsageSummary, UsageTracker
from agentkit.telemetry.tracing import get_tracer, _OTEL_AVAILABLE
from agentkit.telemetry.metrics import llm_token_histogram

logger = logging.getLogger(__name__)


class LLMGateway:
    """LLM 网关 - Provider 注册、模型别名解析、Fallback、Usage 追踪"""

    def __init__(self, config: LLMConfig | None = None):
        self._providers: dict[str, LLMProvider] = {}
        self._usage_tracker = UsageTracker()
        self._config = config or LLMConfig()

    def register_provider(self, name: str, provider: LLMProvider) -> None:
        """注册 Provider"""
        self._providers[name] = provider
        logger.info(f"LLM provider '{name}' registered")

    @property
    def has_providers(self) -> bool:
        """Return True if at least one LLM provider is registered."""
        return bool(self._providers)

    async def chat(
        self,
        messages: list[dict[str, str]],
        model: str,
        agent_name: str = "",
        task_type: str = "",
        tools: list[dict] | None = None,
        tool_choice: str = "auto",
        **kwargs,
    ) -> LLMResponse:
        """发送 chat 请求，自动解析别名和 Fallback"""
        resolved_model = self._resolve_model_alias(model)

        if not self._providers:
            raise LLMProviderError("", "No provider registered")

        # Telemetry: start LLM span
        _span_cm = None
        _span = None
        if _OTEL_AVAILABLE:
            tracer = get_tracer()
            if tracer is not None:
                from opentelemetry.trace import SpanKind
                _span_cm = tracer.start_as_current_span(
                    "gen_ai.chat",
                    kind=SpanKind.CLIENT,
                    attributes={
                        "gen_ai.system": resolved_model.split("/")[0] if "/" in resolved_model else "unknown",
                        "gen_ai.operation.name": "chat",
                        "gen_ai.request.model": resolved_model,
                    },
                )
                _span = _span_cm.__enter__()

        start = time.monotonic()
        models_to_try = self._get_models_to_try(resolved_model)
        last_error: LLMProviderError | None = None

        try:
            for model_name in models_to_try:
                try:
                    provider, actual_model = self._resolve_model(model_name)
                except ModelNotFoundError:
                    continue

                req = LLMRequest(
                    messages=messages,
                    model=actual_model,
                    tools=tools,
                    tool_choice=tool_choice,
                    **kwargs,
                )
                try:
                    response = await provider.chat(req)
                    break
                except LLMProviderError as e:
                    last_error = e
                    logger.warning(f"Model '{model_name}' failed, trying next: {e}")
                    continue
            else:
                raise last_error or LLMProviderError("", f"All models failed for '{resolved_model}'")

            latency_ms = (time.monotonic() - start) * 1000

            # 计算成本
            cost = self._calculate_cost(response.model, response.usage)

            # 记录使用量
            self._usage_tracker.record(
                agent_name=agent_name,
                model=response.model,
                usage=response.usage,
                cost=cost,
                latency_ms=latency_ms,
            )

            # Telemetry: record token usage and end span
            if _span is not None:
                _span.set_attribute("gen_ai.usage.input_tokens", response.usage.prompt_tokens)
                _span.set_attribute("gen_ai.usage.output_tokens", response.usage.completion_tokens)
                _span.set_attribute("gen_ai.response.model", response.model)
                _span.set_attribute("gen_ai.duration_ms", int(latency_ms))
            llm_token_histogram().record(
                response.usage.total_tokens,
                {"gen_ai.request.model": resolved_model},
            )

            return response
        finally:
            if _span_cm is not None:
                _span_cm.__exit__(None, None, None)

    async def chat_stream(
        self,
        messages: list[dict[str, str]],
        model: str,
        agent_name: str = "",
        task_type: str = "",
        tools: list[dict] | None = None,
        tool_choice: str = "auto",
        **kwargs,
    ):
        """Stream chat response with fallback support.

        If the primary model fails before any chunk is yielded, tries fallback
        models. If it fails after chunks have been sent, yields an error chunk
        and terminates (cannot switch mid-stream).
        """
        resolved_model = self._resolve_model_alias(model)

        if not self._providers:
            raise LLMProviderError("", "No provider registered")

        models_to_try = self._get_models_to_try(resolved_model)
        last_error: Exception | None = None

        for model_name in models_to_try:
            try:
                provider, actual_model = self._resolve_model(model_name)
            except ModelNotFoundError:
                continue

            stream_request = LLMRequest(
                messages=messages,
                model=actual_model,
                tools=tools,
                tool_choice=tool_choice,
                **kwargs,
            )

            chunk_yielded = False
            start = time.monotonic()
            total_content = ""
            final_usage = None
            final_model = model_name

            try:
                async for chunk in provider.chat_stream(stream_request):
                    chunk_yielded = True
                    if chunk.content:
                        total_content += chunk.content
                    if chunk.usage:
                        final_usage = chunk.usage
                    if chunk.model:
                        final_model = chunk.model
                    yield chunk

                # Track usage after successful stream
                latency_ms = (time.monotonic() - start) * 1000
                if final_usage is None:
                    final_usage = TokenUsage()
                cost = self._calculate_cost(final_model, final_usage)
                self._usage_tracker.record(
                    agent_name=agent_name,
                    model=final_model,
                    usage=final_usage,
                    cost=cost,
                    latency_ms=latency_ms,
                )
                return  # Success, done
            except Exception as e:
                last_error = e
                if chunk_yielded:
                    # Can't switch mid-stream, terminate gracefully
                    logger.error(f"Stream failed after chunks sent for '{model_name}': {e}")
                    yield StreamChunk(
                        content="",
                        model=final_model,
                        usage=None,
                        is_final=True,
                    )
                    return
                # No chunks yet, try next fallback
                logger.warning(f"Stream failed for '{model_name}', trying fallback: {e}")
                continue

        # All models failed
        raise last_error or LLMProviderError("", f"No provider available for streaming '{resolved_model}'")

    def _get_models_to_try(self, resolved_model: str) -> list[str]:
        """Return [primary_model] + fallback_models for the given resolved model."""
        fallback_models = self._config.fallbacks.get(resolved_model, [])
        return [resolved_model] + fallback_models

    def _resolve_model_alias(self, model: str) -> str:
        """解析模型别名"""
        if model in self._config.model_aliases:
            return self._config.model_aliases[model]
        return model

    def _resolve_model(self, model: str) -> tuple[LLMProvider, str]:
        """解析模型为 (provider, actual_model_name)"""
        # model 格式: "provider/model_name" 或 "model_name"
        if "/" in model:
            provider_name, model_name = model.split("/", 1)
            if provider_name not in self._providers:
                raise ModelNotFoundError(model)
            return self._providers[provider_name], model_name

        # 无 "/" 前缀：仅当只有一个 provider 时自动匹配
        if len(self._providers) == 1:
            provider = next(iter(self._providers.values()))
            return provider, model

        raise ModelNotFoundError(model)

    def _get_fallback_model(self, model: str) -> str | None:
        """获取 Fallback 模型"""
        fallbacks = self._config.fallbacks.get(model, [])
        return fallbacks[0] if fallbacks else None

    def _calculate_cost(self, model: str, usage: TokenUsage) -> float:
        """计算成本"""
        # 在 provider config 的 models 中查找成本配置
        for provider_config in self._config.providers.values():
            if model in provider_config.models:
                model_conf = provider_config.models[model]
                input_cost = usage.prompt_tokens * model_conf.get("cost_per_1k_input", 0) / 1000
                output_cost = usage.completion_tokens * model_conf.get("cost_per_1k_output", 0) / 1000
                return input_cost + output_cost
        return 0.0

    def get_usage(
        self,
        agent_name: str | None = None,
        start_time=None,
        end_time=None,
    ) -> UsageSummary:
        """查询使用量"""
        return self._usage_tracker.get_usage(
            agent_name=agent_name,
            start_time=start_time,
            end_time=end_time,
        )
