import os
from openai import OpenAI
import time


class AI_Agent():
    """AI代理类，负责与AI模型交互生成文本内容"""
    
    def __init__(self, base_url, model_name, api):
        self.url_list = {
            "ali": "https://dashscope.aliyuncs.com/compatible-mode/v1",
            "kimi": "https://api.moonshot.cn/v1",
            "doubao": "https://ark.cn-beijing.volces.com/api/v3/",
            "deepseek": "https://api.deepseek.com",
            "vllm": "http://localhost:8000/v1",
        }
        
        self.base_url = self.url_list[base_url] if base_url in self.url_list else base_url
        self.api = api
        self.model_name = model_name
        
        self.client = OpenAI(
            api_key=self.api,
            base_url=self.base_url,
            # timeout=10
        )
        
    def generate_text(self, system_prompt, user_prompt, temperature, top_p, presence_penalty):
        """生成文本内容，并返回完整响应和token估计值"""
        print("系统提示词:")
        print(system_prompt)
        print("\n用户提示词:")
        print(user_prompt)
        print(f"\nAPI Key: {self.api}")
        print(f"Base URL: {self.base_url}")
        print(f"Model: {self.model_name}")
        
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[{"role": "system", "content": system_prompt}, 
                     {"role": "user", "content": user_prompt}],
            temperature=temperature, 
            top_p=top_p,
            # presence_penalty=presence_penalty,
            stream=True,
            max_tokens=8192,
            extra_body={
                "repetition_penalty": 1.05,
            },
        )
        
        # 收集完整的输出内容
        full_response = ""
        for chunk in response:
            if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None:
                content = chunk.choices[0].delta.content
                full_response += content
                print(content, end="", flush=True)
            if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].finish_reason == "stop":
                break
        
        print("\n完成生成，正在处理结果...")
        
        # 由于使用流式输出，无法获取真实的token计数，因此返回估计值
        estimated_tokens = len(full_response.split()) * 1.3  # 简单估算token数量
        
        return full_response, estimated_tokens

    def read_folder(self, file_folder):
        """读取指定文件夹下的所有文件内容"""
        if not os.path.exists(file_folder):
            return ""
            
        context = ""
        for file in os.listdir(file_folder):
            file_path = os.path.join(file_folder, file)
            if os.path.isfile(file_path):
                with open(file_path, "r", encoding="utf-8") as f:
                    context += f"文件名: {file}\n"
                    for line in f.readlines():
                        context += line
                    context += "\n"
        return context
    
    def work(self, system_prompt, user_prompt, file_folder, temperature, top_p, presence_penalty):
        """完整的工作流程：生成文本并返回结果"""
        # 生成时间戳
        date_time = time.strftime("%Y-%m-%d_%H-%M-%S")
        result_file = f"/root/autodl-tmp/xhsTweetGene/result/{date_time}.md"
        
        # 如果提供了参考文件夹，则读取其内容
        if file_folder:
            context = self.read_folder(file_folder)
            if context:
                user_prompt = f"{user_prompt}\n参考资料:\n{context}"
        
        # 计时
        time_start = time.time()
        
        # 生成文本
        result, tokens = self.generate_text(system_prompt, user_prompt, temperature, top_p, presence_penalty)
        
        # 计算耗时
        time_end = time.time()
        time_cost = time_end - time_start
        
        return result, system_prompt, user_prompt, file_folder, result_file, tokens, time_cost