"""
@file        generate_layout.py
@brief       Vue组件代码生成模块
             基于DeepSeek API生成响应式Vue 3组件布局代码，支持预定义模板

@author      王秀强 (2310460@mail.nankai.edu.cn)
@date        2025.5.19
@version     v0.5.2

@details
本文件主要实现：
 - DeepSeek API调用封装和错误处理
 - Vue 3 Composition API组件代码生成
 - 海报布局的动态排版和样式生成
 - 预定义Vue模板系统
 - 增强Vue代码生成逻辑
 - 代码清理和格式化处理

@note
- 需要配置DEEPSEEK_API_KEY环境变量
- 支持流式和非流式响应模式
- 生成的Vue代码包含完整的template、script和style部分
- 具备指数退避重试机制处理API限流
- 支持基于图片类型的预定义模板

@usage
# 生成Vue组件代码
vue_code = generate_vue_code_enhanced("生成端午节海报Vue组件", parse_imglist, suggestions)
save_code(vue_code, "../outputs/poster.vue")

@copyright
(c) 2025 砚生项目组
"""

import os
import yaml
from openai import OpenAI
from dotenv import load_dotenv
import time
from colorama import init, Fore, Back, Style
from typing import List, Dict, Optional

# 初始化colorama
init(autoreset=True)

# === Config LLM call ===
load_dotenv()
deepseek_url = 'https://api.deepseek.com/v1' # set to be compatible with the OpenAI API
deepseek_api = os.getenv("DEEPSEEK_API_KEY")
if not deepseek_api:
    raise ValueError("DEEPSEEK_API_KEY not set!") 

def call_deepseek(
        messages=None,
        system_prompt="你是一个擅长前端开发的AI，专注于生成Vue.js代码。",
        prompt=None,
        model='deepseek-chat',
        temperature=0.6,
        max_tokens=2000,
        stream=False,
        max_retries=3,
):
    """调用 DeepSeek API,支持多轮对话和流式/非流式响应"""
    client = OpenAI(api_key=deepseek_api, base_url=deepseek_url)

    # 参数验证
    if messages is None:
        if not prompt:
            raise ValueError("必须提供 message 或 prompt 参数")
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]
    elif not isinstance(messages, list) or not messages:
        raise ValueError("message 参数必须是非空列表")

    # 模型验证
    models = ["deepseek-chat", "deepseek-reasoner"]
    if model not in models:
        raise ValueError(f"无效的模型名称: {model}，可用模型: {models}")

    # 调用 API
    for attempt in range(max_retries):
        try:
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=stream
            )

            if stream:
                # 流式响应处理
                def stream_generator():
                    for chunk in response:
                        if hasattr(chunk, 'choices') and chunk.choices[0].delta.content is not None:
                            yield chunk.choices[0].delta.content
                return stream_generator(), None
            else:
                # 非流式响应
                if hasattr(response, 'choices') and response.choices:
                    content = response.choices[0].message.content
                    usage = getattr(response, 'usage', None)
                    return content, usage
                return "", None

        except Exception as e:
            if hasattr(e, 'status_code') and getattr(e, 'status_code') == 429:
                time.sleep(2 ** attempt)  # 指数退避
            elif attempt == max_retries - 1:
                raise
            else:
                time.sleep(1)
    
    raise Exception("达到最大重试次数，API 调用失败")


def load_vue_templates() -> Dict:
    """加载预定义的Vue模板配置"""
    from utils import load_vue_templates as utils_load_templates
    return utils_load_templates()


def get_template_by_images(parse_imglist: List[Dict]) -> Optional[str]:
    """根据图片列表选择合适的预定义模板"""
    templates = load_vue_templates()
    
    if not templates or not parse_imglist:
        return None
    
    # 使用utils中的函数确定模板类型
    from utils import determine_template_type as utils_determine_type
    template_choice = utils_determine_type(parse_imglist)
    
    if template_choice in templates:
        return templates[template_choice]['template']
    
    # 备用：检查特定文件名
    for img_info in parse_imglist:
        img_name = img_info.get('picture_name', '')
        if img_name in templates:
            return templates[img_name]['template']
    
    return None


def determine_template_type(parse_imglist: List[Dict]) -> str:
    """根据图片信息确定模板类型"""
    from utils import determine_template_type as utils_determine_type
    return utils_determine_type(parse_imglist)


def generate_layout_prompt(user_input_analysis_result: Dict, parse_imglist: List[Dict], suggestions: Optional[Dict] = None) -> str:
    """根据用户分析结果动态生成Vue布局提示"""
    width = user_input_analysis_result.get("width", 1080)
    height = user_input_analysis_result.get("height", 1920)
    theme = user_input_analysis_result.get("main_theme", "活动海报")
    
    # 从用户分析中提取更多信息
    style = user_input_analysis_result.get("style", "现代简约")
    color_scheme = user_input_analysis_result.get("color_scheme", "默认配色")
    layout_type = user_input_analysis_result.get("layout_type", "标准布局")
    target_audience = user_input_analysis_result.get("target_audience", "一般用户")

    # 构造图片信息
    images_info = []
    for i, img in enumerate(parse_imglist):
        img_desc = f"图片{i+1}: {img['picture_name']} - {img['picture_description']}"
        images_info.append(img_desc)
    
    images_text = "\n".join(images_info) if images_info else "无特定图片资源"

    # 构造文案信息
    content_parts = []
    if suggestions:
        try:
            if 'layer6_title_content' in suggestions:
                title = suggestions['layer6_title_content'].get('content', theme)
                content_parts.append(f"主标题: {title}")

            if 'layer7_subtitle_content' in suggestions:
                subtitle = suggestions['layer7_subtitle_content'].get('content', '精彩活动，敬请参与')
                content_parts.append(f"副标题: {subtitle}")

            if 'layer5_logo_content' in suggestions:
                logo = suggestions['layer5_logo_content'].get('text', '')
                if logo:
                    content_parts.append(f"Logo文字: {logo}")
        except Exception:
            content_parts = [f"主标题: {theme}", "副标题: 精彩活动，敬请参与"]
    else:
        content_parts = [f"主标题: {theme}", "副标题: 精彩活动，敬请参与"]

    content_info = "\n".join([f"- {part}" for part in content_parts])

    # 根据主题类型调整布局描述
    layout_requirements = []
    if "节日" in theme or "festival" in theme.lower():
        layout_requirements = [
            "背景图层: 使用节日主题背景，营造氛围",
            "标题区域: 突出节日名称，位于画布上方，使用醒目字体",
            "装饰元素: 添加节日相关装饰图案，分布在画面周围",
            "日期信息: 在适当位置显示日期",
            "Logo区域: 位于底部，展示主办方信息"
        ]
    elif "校园" in theme or "大学" in theme or "学术" in theme:
        layout_requirements = [
            "背景图层: 使用校园或学术主题背景",
            "标题区域: 学术风格标题，位于画布上方1/3处",
            "内容区域: 展示学术内容或校园风光，居中布局",
            "信息栏: 显示相关学术信息或活动详情",
            "Logo区域: 展示学校标识，位于底部"
        ]
    else:
        layout_requirements = [
            "背景图层: 使用主题相关背景图片，占据整个组件区域",
            "主标题: 位于画布上方1/3处，居中显示，突出主题",
            "副标题: 位于主标题下方，居中显示，补充说明",
            "内容区域: 合理布局图片和信息，保持视觉平衡",
            "Logo区域: 位于底部，居中显示主办方信息"
        ]

    layout_text = "\n".join([f"{i+1}. {req}" for i, req in enumerate(layout_requirements)])

    system_prompt = f"你是一个专业的前端设计师，擅长根据用户需求生成{style}风格的Vue.js组件。请根据{target_audience}的审美偏好，生成完整的Vue组件代码。"

    prompt = f"""
请为"{theme}"主题设计一个Vue 3组件，要求如下：

【基本信息】
- 组件尺寸: {width}x{height}px
- 设计风格: {style}
- 配色方案: {color_scheme}
- 布局类型: {layout_type}

【图片资源】
{images_text}

【文案内容】
{content_info}

【布局要求】
{layout_text}

【技术要求】
- 使用Vue 3 Composition API
- 使用absolute或flex定位进行精确布局
- 包含完整的template、script setup和style部分
- 确保响应式设计和良好的视觉效果
- 图片使用相对路径引用
- 样式要体现{style}的设计理念

请生成完整可用的Vue组件代码，代码要简洁优雅，符合{style}风格特点。
"""

    try:
        result, _ = call_deepseek(prompt=prompt, system_prompt=system_prompt, temperature=0.4)
        return result if isinstance(result, str) else ""
    except Exception:
        return ""


def fill_template_content(template: str, suggestions: Dict) -> str:
    """填充模板中的动态内容"""
    from utils import extract_content_from_suggestions
    
    try:
        title_content, subtitle_content, logo_content = extract_content_from_suggestions(suggestions)
        
        filled_template = template
        
        # 根据模板类型进行替换
        if 'illustration-theme' in template:
            filled_template = filled_template.replace("'见微知著记录南开'", f"'{title_content}'")
            filled_template = filled_template.replace("'南开大学融媒体中心'", f"'{subtitle_content}'")
            if logo_content:
                filled_template = filled_template.replace("'主办方'", f"'{logo_content}'")
        
        elif 'festival-theme' in template:
            filled_template = filled_template.replace("'南开大学'", "'南开大学'")
            filled_template = filled_template.replace("'中秋节'", f"'{title_content}'")
            filled_template = filled_template.replace("'月圆人团圆'", f"'{subtitle_content}'")
            from datetime import datetime
            current_date = datetime.now().strftime('%Y年%m月%d日')
            filled_template = filled_template.replace("'2025年10月6日'", f"'{current_date}'")
        
        else:
            # 通用替换
            filled_template = filled_template.replace('{{ title_content }}', title_content)
            filled_template = filled_template.replace('{{ subtitle_content }}', subtitle_content)
            filled_template = filled_template.replace('{{ logo_content }}', logo_content)

        return filled_template

    except Exception:
        return template


def generate_vue_code_enhanced(
    user_input_analysis_result: Dict,
    parse_imglist: List[Dict],
    suggestions: Optional[Dict] = None,
    prompt: Optional[str] = None
) -> str:
    """增强的Vue代码生成函数，优先使用预定义模板"""
    
    # 尝试使用预定义模板
    template_code = get_template_by_images(parse_imglist)
    
    if template_code and suggestions:
        vue_code = fill_template_content(template_code, suggestions)
        return vue_code

    # 使用AI生成
    if not prompt:
        prompt = generate_layout_prompt(user_input_analysis_result, parse_imglist, suggestions)

    return generate_vue_code(prompt)


def generate_vue_code(prompt):
    """Vue代码生成函数"""
    system_prompt = (
        "你是一个擅长前端开发的AI，专注于生成Vue.js代码。"
        "请生成完整的Vue 3组件，包含template、script setup和style部分。"
        "确保代码结构清晰，语法正确，可以直接使用。"
        "不要包含任何解释文字，只返回纯代码。"
    )

    try:
        result, usage = call_deepseek(prompt=prompt, system_prompt=system_prompt, temperature=0.4)
        
        # 使用utils中的代码清理函数
        from utils import clean_vue_code
        if result and isinstance(result, str):
            cleaned_result = clean_vue_code(result)
            if cleaned_result:
                return cleaned_result
        
        # 如果API失败，返回简单的错误提示
        return "<template><div>Vue代码生成失败</div></template>"
            
    except Exception:
        return "<template><div>Vue代码生成失败</div></template>"

def save_code(code, file_path="../outputs/generated_code.vue"):
    """保存代码到文件"""
    try:
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(code)
        
        if os.path.exists(file_path):
            file_size = os.path.getsize(file_path)
    except Exception:
        pass

if __name__ == "__main__":
    # 简单测试
    test_analysis = {"width": 1080, "height": 1920, "main_theme": "端午节海报"}
    test_imglist = [
        {"picture_name": "background", "picture_description": "背景图片"},
        {"picture_name": "lotus", "picture_description": "荷花装饰"}
    ]
    test_suggestions = {
        "layer6_title_content": {"content": "端午安康"},
        "layer7_subtitle_content": {"content": "粽叶飘香，龙舟竞渡"},
        "layer5_logo_content": {"text": "主办方"}
    }

    vue_code = generate_vue_code_enhanced(test_analysis, test_imglist, test_suggestions)
    save_code(vue_code)