# -*- coding: utf-8 -*-
import os
import sys
import openai
from datetime import datetime

sys.path.append(os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0] + 'AlphaGenerator'))
PROJECT_PATH = os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0] + 'AlphaGenerator')

NEW_STORY_PROMPT = os.path.join(PROJECT_PATH, 'new_story_prompt')

TEMPERATURE = 0.1

USE_AI = 1

SILICONFLOW_API_KEY = "sk-pvdiisdowmuwkrpnxsrlhxaovicqibmlljwrwwvbbdjaitdl"
SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
MODELS = [
    'Pro/deepseek-ai/DeepSeek-V3.1-Terminus',
    # 'deepseek-ai/DeepSeek-V3.2-Exp',
    # 'Qwen/Qwen3-VL-235B-A22B-Instruct',
    # 'MiniMaxAI/MiniMax-M2',
    # 'zai-org/GLM-4.6',
    # 'inclusionAI/Ring-flash-2.0',
    # 'zai-org/GLM-4.6',
    # 'inclusionAI/Ling-flash-2.0',
    # 'inclusionAI/Ring-flash-2.0',
]

def load_prompt():
    """读取多份financial_logic_X.txt文件（X为0-9的数字）"""
    result = ""
    try:
        for i in range(10):  # 0-9，因为您说数字不会超过1位数
            filename = f"financial_logic_{i}.txt"
            filepath = os.path.join(NEW_STORY_PROMPT, filename)
            
            if os.path.exists(filepath):
                with open(filepath, 'r', encoding='utf-8') as f:
                    content = f.read()
                    result += content + "\n\n"  # 每个文件之间加空行分隔
                    print(f"已加载文件: {filename}")
            else:
                # 如果文件不存在，继续检查下一个
                continue
        
        return result.rstrip("\n\n")  # 移除末尾的空行
        
    except Exception as e:
        print(f"读取提示词文件出错: {e}")
        return ""

def end_prompt():
    """读取结束提示词文件"""
    try:
        # 假设结束提示词文件名为 end_prompt.txt
        end_prompt_file = os.path.join(NEW_STORY_PROMPT, 'end_prompt.txt')
        if os.path.exists(end_prompt_file):
            with open(end_prompt_file, 'r', encoding='utf-8') as f:
                result = f.read()
        else:
            result = ""
        return result
    except Exception as e:
        print(f"读取结束提示词文件出错: {e}")
        return ""

def call_siliconflow(prompt, model):
    try:
        client = openai.OpenAI(
            api_key=SILICONFLOW_API_KEY,
            base_url=SILICONFLOW_BASE_URL
        )

        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "你是一个专业的量化投资专家，编写创新的金融逻辑/策略"},
                {"role": "user", "content": prompt}
                ],
            temperature=TEMPERATURE
        )

        return response.choices[0].message.content

    except openai.AuthenticationError:
        print("API密钥错误")
    except openai.RateLimitError:
        print("调用频率限制")
    except openai.APIError as e:
        print(f"API错误: {e}")
    except Exception as e:
        print(f"其他错误: {e}")
    exit(1)

def save_finally_prompt(finally_prompt):
    """保存最终生成的提示词"""
    try:
        # 获取当前日期和时间
        now = datetime.now()
        year = str(now.year)  # 年份，例如 2026
        month = now.strftime("%m")  # 月份，例如 01
        day = now.strftime("%d")  # 日，例如 10
        
        # 构建目录路径
        base_dir = "manual_new_story_prompt"
        year_dir = os.path.join(base_dir, year)
        month_dir = os.path.join(year_dir, month)
        day_dir = os.path.join(month_dir, day)
        
        # 创建目录（如果不存在）
        for directory in [base_dir, year_dir, month_dir, day_dir]:
            if not os.path.exists(directory):
                os.makedirs(directory)
                print(f"已创建目录: {directory}")
        
        # 生成文件名（包含时间戳，避免重复）
        timestamp = now.strftime("%Y%m%d_%H%M%S")
        filename = f"prompt_{timestamp}.txt"
        filepath = os.path.join(day_dir, filename)
        
        # 保存提示词到文件
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(finally_prompt)
        
        print(f"提示词已保存到: {filepath}")
        return filepath
        
    except Exception as e:
        print(f"保存提示词时出错: {e}")
        return None

def save_ai_result(ai_result, model_name):
    """保存AI生成的结果"""
    try:
        # 获取当前日期和时间
        now = datetime.now()
        year = str(now.year)  # 年份，例如 2026
        month = now.strftime("%m")  # 月份，例如 01
        day = now.strftime("%d")  # 日，例如 10
        
        # 构建目录路径
        base_dir = "ai_generated_new_story"
        year_dir = os.path.join(base_dir, year)
        month_dir = os.path.join(year_dir, month)
        day_dir = os.path.join(month_dir, day)
        
        # 创建目录（如果不存在）
        for directory in [base_dir, year_dir, month_dir, day_dir]:
            if not os.path.exists(directory):
                os.makedirs(directory)
                print(f"已创建目录: {directory}")
        
        # 清理模型名称中的特殊字符，用于文件名
        safe_model_name = model_name.replace('/', '_').replace('\\', '_').replace(':', '_')
        
        # 生成文件名（包含时间戳和模型名）
        timestamp = now.strftime("%Y%m%d_%H%M%S")
        filename = f"result_{safe_model_name}_{timestamp}.txt"
        filepath = os.path.join(day_dir, filename)
        
        # 保存结果到文件
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(ai_result)
        
        print(f"AI结果已保存到: {filepath}")
        return filepath
        
    except Exception as e:
        print(f"保存AI结果时出错: {e}")
        return None

def main():
    finally_prompt = ""
    
    prompt = load_prompt()

    if prompt:
        finally_prompt += prompt

    end_prompt_text = end_prompt()

    if end_prompt_text:
        finally_prompt += end_prompt_text

    save_finally_prompt(finally_prompt)

    if USE_AI:
        for model in MODELS:
            result = call_siliconflow(finally_prompt, model)
            save_ai_result(result, model)
    

if __name__ == "__main__":
    main()