import gradio as gr
import google.generativeai as genai
import os
import json
import re

# 配置 API
api_key = os.getenv("MY_API_KEY")
genai.configure(api_key=api_key)

# --- 定义 CSS 样式 ---
custom_css = """
body { background-color: #0b0f14; }
.report-container { background: #111827; border: 1px solid #1f2937; border-radius: 12px; padding: 20px; color: #e5e7eb; font-family: 'Inter', sans-serif; }
.header-row { display: flex; justify-content: space-between; align-items: center; border-bottom: 1px solid #374151; padding-bottom: 15px; margin-bottom: 20px; }
.score-circle { background: radial-gradient(circle, #3b82f6, #1d4ed8); width: 60px; height: 60px; border-radius: 50%; display: flex; align-items: center; justify-content: center; font-size: 20px; font-weight: bold; color: white; box-shadow: 0 0 15px rgba(59, 130, 246, 0.5); }
.section-title { color: #60a5fa; font-size: 16px; font-weight: bold; margin: 15px 0 10px 0; display: flex; align-items: center; gap: 8px; }
.card { background: #1f2937; border-radius: 8px; padding: 12px; margin-bottom: 10px; border-left: 4px solid #3b82f6; }
.tag-container { display: flex; gap: 8px; margin-top: 10px; }
.tag { background: #374151; padding: 4px 10px; border-radius: 4px; font-size: 12px; color: #9ca3af; border: 1px solid #4b5563; }
.suggestion { background: #064e3b; border-left: 4px solid #10b981; padding: 10px; border-radius: 6px; margin-top: 8px; font-size: 14px; }
"""

def parse_and_format(ai_response):
    """提取 JSON 并渲染 HTML"""
    try:
        # 尝试从回复中找到 JSON 部分
        json_str = re.search(r'\{.*\}', ai_response, re.DOTALL).group()
        data = json.loads(json_str)
        
        html = f"""
        <div class="report-container">
            <div class="header-row">
                <div>
                    <h2 style="margin:0; color:white;">灯光分析报告</h2>
                    <div class="tag-container">
                        <span class="tag">UE5.4</span><span class="tag">Lumen</span><span class="tag">Cinema_Light</span>
                    </div>
                </div>
                <div class="score-circle">{data.get('score', '??')}</div>
            </div>
            
            <div class="section-title">🎨 色彩与氛围</div>
            <div class="card">{data.get('atmosphere', '')}</div>
            
            <div class="section-title">⚙️ 技术优化分析</div>
            <div class="card" style="border-left-color: #f59e0b;">
                <p style="margin:0; font-size:14px;">{data.get('technical', '')}</p>
            </div>
            
            <div class="section-title">🚀 灯光优化建议</div>
            {" ".join([f'<div class="suggestion">💡 {s}</div>' for s in data.get('suggestions', [])])}
        </div>
        """
        return html
    except:
        # 如果解析 JSON 失败，回退到 Markdown 显示
        return f"<div class='report-container'>{ai_response}</div>"

def analyze_lighting(image):
    if image is None: return "请上传图片"
    try:
        model = genai.GenerativeModel('gemini-3-flash-preview')
        # 强制 AI 输出 JSON 格式
        prompt = """
        你是一位顶级游戏灯光艺术总监。请分析这张 UE5 截图，并严格按以下 JSON 格式输出（不要有其他废话）：
        {
          "score": "评分数字",
          "atmosphere": "分析色温对比、氛围感",
          "technical": "分析曝光、死黑、体积雾、GI表现",
          "suggestions": ["建议1", "建议2", "建议3"]
        }
        请用专业中文。
        """
        response = model.generate_content([prompt, image])
        return parse_and_format(response.text)
    except Exception as e:
        return f"<div style='color:red;'>错误: {str(e)}</div>"

# --- 构建 UI ---
with gr.Blocks(css=custom_css) as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📷 场景画面上传")
            input_img = gr.Image(label=None, type="pil")
            btn = gr.Button("生成专业灯光分析报告", variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown("### 📊 分析结果预览")
            output_html = gr.HTML("<div style='color:#666;'>等待分析数据...</div>")
            
    btn.click(analyze_lighting, inputs=[input_img], outputs=[output_html])

demo.launch()