from loguru import logger

from openai import AsyncOpenAI
from dataclasses import dataclass
from tenacity import retry, stop_after_attempt, stop_after_delay, wait_fixed
import asyncio
from core.config import MAX_WORKERS, LLMConfig, use_non_fastgpt_llm


class OpenAITool:
    def __init__(self, llm_config: LLMConfig, max_workers: int = MAX_WORKERS):
        self.max_workers = max_workers
        self.llm_config = llm_config
        self.client = AsyncOpenAI(
            base_url=llm_config.base_url, api_key=llm_config.api_key
        )

    @retry(stop=stop_after_delay(600) | stop_after_attempt(3), wait=wait_fixed(1))
    async def chat(self, msg, tools=None):
        if tools is None:
            extra_body = {}
            # fastgpt专用：如果第一个消息是system角色，则将其内容放入extra_body.variables.system，并从消息列表中移除
            if not use_non_fastgpt_llm:
                if msg[0]["role"] == "system":
                    extra_body = {"variables": {"system": msg[0]["content"]}}
                    msg = msg[1:]
            # deepseek专用关闭思考
            extra_body["thinking"] = {"type": "disabled"}
            try:
                response = await self.client.chat.completions.create(
                    model=self.llm_config.model, messages=msg, extra_body=extra_body
                )
                content = response.choices[0].message.content
                reasoning_content = response.choices[0].message.model_extra.get(
                    "reasoning_content", ""
                )
            except Exception as e:
                logger.error(f"LLM调用失败: {e} | response: {response}")
                raise e
            return content
        else:
            response = await self.client.chat.completions.create(
                model=self.llm_config.model,
                messages=msg,
                tools=tools,
                tool_choice="auto",
            )
            return response.choices[0].message.tool_calls

    async def mul_chat(self, msgs, tools=None):
        sem = asyncio.Semaphore(self.max_workers)

        async def _wrapped(m):
            async with sem:
                return await self.chat(m, tools)

        return await asyncio.gather(*[_wrapped(m) for m in msgs])


if __name__ == "__main__":
    import json

    llm_config = LLMConfig()
    tool = OpenAITool(llm_config)
    messages = [
        {
            "role": "system",
            "content": "你是我的人工智能助手，协助我分析问题并提供建议。",
        },
        {"role": "user", "content": "请分析以下问题：为什么天空是蓝色的？"},
    ]
    response = asyncio.run(tool.chat(messages))
    print("LLM Response:", response)
