import json
from config.logger import setup_logging
import requests
from core.providers.llm.base import LLMProviderBase
from core.providers.llm.system_prompt import get_system_prompt_for_function
from core.utils.util import check_model_key

TAG = __name__
logger = setup_logging()


class LLMProvider(LLMProviderBase):
    def __init__(self, config):
        self.api_key = config["api_key"]
        self.mode = config.get("mode", "chat-messages")
        self.base_url = config.get("base_url", "https://api.dify.ai/v1").rstrip("/")
        self.session_conversation_map = {}  # 存储session_id和conversation_id的映射
        model_key_msg = check_model_key("DifyLLM", self.api_key)
        if model_key_msg:
            logger.bind(tag=TAG).error(model_key_msg)

    def response(self, session_id, dialogue, **kwargs):
        # 取最后一条用户消息
        last_msg = next(m for m in reversed(dialogue) if m["role"] == "user")
        conversation_id = self.session_conversation_map.get(session_id)

        # 发起流式请求
        if self.mode == "chat-messages":
            request_json = {
                "query": last_msg["content"],
                "response_mode": "streaming",
                "user": session_id,
                "inputs": {},
                "conversation_id": conversation_id,
            }
        elif self.mode == "workflows/run":
            request_json = {
                "inputs": {"query": last_msg["content"]},
                "response_mode": "streaming",
                "user": session_id,
            }
        elif self.mode == "completion-messages":
            request_json = {
                "inputs": {"query": last_msg["content"]},
                "response_mode": "streaming",
                "user": session_id,
            }

        with requests.post(
            f"{self.base_url}/{self.mode}",
            headers={"Authorization": f"Bearer {self.api_key}"},
            json=request_json,
            stream=True,
        ) as r:
            if self.mode == "chat-messages":
                for line in r.iter_lines():
                    if line.startswith(b"data: "):
                        event = json.loads(line[6:])
                        # 如果没有找到conversation_id，则获取此次conversation_id
                        if not conversation_id:
                            conversation_id = event.get("conversation_id")
                            self.session_conversation_map[session_id] = (
                                conversation_id  # 更新映射
                            )
                        # 过滤 message_replace 事件，此事件会全量推一次
                        if event.get("event") != "message_replace" and event.get(
                            "answer"
                        ):
                            yield event["answer"]
            elif self.mode == "workflows/run":
                for line in r.iter_lines():
                    if line.startswith(b"data: "):
                        event = json.loads(line[6:])
                        if event.get("event") == "workflow_finished":
                            if event["data"]["status"] == "succeeded":
                                yield event["data"]["outputs"]["answer"]
                            else:
                                yield "【服务响应异常】"
            elif self.mode == "completion-messages":
                for line in r.iter_lines():
                    if line.startswith(b"data: "):
                        event = json.loads(line[6:])
                        # 过滤 message_replace 事件，此事件会全量推一次
                        if event.get("event") != "message_replace" and event.get(
                            "answer"
                        ):
                            yield event["answer"]

    def response_with_functions(self, session_id, dialogue, functions=None):
        if len(dialogue) == 2 and functions is not None and len(functions) > 0:
            # 第一次调用llm， 取最后一条用户消息，附加tool提示词
            last_msg = dialogue[-1]["content"]
            function_str = json.dumps(functions, ensure_ascii=False)
            modify_msg = get_system_prompt_for_function(function_str) + last_msg
            dialogue[-1]["content"] = modify_msg

        # 如果最后一个是 role="tool"，附加到user上
        if len(dialogue) > 1 and dialogue[-1]["role"] == "tool":
            assistant_msg = "\ntool call result: " + dialogue[-1]["content"] + "\n\n"
            while len(dialogue) > 1:
                if dialogue[-1]["role"] == "user":
                    dialogue[-1]["content"] = assistant_msg + dialogue[-1]["content"]
                    break
                dialogue.pop()

        for token in self.response(session_id, dialogue):
            yield token, None