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

import asyncio
import logging
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Protocol

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

if TYPE_CHECKING:
    from agentkit.llm.cache import LitellmCacheManager

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# TYPE_CHECKING Protocols — 避免 Any，描述运行时 lazy import 的对象
# ---------------------------------------------------------------------------


if TYPE_CHECKING:

    class _QuotaServiceLike(Protocol):
        """Quota service 最小契约（仅覆盖 gateway._enforce_quota 用到的方法）。"""

        async def is_model_allowed(
            self, db: Path, department_id: str, model: str
        ) -> tuple[bool, str]: ...

        async def check_quota(
            self,
            db: Path,
            department_id: str,
            quota_type: str,
            period: str,
            current: float,
        ) -> tuple[bool, str]: ...

        async def get_quota(
            self, db: Path, department_id: str, quota_type: str, period: str
        ) -> dict[str, object] | None: ...


class QuotaExceededError(Exception):
    """Raised when a department's LLM quota is exceeded.

    Carries enough metadata for the API layer to return a structured
    429 response (department_id, quota_type, period, limit, current).
    """

    def __init__(
        self,
        department_id: str,
        quota_type: str,
        period: str,
        limit: object,
        current: object,
    ) -> None:
        self.department_id = department_id
        self.quota_type = quota_type
        self.period = period
        self.limit = limit
        self.current = current
        super().__init__(
            f"Quota exceeded for department {department_id}: "
            f"{quota_type} {period} (limit={limit}, current={current})"
        )


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

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

        # Cache (U17 — LiteLLM 缓存管理器，opt-in，默认禁用)
        self._cache_manager: "LitellmCacheManager | None" = None
        if self._config.cache and self._config.cache.enabled:
            from agentkit.llm.cache import LitellmCacheConfig, LitellmCacheManager

            litellm_config = LitellmCacheConfig.from_cache_config(self._config.cache)
            self._cache_manager = LitellmCacheManager(litellm_config)
            self._cache_manager.enable()
            logger.info(
                f"LLM cache enabled (LiteLLM, backend={self._config.cache.backend}, "
                f"similarity_threshold={litellm_config.similarity_threshold})"
            )

    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",
        timeout: float | None = None,
        user_id: str | None = None,
        department_ids: list[str] | None = None,
        db_path: Path | str | None = None,
        kb_id: str | None = None,
        kb_acl_hash: str | None = None,
        **kwargs,
    ) -> LLMResponse:
        """发送 chat 请求，自动解析别名和 Fallback"""
        resolved_model = self._resolve_model_alias(model)

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

        # ── Quota enforcement ──
        # Only enforce when department_ids + db_path are provided
        # (other call sites pass None — no quota check).
        if department_ids and db_path:
            await self._enforce_quota(db_path, department_ids, resolved_model)

        # 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()

        # ── Cache check (U17 — LiteLLM cache via cache_key in request) ──
        # LiteLLM 在 litellm.acompletion 内部处理缓存读写，gateway 只需：
        # 1. 构建 per-user + ACL-scoped cache_key（安全要求 a, b）
        # 2. 将 cache 参数注入 kwargs 透传到 provider
        # 3. 检测响应的 cache_hit 标志，用于 usage tracking（cost=0）
        if self._cache_manager is not None:
            from agentkit.llm.cache import LitellmCacheManager

            # 解析 KB caching_disabled（安全要求 c）
            # 非 RAG 请求（kb_id=None）→ 默认启用缓存（无 KB 数据需保护）。
            # RAG 请求（kb_id!=None）→ fail-closed：默认禁用缓存，仅在 settings
            # 明确返回 caching_disabled=False 时启用。防止 DB 异常时 fail-open
            # 导致禁用缓存的 KB 数据泄漏到缓存。
            kb_caching_disabled = kb_id is not None
            if kb_id is not None:
                try:
                    from agentkit.rag_platform.settings import get_settings_store

                    settings = await get_settings_store().get_settings(kb_id)
                    if settings is not None:
                        kb_caching_disabled = settings.caching_disabled
                    # settings 为 None（KB 不存在）→ 保持 True（fail-closed）
                except Exception as e:
                    logger.warning(f"Failed to read KB cache settings for kb_id={kb_id}: {e}")
                    # 读取异常 → 保持 True（fail-closed，禁用缓存）

            if self._cache_manager.should_cache(kb_caching_disabled, user_id):
                cache_key = self._cache_manager.build_cache_key(
                    model=resolved_model,
                    messages=messages,
                    temperature=kwargs.get("temperature", 0.7),
                    tools=tools,
                    tool_choice=tool_choice,
                    max_tokens=kwargs.get("max_tokens", 2000),
                    user_id=user_id,
                    kb_acl_hash=kb_acl_hash,
                )
                kwargs["cache"] = LitellmCacheManager.cache_params_for_hit(cache_key)
            else:
                kwargs["cache"] = LitellmCacheManager.cache_params_for_no_cache()

        # ── Normal provider call ──
        models_to_try = self._get_models_to_try(resolved_model)
        last_error: LLMProviderError | None = None
        response: LLMResponse | 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,
                    timeout=timeout,
                    **kwargs,
                )
                try:
                    response = await provider.chat(req)
                    # Empty response detection: if content is None/empty and no tool_calls,
                    # treat as failure and try next fallback model.
                    # This handles the common case where providers return 200 OK but empty body.
                    if (
                        response.content is None or not response.content.strip()
                    ) and not response.tool_calls:
                        # Record usage for billing before discarding this response
                        if response.usage:
                            latency_ms = (time.monotonic() - start) * 1000
                            cost = self._calculate_cost(model_name, response.usage)
                            await self._record_usage(
                                agent_name=agent_name,
                                model=model_name,
                                usage=response.usage,
                                cost=cost,
                                latency_ms=latency_ms,
                                user_id=user_id,
                                department_ids=department_ids,
                            )
                        logger.warning(
                            f"Model '{model_name}' returned empty content with no tool_calls, "
                            f"trying next fallback"
                        )
                        last_error = LLMProviderError(
                            model_name,
                            f"Empty response from {model_name} (no content, no tool_calls)",
                        )
                        continue
                    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

            # U17 — 检测 LiteLLM 缓存命中（用于 usage tracking cost=0）
            is_cache_hit = getattr(response, "cache_hit", False)
            if self._cache_manager is not None:
                self._cache_manager.record_cache_result(is_cache_hit)

            # 计算成本（缓存命中时 cost=0）
            cost = 0.0 if is_cache_hit else self._calculate_cost(response.model, response.usage)

            # 记录使用量
            await self._record_usage(
                agent_name=agent_name,
                model=response.model,
                usage=response.usage,
                cost=cost,
                latency_ms=latency_ms,
                user_id=user_id,
                department_ids=department_ids,
            )

            # 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))
                if self._cache_manager is not None:
                    _span.set_attribute("gen_ai.cache.hit", is_cache_hit)
            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",
        timeout: float | None = None,
        user_id: str | None = None,
        department_ids: list[str] | None = None,
        db_path: Path | str | None = None,
        **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).

        Note: Streaming responses are NOT cached in this iteration.
        """
        resolved_model = self._resolve_model_alias(model)

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

        # ── Quota enforcement ──
        if department_ids and db_path:
            await self._enforce_quota(db_path, department_ids, resolved_model)

        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,
                timeout=timeout,
                **kwargs,
            )

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

            try:
                stream_obj = provider.chat_stream(stream_request)
                # Defensive: guard against misconfigured providers (e.g. an
                # AsyncMock in tests, or a future refactor that accidentally
                # turns chat_stream into a regular ``async def``) that return
                # a coroutine instead of an async generator. The original
                # cryptic error ``'async for' requires an object with
                # __aiter__ method, got coroutine`` becomes a clear,
                # actionable message naming the offending provider+model.
                if asyncio.iscoroutine(stream_obj):
                    logger.error(
                        f"Provider '{model_name}'.chat_stream returned a "
                        f"coroutine instead of an async generator. "
                        f"Check that the method is defined as "
                        f"``async def chat_stream(...): ...; yield ...``."
                    )
                    raise TypeError(
                        f"Provider '{model_name}' returned a coroutine "
                        f"from chat_stream() — expected an async "
                        f"generator. This indicates a provider "
                        f"implementation bug."
                    )
                async for chunk in stream_obj:
                    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)
                await self._record_usage(
                    agent_name=agent_name,
                    model=final_model,
                    usage=final_usage,
                    cost=cost,
                    latency_ms=latency_ms,
                    user_id=user_id,
                    department_ids=department_ids,
                )

                # Empty stream detection: if no content was produced,
                # raise error so the caller (ReActEngine) can retry with a different model.
                # We cannot continue to next model here because chunks may have already
                # been yielded to the client, which would cause mixed output.
                # Note: stream tool_calls are not tracked in chunks, so we only check content.
                if not total_content.strip():
                    logger.warning(f"Stream from '{model_name}' produced empty content")
                    raise LLMProviderError(
                        model_name,
                        f"Empty stream from {model_name}",
                    )

                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 get_provider_name_for_model(self, model: str) -> str | None:
        """返回 model 对应的 provider 名(用于 provider-specific 优化如 cache_control)。

        ponytail: 仅做 alias 解析 + provider 前缀提取,不查内部状态。
        升级路径:ServerConfig 显式声明 provider per model。
        返回 None 表示无法确定(多 provider + 无 "/" 前缀),调用方应回退到字符串拼接。
        """
        resolved = self._resolve_model_alias(model)
        if "/" in resolved:
            provider_name = resolved.split("/", 1)[0]
            if provider_name in self._providers:
                return provider_name
            return None
        # 无 "/" 前缀:仅当只有一个 provider 时能确定
        if len(self._providers) == 1:
            return next(iter(self._providers))
        return None

    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,
        )

    # ------------------------------------------------------------------
    # Quota enforcement helpers (U7)
    # ------------------------------------------------------------------

    async def _record_usage(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None,
        department_ids: list[str] | None,
    ) -> None:
        """Record a usage event via the async store interface (KTD-6).

        Multi-department attribution (U2): when a user belongs to
        multiple departments, a separate :class:`UsageRecord` is created
        for each department. This ensures ``get_usage(dept_id)`` returns
        the correct total for every department the user belongs to,
        matching the quota check scope (which checks all departments).

        TOCTOU (KTD-2): This method is called *after* the LLM response
        is received. Between ``_enforce_quota`` (before the call) and
        this recording, concurrent requests may push usage over the
        limit. This race window is accepted; post-hoc reconciliation
        (periodic scans for over-limit users) handles violations.
        """
        if not department_ids:
            # API key users (no departments) — record once with dept=None.
            await self._usage_tracker.record_async(
                agent_name=agent_name,
                model=model,
                usage=usage,
                cost=cost,
                latency_ms=latency_ms,
                user_id=user_id,
                department_id=None,
            )
            return

        # Record one entry per department so each department's aggregate
        # includes this usage. The cost is attributed in full to each
        # department (not split) — this matches how quota checks work
        # (each department is checked against the full usage).
        for dept_id in department_ids:
            await self._usage_tracker.record_async(
                agent_name=agent_name,
                model=model,
                usage=usage,
                cost=cost,
                latency_ms=latency_ms,
                user_id=user_id,
                department_id=dept_id,
            )

    async def _enforce_quota(
        self,
        db_path: Path | str,
        department_ids: list[str],
        resolved_model: str,
    ) -> None:
        """Run all quota checks for the given departments.

        Strictest-wins: if ANY department fails ANY check, raises
        :class:`QuotaExceededError` and the request is rejected.

        Both daily and monthly periods are checked (U2): for each
        department, ``token_limit`` and ``cost_limit`` are evaluated
        against both ``daily`` and ``monthly`` windows.

        Fail-closed (KTD-1): if the usage store is unavailable (Redis
        degraded), raises :class:`UsageStoreUnavailableError`. The
        caller must translate this to HTTP 503.

        TOCTOU (KTD-2): quota is checked *before* the LLM call, and
        usage is recorded *after*. Concurrent requests in this window
        may exceed the limit. This race is accepted; see
        :meth:`_record_usage` for the reconciliation strategy.
        """
        # Lazy import to avoid circular dependency (admin → ... → gateway).
        from agentkit.server.admin.quota_service import get_quota_service

        quota_service = get_quota_service()
        db = Path(db_path)

        for dept_id in department_ids:
            # 1. Model whitelist
            allowed, _reason = await quota_service.is_model_allowed(db, dept_id, resolved_model)
            if not allowed:
                raise QuotaExceededError(
                    department_id=dept_id,
                    quota_type="model_whitelist",
                    period="",
                    limit="",
                    current=resolved_model,
                )

            # 2. Token + cost limits (daily AND monthly)
            # 优化：每个 period 只查一次 get_usage，复用 summary 检查 token + cost
            for period in ("daily", "monthly"):
                summary = self._get_usage_summary(dept_id, period)
                current_tokens = int(summary.total_tokens)
                current_cost_cents = float(summary.total_cost) * 100.0
                await self._check_quota_value(
                    quota_service, db, dept_id, period, "token_limit", current_tokens
                )
                await self._check_quota_value(
                    quota_service, db, dept_id, period, "cost_limit", current_cost_cents
                )

    def _get_usage_summary(self, department_id: str, period: str) -> UsageSummary:
        """返回 department_id 在当前 period 的 usage summary（单次查询，供 token+cost 复用）。"""
        now = datetime.now(timezone.utc)
        if period == "monthly":
            start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        else:
            start = now.replace(hour=0, minute=0, second=0, microsecond=0)
        return self._usage_tracker.get_usage(
            department_id=department_id, start_time=start, end_time=now
        )

    async def _check_quota_value(
        self,
        quota_service: _QuotaServiceLike,
        db: Path,
        dept_id: str,
        period: str,
        quota_type: str,
        current: float,
    ) -> None:
        """检查单个配额（token_limit 或 cost_limit）— current 由调用方预计算传入。"""
        allowed, _reason = await quota_service.check_quota(db, dept_id, quota_type, period, current)
        if not allowed:
            quota = await quota_service.get_quota(db, dept_id, quota_type, period)
            limit = quota["limit_value"] if quota else None
            raise QuotaExceededError(
                department_id=dept_id,
                quota_type=quota_type,
                period=period,
                limit=limit,
                current=current,
            )
