#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2026/2/10 15:20
# @Author  : linghu
# @File    : LLMClient.py
# @Software: PyCharm
import os

from dotenv import load_dotenv

from phad_llm_sdk.dto.req.llm_send import LLMRequest
from phad_llm_sdk.utils.HttpUtils import request_llm

load_dotenv()

VLLM_LLM_URL = os.getenv("VLLM_LLM_URL")
LLM_QWEN_PATH = os.getenv("LLM_QWEN_PATH")

class MessageBuilder:
    """
    消息构建器，用于方便地构建消息列表
    """
    def __init__(self):
        """
        初始化消息构建器
        """
        self.messages = []
    
    def add_message(self, role, content):
        """
        添加一条消息
        
        Args:
            role: 角色，例如 "system"、"user"、"assistant"
            content: 消息内容
        
        Returns:
            self: 返回自身，支持链式调用
        """
        self.messages.append({"role": role, "content": content})
        return self
    
    def add_system_message(self, content):
        """
        添加系统消息
        
        Args:
            content: 系统消息内容
        
        Returns:
            self: 返回自身，支持链式调用
        """
        return self.add_message("system", content)
    
    def add_user_message(self, content):
        """
        添加用户消息
        
        Args:
            content: 用户消息内容
        
        Returns:
            self: 返回自身，支持链式调用
        """
        return self.add_message("user", content)
    
    def add_assistant_message(self, content):
        """
        添加助手消息
        
        Args:
            content: 助手消息内容
        
        Returns:
            self: 返回自身，支持链式调用
        """
        return self.add_message("assistant", content)
    
    def build(self):
        """
        构建消息列表
        
        Returns:
            list: 消息列表
        """
        return self.messages

class LLMClient:
    """
    大模型客户端，统一调用接口
    """
    _instance = None
    
    def __new__(cls):
        """
        单例模式
        """
        if cls._instance is None:
            cls._instance = super(LLMClient, cls).__new__(cls)
        return cls._instance
    
    def __init__(self):
        """
        初始化客户端
        """
        # 拼接完整的URL路径
        if LLM_QWEN_PATH:
            # 确保URL末尾没有斜杠
            base_url = VLLM_LLM_URL.rstrip('/')
            # 确保路径开头有斜杠
            path = LLM_QWEN_PATH.lstrip('/')
            # 拼接URL
            self.url = f"{base_url}/{path}"
        else:
            self.url = VLLM_LLM_URL
    
    def chat(self, messages, stream=False, max_tokens=1000, temperature=0.7, top_p=0.1, think=False):
        """
        调用大模型进行对话
        
        Args:
            messages: 消息列表，格式为 [{"role": "system", "content": "..."}, {"role": "user", "content": "..."}]
            stream: 是否使用流式响应
            max_tokens: 最大生成 token 数
            temperature: 温度参数
            top_p: 采样的核概率
            think: 是否返回思考过程
        
        Returns:
            如果 stream=True，返回响应对象
            如果 stream=False，返回 {"message": {"content": "回答内容"}}
        """
        # 创建LLMRequest对象
        request_obj = LLMRequest(
            messages=messages,
            stream=stream,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            think=think
        )
        
        # 调用request_llm函数
        return request_llm(self.url, request_obj)
    
    def create_message_builder(self):
        """
        创建消息构建器
        
        Returns:
            MessageBuilder: 消息构建器实例
        """
        return MessageBuilder()

# 创建全局实例
llm_client = LLMClient()

# 示例用法
if __name__ == "__main__":
    # 示例1: 使用消息构建器构建消息列表（非流式请求）
    message_builder = llm_client.create_message_builder()
    messages = message_builder\
        .add_system_message("你是一个AI助手，喜欢帮助别人。")\
        .add_user_message("你好，能帮我解释一下什么是机器学习吗？")\
        .build()

    # 使用全局实例调用
    response = llm_client.chat(messages)
    print("非流式响应:")
    print(response)
    
    # 示例2: 使用消息构建器构建消息列表（流式请求）
    print("\n流式响应:")
    stream_message_builder = llm_client.create_message_builder()
    stream_messages = stream_message_builder\
        .add_system_message("你是一个AI助手，喜欢帮助别人。")\
        .add_user_message("你好，能帮我解释一下什么是深度学习吗？")\
        .build()
    
    stream_response = llm_client.chat(stream_messages, stream=True)
    if stream_response:
        import json
        for chunk in stream_response.iter_lines():
            if chunk:
                chunk = chunk.decode('utf-8')
                if chunk.startswith('data: '):
                    chunk = chunk[6:]
                if chunk == '[DONE]':
                    break
                try:
                    chunk_data = json.loads(chunk)
                    content = chunk_data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                    if content:
                        print(content, end="")
                except json.JSONDecodeError:
                    pass
        print()
    else:
        print("流式请求失败")