import os
from dotenv import load_dotenv
from generate_layout import call_deepseek
import json
from colorama import init, Fore, Back, Style

# 初始化colorama
init(autoreset=True)

# 加载环境变量
load_dotenv()

def llm_user_analysis(user_input):
    """
    使用DeepSeek动态分析用户输入，提取海报设计参数
    """
    if not user_input:
        user_input = "端午节海报，包含背景、活动亮点和图标"
    
    print(f"{Fore.CYAN}🔍 正在分析用户输入: {Style.BRIGHT}{user_input}{Style.RESET_ALL}")
    
    # 构建分析提示词
    analysis_prompt = f"""
    请分析以下用户输入的海报需求，提取关键信息并以JSON格式返回：
    
    用户输入：{user_input}
    
    请严格按照以下JSON格式返回：
    {{
        "analyzed_prompt": "原始用户输入",
        "keywords": ["提取的关键词1", "关键词2", "关键词3"],
        "width": 1080,
        "height": 1920,
        "batch_size": 2,
        "poster_type": "海报类型（如：节日海报、活动海报、产品海报等）",
        "main_theme": "主要主题",
        "style_preference": "风格偏好（如：现代、传统、简约等）",
        "color_preference": "颜色偏好（如：暖色调、冷色调、单色等）"
    }}
    
    注意：
    1. keywords应该包含3-5个最重要的关键词
    2. 根据用户输入推断合适的海报类型
    3. 尺寸默认为1080x1920，除非用户明确指定
    4. batch_size根据需求调整，通常为1-4
    5. 分析用户的风格和颜色偏好
    """
    
    system_prompt = "你是一个专业的设计需求分析师，擅长从用户描述中提取海报设计的关键参数。请严格按照JSON格式返回结果，确保输出的JSON格式正确且完整。"
    
    try:
        result, _ = call_deepseek(prompt=analysis_prompt, system_prompt=system_prompt, temperature=0.3)
        
        # 解析JSON
        json_str = result.strip()
        if "```json" in json_str:
            json_str = json_str.split("```json")[1].split("```")[0].strip()
        elif json_str.startswith("```") and json_str.endswith("```"):
            json_str = json_str[3:-3].strip()
            
        analysis_result = json.loads(json_str)
        
        print(f"{Fore.GREEN}✅ 分析完成! {Style.RESET_ALL}")
        print(f"{Fore.YELLOW}📊 主题: {Style.BRIGHT}{analysis_result.get('main_theme', '未知')}{Style.RESET_ALL}")
        print(f"{Fore.YELLOW}🎨 风格: {analysis_result.get('style_preference', '未设置')}")
        print(f"{Fore.YELLOW}🔖 关键词: {', '.join(analysis_result.get('keywords', []))}{Style.RESET_ALL}")
        
        return analysis_result
        
    except Exception as e:
        print(f"{Fore.RED}❌ 分析失败: {str(e)}{Style.RESET_ALL}")
        # 返回默认值
        return {
            "analyzed_prompt": user_input,
            "keywords": ["海报", "设计", "活动"],
            "width": 1080,
            "height": 1920,
            "batch_size": 2,
            "poster_type": "通用海报",
            "main_theme": "默认主题", 
            "style_preference": "现代",
            "color_preference": "暖色调"
        }

if __name__ == "__main__":
    # 测试函数
    test_input = "春节海报，红色背景，包含灯笼和烟花，现代风格"
    result = llm_user_analysis(test_input)
    print(f"\n{Fore.GREEN}测试结果:")
    print(json.dumps(result, indent=2, ensure_ascii=False))