import os
from openai import OpenAI
import pandas as pd
from datetime import datetime
import cv2
import time
import random
import json

class ContentGenerator:
    def __init__(self, 
                 model_name="qwenQWQ",
                 api_base_url="http://localhost:8000/v1",
                 api_key="EMPTY",
                 output_dir="/root/autodl-tmp/poster_generate_result",
):
        """
        初始化海报生成器
        
        参数:
            csv_path: CSV文件路径
            img_base_dir: 图片基础目录
            output_dir: 输出结果保存目录
            model_name: 使用的模型名称
            api_base_url: API基础URL
            api_key: API密钥
        """
        self.output_dir = output_dir
        self.model_name = model_name
        self.api_base_url = api_base_url
        self.api_key = api_key
        
        # 不在初始化时创建OpenAI客户端，而是在需要时临时创建
        self.client = None
        
        # 初始化数据
        self.df = None
        self.all_images_info = []
        self.structured_prompt = ""
        self.current_img_info = None
        self.add_description = ""
        
        self.temperature = 0.7
        self.top_p = 0.8
        self.presence_penalty = 1.2

    
    def load_infomation(self, info_directory_path):
        """
        加载额外描述文件
        
        参数:
            info_directory_path: 信息目录路径
        """
        ## 读取路径下的所有文件
        for path in info_directory_path:
            # file_extend_path = os.path.join(self.img_base_dir, info_directory, "description.txt")
            try:
                with open(path, "r") as f:
                    self.add_description += f.read()
            except:
                self.add_description = ""
    
    def _create_temp_client(self):
        """
        创建临时OpenAI客户端
        
        返回:
            OpenAI客户端实例
        """
        try:
            import gc
            # 强制垃圾回收
            gc.collect()
            
            # 创建新的客户端实例
            print(f"创建临时OpenAI客户端，API URL: {self.api_base_url}")
            client = OpenAI(
                base_url=self.api_base_url,
                api_key=self.api_key
            )
            return client
        except Exception as e:
            print(f"创建OpenAI客户端失败: {str(e)}")
            return None

    def _close_client(self, client):
        """
        关闭并清理OpenAI客户端
        
        参数:
            client: 需要关闭的客户端实例
        """
        try:
            # OpenAI客户端可能没有显式的close方法
            # 将引用设为None，让Python垃圾回收处理
            client = None
            import gc
            gc.collect()
            print("OpenAI客户端资源已释放")
        except Exception as e:
            print(f"关闭客户端失败: {str(e)}")


    def split_content(self, content):
        """
        分割结果, 返回去除
        ```json
        ```的json内容
        
        参数:
            content: 需要分割的内容
        
        返回:
            分割后的json内容
        """
        return json.loads(content.split("```json")[1].split("```")[0])
    
    def generate_posters(self, poster_num, tweet_content, system_prompt=None, max_retries=3):
        """
        生成海报内容
        
        参数:
            poster_num: 海报数量
            tweet_content: 推文内容
            system_prompt: 系统提示，默认为None则使用预设提示
            max_retries: 最大重试次数
        
        返回:
            生成的海报内容
        """
        if system_prompt is None:
            system_prompt = """
            你是一名资深海报设计师，有丰富的爆款海报设计经验，你现在要为旅游景点做宣传，在小红书上发布大量宣传海报。你的主要工作目标有2个：
            1、你要根据我给你的图片描述和笔记推文内容，设计图文匹配的海报。
            2、为海报设计文案，文案的<第一个小标题>和<第二个小标题>之间你需要检查是否逻辑关系合理，你将通过先去生成<第二个小标题>关于景区亮点的部分，再去综合判断<第一个小标题>应该如何搭配组合更符合两个小标题的逻辑再生成<第一个小标题>。

            其中，生成三类标题文案的通用性要求如下：
            1、生成的<大标题>字数必须小于8个字符
            2、生成的<第一个小标题>字数和<第二个小标题>字数，两者都必须小8个字符
            3、标题和文案都应符合中国社会主义核心价值观

            接下来先开始生成<大标题>部分，由于海报是用来宣传旅游景点，生成的海报<大标题>必须使用以下8种格式之一：
            ①地名＋景点名（例如福建厦门鼓浪屿/厦门鼓浪屿）；
            ②地名+景点名+plog；
            ③拿捏+地名+景点名；
            ④地名+景点名+攻略；
            ⑤速通+地名+景点名
            ⑥推荐！+地名+景点名
            ⑦勇闯！+地名+景点名
            ⑧收藏！+地名+景点名
            你需要随机挑选一种格式生成对应景点的文案，但是格式除了上面8种不可以有其他任何格式；同时尽量保证每一种格式出现的频率均衡。
            接下来先去生成<第二个小标题>，<第二个小标题>文案的创作必须遵循以下原则：
            请根据笔记内容和图片识别，用极简的文字概括这篇笔记和图片中景点的特色亮点，其中你可以参考以下词汇进行创作，这段文案字数控制6-8字符以内；

            特色亮点可能会出现的词汇不完全举例：非遗、古建、绝佳山水、祈福圣地、研学圣地、解压天堂、中国小瑞士、秘境竹筏游等等类型词汇

            接下来再去生成<第一个小标题>，<第一个小标题>文案的创作必须遵循以下原则：
            这部分文案创作公式有5种，分别为：
            ①<受众人群画像>+<痛点词>
            ②<受众人群画像>
            ③<痛点词>
            ④<受众人群画像>+ | +<痛点词>
            ⑤<痛点词>+ | +<受众人群画像>
            请你根据实际笔记内容，结合这部分文案创作公式，需要结合<受众人群画像>和<痛点词>时，必须根据<第二个小标题>的景点特征和所对应的完整笔记推文内容主旨，特征挑选对应<受众人群画像>和<痛点词>。

            我给你提供受众人群画像库和痛点词库如下：
            1、受众人群画像库：情侣党、亲子游、合家游、银发族、亲子研学、学生党、打工人、周边游、本地人、穷游党、性价比、户外人、美食党、出片
            2、痛点词库：3天2夜、必去、看了都哭了、不能错过、一定要来、问爆了、超全攻略、必打卡、强推、懒人攻略、必游榜、小众打卡、狂喜等等。

            你需要为每个请求至少生成{poster_num}个海报设计。请使用JSON格式输出结果，结构如下：

            ```json
            [
                {
                    "index": 1,
                    "main_title": "主标题内容",
                    "texts": ["第一个小标题", "第二个小标题"]
                },
                {
                    "index": 2,
                    "main_title": "主标题内容",
                    "texts": ["第一个小标题", "第二个小标题"]
                }
                // ... 更多海报
            ]
            ```

            确保生成的数量与用户要求的数量一致。只生成上述JSON格式内容，不要有其他任何额外内容。
            """
        
        user_content = f"""
        海报数量：{poster_num};
        景区介绍：{self.add_description};
        推文内容：{tweet_content};
        """
        
        # 最终响应内容
        full_response = ""
        
        # 创建临时客户端
        temp_client = self._create_temp_client()
        
        # 如果创建客户端失败，直接使用备用方案
        if temp_client is None:
            print("创建OpenAI客户端失败，使用备用方案生成内容")
            return 404
        else:
            # 添加重试机制
            for retry in range(max_retries):
                try:
                    print(f"尝试连接API (尝试 {retry+1}/{max_retries})...")
                    
                    # 计算退避时间（指数退避策略）：0, 2, 4, 8, 16...秒
                    if retry > 0:
                        backoff_time = min(2 ** (retry - 1) * 2, 30)  # 最大等待30秒
                        print(f"等待 {backoff_time} 秒后重试...")
                        time.sleep(backoff_time)
                    
                    # 设置超时时间随重试次数递增
                    timeout = 30 + (retry * 30)  # 30, 60, 90, ...秒
                    
                    chat_response = temp_client.chat.completions.create(
                        model=self.model_name,
                        messages=[
                            {"role": "system", "content": system_prompt},
                            {"role": "user", "content": user_content}
                        ],
                        stream=True,
                        temperature=self.temperature,
                        top_p=self.top_p,
                        presence_penalty=self.presence_penalty,
                        timeout=timeout  # 设置请求超时时间
                    )
                    
                    # 获取响应内容
                    for chunk in chat_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")  # 输出完成后换行
                    
                    # 成功获取响应，跳出重试循环
                    break
                    
                except Exception as e:
                    error_msg = str(e)
                    print(f"API连接错误 (尝试 {retry+1}/{max_retries}): {error_msg}")
                    
                    # 如果已经达到最大重试次数
                    if retry + 1 >= max_retries:
                        print("已达到最大重试次数，使用备用方案...")
                        # 生成备用内容（简单模板）
                        full_response = self._generate_fallback_content(poster_num)
                    else:
                        print(f"将在稍后重试，还剩 {max_retries - retry - 1} 次重试机会")
            
            # 关闭临时客户端
            self._close_client(temp_client)
        
        # 生成时间戳
        return full_response
    
    def save_result(self, full_response):
        """
        保存生成结果到文件
        
        参数:
            full_response: 生成的完整响应内容
        
        返回:
            结果文件路径
        """
        # 生成时间戳
        print(full_response)
        date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        full_response = self.split_content(full_response)
        # 创建结果文件路径
        result_path = os.path.join(self.output_dir, f"{date_time}.json")
        os.makedirs(os.path.dirname(result_path), exist_ok=True)
        
        # 保存结果到文件
        with open(result_path, "w", encoding="utf-8") as f:
            json.dump(full_response, f, ensure_ascii=False)
        
        print(f"结果已保存到: {result_path}")
        return result_path

    def run(self, info_directory, poster_num, tweet_content):
        """
        运行完整的海报生成流程
        
        参数:
            info_directory: 信息目录
            poster_num: 海报数量
            tweet_content: 推文内容
        
        返回:
            结果保存路径
        """
        ## 读取资料文件        
        self.load_infomation(info_directory)
        
        ## 生成海报内容
        full_response = self.generate_posters(poster_num, tweet_content)
        
        if self.output_dir:
            ## 保存结果
            return self.save_result(full_response)
        else:
            return full_response

    def set_temperature(self, temperature):
        self.temperature = temperature

    def set_top_p(self, top_p):
        self.top_p = top_p

    def set_presence_penalty(self, presence_penalty):
        self.presence_penalty = presence_penalty
    
    def set_model_para(self, temperature, top_p, presence_penalty):
        self.temperature = temperature
        self.top_p = top_p
        self.presence_penalty = presence_penalty

# def main():
#     # 配置参数
#     info_directory = [
#         "/root/autodl-tmp/sanming_img/相机/甘露寺/description.txt"
#         ]  # 信息目录
#     poster_num = 4  # 海报数量
    
#     # 推文内容
#     tweet_content = """<title>
#     🌿清明遛娃天花板！悬空古寺+非遗探秘
#     </title>
#     <content>
#     清明假期带娃哪里玩？泰宁甘露寺藏着明代建筑奇迹！一柱擎天的悬空阁楼+状元祈福传说，让孩子边玩边涨知识✨

#     🎒行程亮点：
#     ✅ 安全科普第一站：讲解"一柱插地"千年不倒的秘密，用乐高积木模型让孩子理解力学原理
#     ✅ 文化沉浸体验：穿汉服听"叶状元还愿建寺"故事，触摸3.38米粗的"状元柱"许愿
#     ✅ 自然探索路线：连接金湖栈道徒步，观察丹霞地貌与古建筑的巧妙融合

#     📌实用攻略：
#     📍位置：福建省三明市泰宁县金湖西路（导航搜"甘露岩寺"）
#     🕒最佳时段：上午10点前抵达避开人流，下午可衔接参观明清园（80元/人）
#     ⚠️注意事项：悬空栈道设置儿童安全绳租赁点，建议穿防滑鞋

#     💡亲子彩蛋：
#     1️⃣ 在"右鼓左钟"景观区玩声音实验，敲击不同岩石听回声差异
#     2️⃣ 领取任务卡完成"寻找建筑中的T形拱"小游戏，集章兑换非遗木雕书签
#     3️⃣ 结合清明节俗，用竹简模板书写祈福语系在古松枝头

#     周边推荐：游览完可直奔尚书第明代古民居群，对比不同时期建筑特色，晚餐推荐尝泰宁特色"灯盏糕"，亲子套票更划算！

#     清明带着孩子来场穿越850年的建筑探险，把课本里的力学知识变成触手可及的历史课堂！🌸

#     #清明节周边游 #亲子科普游 #福建遛娃 #泰宁旅行攻略 #建筑启蒙
#     </content>
#     """
    
#     # 创建海报生成器
#     generator = ContentGenerator()
    
#     # 运行生成流程
#     generator.run(info_directory, poster_num, tweet_content)


# if __name__ == "__main__":
#     main()