"""The LLM gateway (架构文档 §8.1).

Built-in capabilities, all exercised on every call:

===================  =========================================================
场景化路由            ``scene → (provider, model, temperature, max_tokens, timeout)``
故障转移 + 熔断        primary 超时/5xx/限流 → 自动切备用；连续失败熔断，冷却后半开
结构化输出            传入 Pydantic Schema → JSON 约束 + 解析失败重试 + 校验
Prompt 版本           模板文件按 id+version 加载，版本号随每次调用落库
缓存前缀纪律          固定前缀含时间戳/随机 ID/昵称 → 拒绝或告警 (§7.2, TD-04)
成本计量              每次 provider 调用（含失败与重试）写一行 ``ob_ai_call``
预算保护              平台级日成本四档分级降级：warn/secondary/generation/block_all (§13.3)
===================  =========================================================

Deliberately not here yet: streaming (SSE), 审核联动 (input/output moderation),
per-user rate limiting, A/B prompt split by user hash. See packages/llm/README.md.
"""

from __future__ import annotations

import re
import time
from collections.abc import Callable, Mapping, Sequence
from decimal import Decimal

from pydantic import BaseModel, ValidationError

from packages.domain.ids import next_id
from packages.llm.budget import DailyBudgetGuard
from packages.llm.circuit import CircuitBreakerRegistry
from packages.llm.cost import AICallRecord, AICallRecorder, RateResolver, compute_cost_cents
from packages.llm.errors import (
    AllProvidersFailedError,
    CircuitOpenError,
    ProviderError,
    SceneNotRoutedError,
    StructuredOutputError,
)
from packages.llm.prompts import PromptRegistry, validate_cache_prefix
from packages.llm.providers.base import LLMProvider
from packages.llm.routing import SceneRouter
from packages.llm.types import (
    Completion,
    Message,
    ProviderRequest,
    Scene,
    SceneConfig,
    Usage,
)
from packages.shared.logging import get_logger

logger = get_logger(__name__)

_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL)
_STRUCTURED_INVALID = "llm_structured_output_invalid"


def extract_json(text: str) -> str:
    """Pull the JSON object out of a model reply that may be fenced or chatty."""
    match = _JSON_FENCE_RE.search(text)
    candidate = (match.group(1) if match else text).strip()
    start, end = candidate.find("{"), candidate.rfind("}")
    if start != -1 and end > start:
        return candidate[start : end + 1]
    return candidate


class LLMGateway:
    """Single entry point for every LLM call in the platform."""

    def __init__(
        self,
        *,
        providers: Mapping[str, LLMProvider],
        router: SceneRouter,
        prompts: PromptRegistry,
        rates: RateResolver,
        recorder: AICallRecorder,
        budget: DailyBudgetGuard,
        breakers: CircuitBreakerRegistry | None = None,
        id_factory: Callable[[], int] = next_id,
        timer: Callable[[], float] = time.perf_counter,
    ) -> None:
        self._providers = dict(providers)
        self._router = router
        self._prompts = prompts
        self._rates = rates
        self._recorder = recorder
        self._budget = budget
        self._breakers = breakers or CircuitBreakerRegistry()
        self._id_factory = id_factory
        self._timer = timer

    # -- public API ---------------------------------------------------------

    async def complete(
        self,
        *,
        scene: Scene,
        messages: Sequence[Message] = (),
        prompt_vars: Mapping[str, str] | None = None,
        schema: type[BaseModel] | None = None,
        user_id: int | None = None,
        ref_type: str | None = None,
        ref_id: str | None = None,
    ) -> Completion:
        config = self._router.resolve(scene)

        # Gate 1: four-tier graded daily cost breaker (架构文档 §13.3),
        # checked before any spend — 在途任务不受影响，这里只拦新请求。
        await self._budget.ensure_scene_allowed(scene=scene.value, critical=config.critical)

        template = self._prompts.get(config.prompt_id, config.prompt_version)
        violations = validate_cache_prefix(
            template.system_prefix, mode=self._prompts.cache_prefix_mode, ref=template.ref
        )
        if violations:
            logger.warning(
                "llm_prompt_cache_prefix_violation",
                prompt=template.ref,
                violations=[str(v) for v in violations],
            )

        composed = self._compose(
            system_prefix=template.system_prefix,
            user_text=template.render_user(dict(prompt_vars or {})),
            messages=messages,
            schema=schema,
        )

        failures: list[Exception] = []
        for provider_name in config.providers:
            provider = self._providers.get(provider_name)
            if provider is None:
                failures.append(
                    SceneNotRoutedError(f"provider '{provider_name}' is not registered")
                )
                continue

            breaker = self._breakers.get(provider_name)
            if not breaker.allow():
                logger.warning(
                    "llm_circuit_open", provider=provider_name, state=breaker.state.value
                )
                failures.append(CircuitOpenError(provider_name, f"circuit {breaker.state.value}"))
                continue

            try:
                completion = await self._call_provider(
                    provider=provider,
                    config=config,
                    messages=composed,
                    schema=schema,
                    user_id=user_id,
                    prompt_id=template.prompt_id,
                    prompt_version=template.version,
                    ref_type=ref_type,
                    ref_id=ref_id,
                )
            except ProviderError as exc:
                if not exc.retryable:
                    # A malformed request would fail on the backup too.
                    raise
                breaker.record_failure()
                failures.append(exc)
                logger.warning(
                    "llm_provider_failed",
                    provider=provider_name,
                    scene=scene.value,
                    error_code=exc.code,
                    circuit=breaker.state.value,
                )
                continue
            except StructuredOutputError as exc:
                breaker.record_failure()
                failures.append(exc)
                continue

            breaker.record_success()
            return completion

        raise AllProvidersFailedError(scene.value, failures)

    def breaker_states(self) -> dict[str, str]:
        return self._breakers.snapshot()

    async def aclose(self) -> None:
        for provider in self._providers.values():
            await provider.aclose()

    # -- internals ----------------------------------------------------------

    @staticmethod
    def _compose(
        *,
        system_prefix: str,
        user_text: str,
        messages: Sequence[Message],
        schema: type[BaseModel] | None,
    ) -> list[Message]:
        """架构文档 §7.2 ③: 固定前缀放最前以命中缓存，变量部分一律排在后面。"""
        composed: list[Message] = [Message(role="system", content=system_prefix)]
        composed.extend(messages)
        if user_text:
            composed.append(Message(role="user", content=user_text))
        if schema is not None:
            composed.append(Message(role="user", content=_schema_instruction(schema)))
        return composed

    async def _call_provider(
        self,
        *,
        provider: LLMProvider,
        config: SceneConfig,
        messages: Sequence[Message],
        schema: type[BaseModel] | None,
        user_id: int | None,
        prompt_id: str,
        prompt_version: str,
        ref_type: str | None,
        ref_id: str | None,
    ) -> Completion:
        model = config.model_for(provider.name, provider.default_model)
        conversation: list[Message] = list(messages)
        max_attempts = config.structured_max_retries + 1 if schema is not None else 1
        schema_name = schema.__name__ if schema is not None else "<none>"
        last_detail = "no attempt was made"

        for attempt in range(1, max_attempts + 1):
            request = ProviderRequest(
                model=model,
                messages=tuple(conversation),
                temperature=config.temperature,
                max_tokens=config.max_tokens,
                timeout_seconds=config.timeout_seconds,
                json_mode=schema is not None,
                thinking=config.thinking,
            )
            call_id = self._id_factory()
            started = self._timer()

            try:
                response = await provider.complete(request)
            except ProviderError as exc:
                await self._write_record(
                    AICallRecord(
                        call_id=call_id,
                        scene=config.scene.value,
                        provider=provider.name,
                        model=model,
                        success=False,
                        usage=Usage(),
                        latency_ms=self._elapsed_ms(started),
                        cost_cents=Decimal("0"),
                        user_id=user_id,
                        prompt_id=prompt_id,
                        prompt_version=prompt_version,
                        rate_tier=self._rates.current_tier(),
                        error_code=exc.code,
                        error_message=str(exc)[:1000],
                        ref_type=ref_type,
                        ref_id=ref_id,
                    )
                )
                raise

            latency_ms = self._elapsed_ms(started)
            # 单价与档位在同一刻取，成本落库后不再回算 —— 跨时段重算会改掉历史成本。
            rate_tier = self._rates.current_tier()
            cost = compute_cost_cents(response.usage, self._rates.for_model(response.model))

            parsed: BaseModel | None = None
            success = True
            error_code: str | None = None
            if schema is not None:
                try:
                    parsed = schema.model_validate_json(extract_json(response.text))
                except (ValidationError, ValueError) as exc:
                    success = False
                    error_code = _STRUCTURED_INVALID
                    last_detail = str(exc)[:400]

            await self._write_record(
                AICallRecord(
                    call_id=call_id,
                    scene=config.scene.value,
                    provider=provider.name,
                    model=response.model,
                    success=success,
                    usage=response.usage,
                    latency_ms=latency_ms,
                    cost_cents=cost,
                    user_id=user_id,
                    prompt_id=prompt_id,
                    prompt_version=prompt_version,
                    rate_tier=rate_tier,
                    error_code=error_code,
                    error_message=last_detail if not success else None,
                    ref_type=ref_type,
                    ref_id=ref_id,
                )
            )
            # Budget accrues on every attempt, including ones we then discard —
            # a validation-retry loop costs real money and must be visible.
            await self._budget.add_cost(cost)

            if success:
                return Completion(
                    text=response.text,
                    scene=config.scene,
                    provider=provider.name,
                    model=response.model,
                    prompt_id=prompt_id,
                    prompt_version=prompt_version,
                    usage=response.usage,
                    latency_ms=latency_ms,
                    cost_cents=cost,
                    call_id=call_id,
                    parsed=parsed,
                    attempts=attempt,
                )

            logger.warning(
                "llm_structured_output_invalid",
                provider=provider.name,
                scene=config.scene.value,
                schema=schema_name,
                attempt=attempt,
            )
            conversation.append(Message(role="assistant", content=response.text))
            conversation.append(Message(role="user", content=_repair_instruction(last_detail)))

        raise StructuredOutputError(provider.name, schema_name, max_attempts, last_detail)

    def _elapsed_ms(self, started: float) -> int:
        return max(int((self._timer() - started) * 1000), 0)

    async def _write_record(self, record: AICallRecord) -> None:
        """Cost attribution is mandatory but must never break the user's request."""
        try:
            await self._recorder.record(record.with_defaults())
        except Exception:
            logger.exception("ai_call_record_failed", call_id=record.call_id, scene=record.scene)


def _schema_instruction(schema: type[BaseModel]) -> str:
    import json

    return (
        "请只输出一个 JSON 对象，不要包含解释文字或代码块标记。"
        "该对象必须严格符合以下 JSON Schema：\n"
        f"{json.dumps(schema.model_json_schema(), ensure_ascii=False)}"
    )


def _repair_instruction(detail: str) -> str:
    return (
        "上一次输出不符合要求的 JSON Schema，校验错误如下：\n"
        f"{detail}\n"
        "请重新只输出修正后的 JSON 对象，不要包含任何其他文字。"
    )