"""可复用的 Jev 决策客户端；重试和 HTTP 连接管理交给官方 SDK"""

from functools import cache
from math import isclose
from typing import Self

from pydantic import ValidationError
from typesafe_sdk import (
    AsyncTypeSafeClient,
    RetryPolicy,
    SystemOneResponse,
    TypeSafeAPIConnectionError,
    TypeSafeAPIError,
    TypeSafeAPIResponseValidationError,
    TypeSafeAPITimeoutError,
    TypeSafeError,
)

from app.config import Settings, settings
from app.exceptions import DecisionError
from app.schemas.decision import (
    ChoiceDecision,
    ChoiceResult,
    DecisionRequest,
    DecisionResponse,
    ScoreDecision,
    ScoreResult,
)


class DecisionEngine:
    def __init__(
        self,
        config: Settings | None = None,
        *,
        client: AsyncTypeSafeClient | None = None,
    ):
        self._config = config if config is not None else settings
        self._client = client
        self._owns_client = client is None
        self._closed = False

    @property
    def closed(self) -> bool:
        return self._closed

    def _get_client(self) -> AsyncTypeSafeClient:
        if self._closed:
            raise DecisionError("决策客户端已关闭", kind="closed")
        if self._client is None:
            config = self._config
            if not config.typesafe_api_key.strip():
                raise DecisionError("请配置 TYPESAFE_API_KEY", kind="configuration")
            try:
                self._client = AsyncTypeSafeClient(
                    api_key=config.typesafe_api_key,
                    model=config.typesafe_model,
                    base_url=config.typesafe_base_url,
                    timeout=config.typesafe_timeout,
                    retry=RetryPolicy(
                        max_retries=config.typesafe_max_retries,
                        timeout=config.typesafe_timeout,
                    ),
                )
            except (TypeSafeError, ValueError) as error:
                raise DecisionError(
                    "Jev 客户端配置无效", kind="configuration"
                ) from error
        return self._client

    async def run(self, request: DecisionRequest) -> DecisionResponse:
        # 即便模型实例在创建后被修改，也必须重新经过请求校验
        request = DecisionRequest.model_validate(request.model_dump(mode="json"))
        client = self._get_client()
        questions = {
            name: question.model_dump(mode="json", exclude_unset=True)
            | {"type": question.type}
            for name, question in request.decisions.items()
        }
        from app.observability.tracing import observe_decision, record_decision

        # 观测包在请求外面：成功才记结果，失败仍抛原来的 DecisionError
        # generation 为 None 时 record_decision 什么都不做
        with observe_decision(request) as generation:
            try:
                response = await client.system_one(
                    state=request.state,
                    questions=questions,
                    model=request.model,
                    extra_body=request.extra_body or None,
                )
            except TypeSafeAPIResponseValidationError as error:
                raise DecisionError(
                    "Jev 返回的数据格式无效",
                    kind="response_validation",
                    status=error.status,
                    request_id=error.request_id,
                ) from error
            except TypeSafeAPIError as error:
                raise DecisionError(
                    f"Jev 请求失败（HTTP {error.status}）",
                    kind="api",
                    status=error.status,
                    request_id=error.request_id,
                ) from error
            except TypeSafeAPITimeoutError as error:
                raise DecisionError("Jev 请求超时", kind="timeout") from error
            except TypeSafeAPIConnectionError as error:
                raise DecisionError("无法连接 Jev 服务", kind="connection") from error
            except TypeSafeError as error:
                raise DecisionError("Jev SDK 调用失败", kind="sdk") from error
            parsed = self._parse_response(request, response)
            record_decision(generation, parsed)
            return parsed

    @staticmethod
    def _parse_response(
        request: DecisionRequest, response: SystemOneResponse
    ) -> DecisionResponse:
        request_id = response.raw_http_response.headers.get("x-typesafe-request-id")
        try:
            result = DecisionResponse.model_validate(
                {
                    "results": {
                        name: answer.model_dump(mode="json")
                        for name, answer in response.answers.items()
                    },
                    "model": response.model,
                    "usage": response.usage.model_dump(mode="json"),
                    "request_id": request_id,
                }
            )
            if result.results.keys() != request.decisions.keys():
                raise ValueError("返回的问题 ID 与请求不一致")
            for name, question in request.decisions.items():
                answer = result.results[name]
                if answer.type != question.type:
                    raise ValueError("返回的决策类型与请求不一致")
                if isinstance(question, ChoiceDecision) and isinstance(
                    answer, ChoiceResult
                ):
                    if answer.choice not in question.criteria:
                        raise ValueError("返回了候选集合之外的选项")
                    if answer.probabilities.keys() != question.criteria.keys():
                        raise ValueError("选项概率不完整")
                if isinstance(question, ScoreDecision) and isinstance(
                    answer, ScoreResult
                ):
                    levels = {str(index) for index in range(len(question.criteria))}
                    if (
                        answer.probabilities.keys() != levels
                        or answer.legend.keys() != levels
                    ):
                        raise ValueError("评分等级不完整")
                    if not 0 <= answer.score <= len(question.criteria) - 1:
                        raise ValueError("评分超出等级范围")
                if isinstance(answer, (ChoiceResult, ScoreResult)):
                    if not isclose(sum(answer.probabilities.values()), 1, abs_tol=1e-4):
                        raise ValueError("概率之和不为 1")
            return result
        except (ValidationError, ValueError) as error:
            raise DecisionError(
                "Jev 返回结果与决策协议不匹配",
                kind="response_validation",
                request_id=request_id,
            ) from error

    async def aclose(self) -> None:
        if not self._closed:
            if self._owns_client and self._client is not None:
                await self._client.aclose()
            self._closed = True

    async def __aenter__(self) -> Self:
        self._get_client()
        return self

    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        await self.aclose()


@cache
def _cached_engine() -> DecisionEngine:
    return DecisionEngine()


def get_decision_engine() -> DecisionEngine:
    """与 get_llm 类似的共享入口；适用于单个服务事件循环"""
    engine = _cached_engine()
    if engine.closed:
        _cached_engine.cache_clear()
        engine = _cached_engine()
    return engine


async def close_decision_engine() -> None:
    """在应用关闭阶段调用；跨 asyncio.run 时使用独立实例"""
    if _cached_engine.cache_info().currsize:
        await _cached_engine().aclose()
        _cached_engine.cache_clear()