from typing import Any, cast
import openai

from loguru import logger
from openai.types.chat import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageParam
from pydantic import BaseModel, Field

from konabot.common.path import CONFIG_PATH

LLM_CONFIG_PATH = CONFIG_PATH / 'llm.json'

if not LLM_CONFIG_PATH.exists():
    LLM_CONFIG_PATH.write_text("{}")


class LLMInfo(BaseModel):
    base_url: str
    api_key: str
    model_name: str

    def get_openai_client(self):
        return openai.AsyncClient(
            api_key=self.api_key,
            base_url=self.base_url,
        )

    async def chat(
        self,
        messages: list[ChatCompletionMessageParam] | list[dict[str, Any]],
        timeout: float | None = 30.0,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> ChatCompletionMessage:
        logger.info(f"调用 LLM: BASE_URL={self.base_url} MODEL_NAME={self.model_name}")
        completion: ChatCompletion = await self.get_openai_client().chat.completions.create(
            messages=cast(Any, messages),
            model=self.model_name,
            max_tokens=max_tokens,
            timeout=timeout,
            stream=False,
            **kwargs,
        )
        choice = completion.choices[0]
        logger.info(
            f"调用 LLM 完成: BASE_URL={self.base_url} MODEL_NAME={self.model_name} REASON={choice.finish_reason}"
        )
        return choice.message


class LLMConfig(BaseModel):
    llms: dict[str, LLMInfo] = Field(default_factory=dict)
    default_llm: str = "Qwen2.5-7B-Instruct"


llm_config = LLMConfig.model_validate_json(LLM_CONFIG_PATH.read_text())


def get_llm(llm_model: str | None = None):
    if llm_model is None:
        llm_model = llm_config.default_llm
    if llm_model not in llm_config.llms:
        if llm_config.default_llm in llm_config.llms:
            logger.warning(f"[LLM] 需求的 LLM 不存在，回退到默认模型 REQUIRED={llm_model}")
            return llm_config.llms[llm_config.default_llm]
        raise NotImplementedError("[LLM] LLM 未配置，该功能无法使用")
    return llm_config.llms[llm_model]