"""
本地 CustomLLM 服务。

当前实现直接在同一个 FastAPI 进程内调用方舟 SDK，
并由 /api/chat_callback 对外提供火山要求的 SSE 回调接口。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Iterator

from utils.env import env_float, env_int, env_str

DEFAULT_ARK_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
DEFAULT_ARK_TIMEOUT_SECONDS = 1800
DEFAULT_ARK_TEMPERATURE = 0


def _coalesce(*values):
    for value in values:
        if value is not None:
            return value
    return None


@dataclass(frozen=True)
class LocalLLMSettings:
    api_key: str
    endpoint_id: str
    base_url: str
    timeout_seconds: int
    system_prompt: str
    default_temperature: float


def _load_settings() -> LocalLLMSettings:
    api_key = env_str("ARK_API_KEY")
    endpoint_id = env_str("ARK_ENDPOINT_ID")

    if not api_key:
        raise ValueError("ARK_API_KEY 不能为空")
    if not endpoint_id:
        raise ValueError("ARK_ENDPOINT_ID 不能为空")

    return LocalLLMSettings(
        api_key=api_key,
        endpoint_id=endpoint_id,
        base_url=env_str("ARK_BASE_URL", DEFAULT_ARK_BASE_URL),
        timeout_seconds=env_int("ARK_TIMEOUT_SECONDS", DEFAULT_ARK_TIMEOUT_SECONDS),
        system_prompt=env_str(
            "LOCAL_LLM_SYSTEM_PROMPT",
            env_str("CUSTOM_LLM_SYSTEM_MESSAGE"),
        ),
        default_temperature=env_float(
            "LOCAL_LLM_TEMPERATURE", DEFAULT_ARK_TEMPERATURE
        ),
    )


class LocalLLMService:
    def __init__(self):
        self._client = None
        self._settings: LocalLLMSettings | None = None

    @property
    def settings(self) -> LocalLLMSettings:
        if self._settings is None:
            self._settings = _load_settings()
        return self._settings

    def _get_client(self):
        if self._client is not None:
            return self._client

        try:
            from volcenginesdkarkruntime import Ark
        except ImportError as exc:
            raise RuntimeError(
                "未安装 volcenginesdkarkruntime，请先执行 uv sync 安装依赖"
            ) from exc

        s = self.settings
        self._client = Ark(
            base_url=s.base_url,
            api_key=s.api_key,
            timeout=s.timeout_seconds,
        )
        return self._client

    def chat_stream(
        self,
        history_messages: list[dict[str, Any]],
        rag_context: str = "",
        request_options: dict[str, Any] | None = None,
    ) -> Iterator[Any]:
        settings = self.settings
        client = self._get_client()
        request_options = request_options or {}

        system_blocks = [settings.system_prompt]
        if rag_context:
            system_blocks.append(f"### 参考知识库（绝对准则）\n{rag_context.strip()}")

        messages = [
            {
                "role": "system",
                "content": "\n\n".join(block for block in system_blocks if block),
            }
        ]
        messages.extend(history_messages)

        payload: dict[str, Any] = {
            "model": settings.endpoint_id,
            "messages": messages,
            "temperature": _coalesce(
                request_options.get("temperature"),
                settings.default_temperature,
            ),
            "stream": True,
            "stream_options": {"include_usage": True},
        }

        if request_options.get("max_tokens") is not None:
            payload["max_tokens"] = request_options["max_tokens"]
        if request_options.get("top_p") is not None:
            payload["top_p"] = request_options["top_p"]

        print(f"🚀 发起流式调用 (Endpoint: {settings.endpoint_id})")
        try:
            stream = client.chat.completions.create(**payload)
            for chunk in stream:
                yield chunk
        except Exception as exc:
            print(f"❌ LLM 调用失败: {exc}")
            raise


local_llm_service = LocalLLMService()