import gradio as gr
import os
import tempfile
import numpy as np
from PIL import Image
from src.predict import process_single_image
import sys
sys.path.insert(0, "./src")

# 自定义主题 - 炫彩现代化
custom_theme = gr.themes.Default(
    primary_hue="purple",
    secondary_hue="pink",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Poppins"), gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"]
).set(
    button_primary_background_fill="linear-gradient(45deg, #667eea 0%, #764ba2 100%)",
    button_primary_background_fill_hover="linear-gradient(45deg, #764ba2 0%, #667eea 100%)",
    button_primary_text_color="white",
    button_secondary_background_fill="linear-gradient(45deg, #f093fb 0%, #f5576c 100%)",
    button_secondary_background_fill_hover="linear-gradient(45deg, #f5576c 0%, #f093fb 100%)",
    button_secondary_text_color="white"
)

def get_example_images(folder_path="ffhq"):
    """获取示例图片列表"""
    return sorted([os.path.join(folder_path, f) for f in os.listdir(folder_path) 
            if f.lower().endswith(('.png', '.jpg', '.jpeg'))])

def safe_extract_prob(cls_probs):
    """安全地从cls_probs中提取概率值"""
    try:
        if cls_probs is None:
            return 0.0
        elif isinstance(cls_probs, (list, np.ndarray)) and len(cls_probs) > 0:
            return float(cls_probs[0])
        elif hasattr(cls_probs, '__getitem__'):
            return float(cls_probs[0])
        else:
            return float(cls_probs)
    except (TypeError, IndexError, ValueError) as e:
        print(f"Error extracting probability: {e}")
        return 0.0

# 创建主界面
with gr.Blocks(
    title="Loupe - AI图像伪造检测系统",
    theme=custom_theme,
    css="""
    /* 全局样式 */
    body {
        background: linear-gradient(-45deg, #ee7752, #e73c7e, #23a6d5, #23d5ab);
        background-size: 400% 400%;
        animation: gradientBG 15s ease infinite;
        min-height: 100vh;
    }
    
    @keyframes gradientBG {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    
    /* 主容器样式 */
    .gradio-container {
        background: rgba(255, 255, 255, 0.95);
        backdrop-filter: blur(10px);
        border-radius: 20px;
        box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
        margin: 20px;
        padding: 20px;
    }
    
    /* 标题样式 */
    .title-box {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 30px;
        border-radius: 15px;
        margin-bottom: 30px;
        box-shadow: 0 15px 35px rgba(102, 126, 234, 0.3);
        position: relative;
        overflow: hidden;
    }
    
    .title-box::before {
        content: '';
        position: absolute;
        top: -50%;
        left: -50%;
        width: 200%;
        height: 200%;
        background: linear-gradient(45deg, transparent, rgba(255, 255, 255, 0.1), transparent);
        animation: shine 3s infinite;
    }
    
    @keyframes shine {
        0% { transform: translateX(-100%) translateY(-100%) rotate(45deg); }
        100% { transform: translateX(100%) translateY(100%) rotate(45deg); }
    }
    
    .title-text {
        font-weight: 700;
        font-size: 32px;
        color: white;
        margin-bottom: 8px;
        text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.3);
        background: linear-gradient(45deg, #fff, #f0f8ff);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
    }
    
    .subtitle-text {
        color: rgba(255, 255, 255, 0.9);
        font-size: 18px;
        font-weight: 300;
        text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.2);
    }
    
    /* 输入和结果框样式 */
    .input-box, .result-box {
        background: linear-gradient(145deg, rgba(255, 255, 255, 0.9), rgba(248, 250, 252, 0.9));
        padding: 25px;
        border-radius: 15px;
        margin-bottom: 20px;
        border: 1px solid rgba(255, 255, 255, 0.3);
        box-shadow: 0 10px 30px rgba(0, 0, 0, 0.1);
        backdrop-filter: blur(10px);
        transition: all 0.3s ease;
    }
    
    .input-box:hover, .result-box:hover {
        transform: translateY(-5px);
        box-shadow: 0 20px 40px rgba(0, 0, 0, 0.15);
    }
    
    .input-title, .result-title {
        font-weight: 700;
        background: linear-gradient(45deg, #667eea, #764ba2);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        margin-bottom: 15px;
        font-size: 20px;
        text-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
    }
    
    /* 按钮样式 */
    .btn-primary {
        background: linear-gradient(45deg, #667eea 0%, #764ba2 100%);
        border: none;
        border-radius: 25px;
        padding: 12px 30px;
        font-weight: 600;
        text-transform: uppercase;
        letter-spacing: 1px;
        box-shadow: 0 10px 20px rgba(102, 126, 234, 0.3);
        transition: all 0.3s ease;
    }
    
    .btn-primary:hover {
        transform: translateY(-3px);
        box-shadow: 0 15px 30px rgba(102, 126, 234, 0.4);
        background: linear-gradient(45deg, #764ba2 0%, #667eea 100%);
    }
    
    .btn-secondary {
        background: linear-gradient(45deg, #f093fb 0%, #f5576c 100%);
        border: none;
        border-radius: 25px;
        padding: 10px 25px;
        font-weight: 600;
        box-shadow: 0 8px 16px rgba(240, 147, 251, 0.3);
        transition: all 0.3s ease;
    }
    
    .btn-secondary:hover {
        transform: translateY(-2px);
        box-shadow: 0 12px 24px rgba(240, 147, 251, 0.4);
    }
    
    /* 图片上传区域 */
    #upload_image {
        min-height: 350px;
        border: 3px dashed rgba(102, 126, 234, 0.3);
        border-radius: 15px;
        background: linear-gradient(45deg, rgba(102, 126, 234, 0.05), rgba(118, 75, 162, 0.05));
        transition: all 0.3s ease;
    }
    
    #upload_image:hover {
        border-color: rgba(102, 126, 234, 0.6);
        background: linear-gradient(45deg, rgba(102, 126, 234, 0.1), rgba(118, 75, 162, 0.1));
        transform: scale(1.02);
    }
    
    /* 概率显示 */
    #probability input {
        font-weight: bold;
        background: linear-gradient(45deg, #667eea, #764ba2);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        font-size: 1.2em;
    }
    
    #result_text input {
        font-size: 1.1em;
        font-weight: 600;
        background: linear-gradient(45deg, rgba(102, 126, 234, 0.1), rgba(118, 75, 162, 0.1));
        border-radius: 10px;
        border: 2px solid rgba(102, 126, 234, 0.2);
    }
    
    /* 画廊样式 */
    .gallery-item {
        border-radius: 12px !important;
        transition: all 0.3s ease;
        box-shadow: 0 5px 15px rgba(0, 0, 0, 0.1);
    }
    
    .gallery-item:hover {
        transform: scale(1.05);
        box-shadow: 0 10px 25px rgba(0, 0, 0, 0.2);
    }
    
    /* 示例按钮 */
    .example-btn {
        margin-top: 15px;
        width: 100%;
        background: linear-gradient(45deg, #23a6d5 0%, #23d5ab 100%);
        border-radius: 20px;
        font-weight: 600;
        box-shadow: 0 8px 16px rgba(35, 166, 213, 0.3);
        transition: all 0.3s ease;
    }
    
    .example-btn:hover {
        transform: translateY(-2px);
        box-shadow: 0 12px 24px rgba(35, 166, 213, 0.4);
    }
    
    /* Tab 样式 */
    .tab-nav button {
        border-radius: 15px 15px 0 0;
        background: linear-gradient(45deg, rgba(102, 126, 234, 0.8), rgba(118, 75, 162, 0.8));
        color: white;
        font-weight: 600;
        transition: all 0.3s ease;
    }
    
    .tab-nav button:hover {
        background: linear-gradient(45deg, rgba(118, 75, 162, 0.9), rgba(102, 126, 234, 0.9));
        transform: translateY(-2px);
    }
    
    /* 滑块样式 */
    .gr-slider input[type="range"] {
        background: linear-gradient(45deg, #667eea, #764ba2);
        border-radius: 10px;
    }
    
    /* 手风琴样式 */
    .gr-accordion {
        background: linear-gradient(145deg, rgba(255, 255, 255, 0.8), rgba(248, 250, 252, 0.8));
        border-radius: 15px;
        border: 1px solid rgba(102, 126, 234, 0.2);
        box-shadow: 0 5px 15px rgba(0, 0, 0, 0.1);
    }
    
    /* 炫彩加载动画 */
    @keyframes rainbow {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    
    .processing {
        background: linear-gradient(-45deg, #ee7752, #e73c7e, #23a6d5, #23d5ab);
        background-size: 400% 400%;
        animation: rainbow 2s ease infinite;
    }
    
    /* 响应式设计 */
    @media (max-width: 768px) {
        .title-text { font-size: 24px; }
        .subtitle-text { font-size: 16px; }
        .input-box, .result-box { padding: 20px; }
    }
    """
) as demo:
    
    # 标题部分 - 炫彩渐变设计
    with gr.Column(elem_classes="title-box"):
        gr.Markdown("""
        <div class="title-text">🔍 Loupe 图像伪造检测系统</div>
        <div class="subtitle-text">✨ 基于深度学习的图像伪造检测与定位技术</div>
        """)
    
    # 添加装饰性分割线
    gr.HTML("""
    <div style="height: 4px; background: linear-gradient(90deg, #667eea, #764ba2, #f093fb, #f5576c, #23a6d5, #23d5ab); 
                border-radius: 2px; margin: 20px 0; box-shadow: 0 2px 10px rgba(0,0,0,0.2);"></div>
    """)
    
    # 主界面组件
    with gr.Row(equal_height=True):
        with gr.Column(scale=1, min_width=300):
            # 输入图像区域 - 炫彩设计
            with gr.Column(elem_classes="input-box"):
                gr.Markdown("""<div class="input-title">🎨 输入图像</div>""")
                with gr.Tabs():
                    with gr.Tab("📤 上传图片", id="upload_tab"):
                        image_input = gr.Image(type="pil", label="", elem_id="upload_image")
                        upload_button = gr.Button("🚀 开始检测", variant="primary", size="lg", elem_classes="btn-primary")
                    
                    with gr.Tab("🖼️ 示例图片", id="example_tab"):
                        example_images = get_example_images()
                        example_gallery = gr.Gallery(
                            value=example_images,
                            label="",
                            columns=4,
                            rows=None,
                            height="auto",
                            object_fit="contain",
                            allow_preview=True,
                            selected_index=None
                        )
                        # 添加炫彩检测按钮
                        example_button = gr.Button(
                            "✨ 检测选中的示例图片", 
                            variant="primary", 
                            elem_classes="example-btn"
                        )
                        # 隐藏组件用于存储选中索引
                        selected_index = gr.Number(visible=False)
                
            with gr.Accordion("⚙️ 高级设置", open=False):
                threshold = gr.Slider(0, 1, value=0.5, step=0.01, label="🎯 检测敏感度")
                gr.HTML("""
                <div style="background: linear-gradient(45deg, rgba(102,126,234,0.1), rgba(118,75,162,0.1)); 
                           padding: 10px; border-radius: 8px; margin-top: 10px;">
                    <small style="color: #667eea; font-weight: 500;">💡 调整数值可改变检测的严格程度</small>
                </div>
                """)
        
        with gr.Column(scale=1.5, min_width=500):
            # 检测结果区域 - 炫彩设计
            with gr.Column(elem_classes="result-box"):
                gr.Markdown("""<div class="result-title">🎯 检测结果</div>""")
                with gr.Tabs():
                    with gr.Tab("🔍 检测效果", id="result_tab"):
                        output_image = gr.Image(label="伪造区域标记", interactive=False)
                    
                    with gr.Tab("⚖️ 对比视图", id="compare_tab"):
                        with gr.Row():
                            original_display = gr.Image(label="原始图像", interactive=False)
                            processed_display = gr.Image(label="检测结果", interactive=False)
                
                with gr.Group():
                    with gr.Row():
                        fake_prob = gr.Number(label="🎲 伪造概率", precision=2, elem_id="probability")
                        result_text = gr.Textbox(label="📝 检测结论", interactive=False, elem_id="result_text")
                
                with gr.Row():
                    save_button = gr.Button("💾 保存结果", variant="secondary", elem_classes="btn-secondary")
                    clear_button = gr.Button("🧹 清除", variant="secondary", elem_classes="btn-secondary")
    
    # 关于部分 - 炫彩设计
    with gr.Accordion("🌟 关于系统", open=False):
        gr.HTML("""
        <div style="background: linear-gradient(135deg, rgba(102,126,234,0.1), rgba(118,75,162,0.1), rgba(240,147,251,0.1)); 
                   padding: 20px; border-radius: 15px; border: 1px solid rgba(102,126,234,0.2);">
            <h3 style="background: linear-gradient(45deg, #667eea, #764ba2); -webkit-background-clip: text; 
                      -webkit-text-fill-color: transparent; margin-bottom: 15px;">
                ✨ Loupe 伪造图像检测系统
            </h3>
            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 15px;">
                <div style="background: rgba(102,126,234,0.1); padding: 15px; border-radius: 10px;">
                    <strong style="color: #667eea;">🚀 技术</strong><br>
                    基于深度学习的图像伪造检测与定位
                </div>
                <div style="background: rgba(118,75,162,0.1); padding: 15px; border-radius: 10px;">
                    <strong style="color: #764ba2;">⭐ 特点</strong><br>
                    高精度、实时处理、可解释性强
                </div>
                <div style="background: rgba(240,147,251,0.1); padding: 15px; border-radius: 10px;">
                    <strong style="color: #f093fb;">📱 版本</strong><br>
                    v2.0.0 炫彩版
                </div>
                <div style="background: rgba(245,87,108,0.1); padding: 15px; border-radius: 10px;">
                    <strong style="color: #f5576c;">👥 开发者</strong><br>
                    EVOL Lab (jyc, xxw)
                </div>
            </div>
            <div style="margin-top: 20px; padding: 15px; background: linear-gradient(45deg, rgba(35,166,213,0.1), rgba(35,213,171,0.1)); 
                       border-radius: 10px; border-left: 4px solid #23a6d5;">
                <strong style="color: #23a6d5;">💡 系统介绍</strong><br>
                本系统可检测多种图像篡改痕迹，包括复制-移动、拼接、擦除等操作。采用最新的深度学习算法，提供高精度的检测结果和直观的可视化分析。
            </div>
        </div>
        """)
    
    # 页脚 - 炫彩设计
    gr.HTML("""
    <div style="margin-top: 40px; padding: 20px; text-align: center; 
               background: linear-gradient(135deg, rgba(102,126,234,0.1), rgba(118,75,162,0.1)); 
               border-radius: 15px; border-top: 2px solid rgba(102,126,234,0.3);">
        <div style="background: linear-gradient(45deg, #667eea, #764ba2); -webkit-background-clip: text; 
                   -webkit-text-fill-color: transparent; font-weight: 600; margin-bottom: 10px;">
            ✨ 感谢使用 Loupe 图像伪造检测系统 ✨
        </div>
        <div style="color: #64748b; font-size: 14px;">
            © 2025 EVOL Lab | 让AI守护图像真实性 🛡️
        </div>
        <div style="margin-top: 10px;">
            <span style="background: linear-gradient(45deg, #f093fb, #f5576c); -webkit-background-clip: text; 
                        -webkit-text-fill-color: transparent; font-weight: 500;">
                🌟 科技点亮未来，智能守护真实 🌟
            </span>
        </div>
    </div>
    """)

    def process_image(image, threshold_value):
        """处理上传的图像"""
        if image is None:
            return {
                output_image: None, 
                fake_prob: 0.0, 
                result_text: "❌ 请上传有效图像"
            }
        
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
            image_path = tmp_file.name
            image.save(image_path)

        try:
            processed_img, cls_probs = process_single_image(image_path)
            prob = safe_extract_prob(cls_probs)
            
            # 根据概率生成炫彩结论
            if prob > threshold_value + 0.2:
                conclusion = "🚨 高度疑似伪造"
                emoji = "🔴"
            elif prob > threshold_value:
                conclusion = "⚠️ 可能伪造"
                emoji = "🟡"
            else:
                conclusion = "✅ 未检测到伪造"
                emoji = "🟢"
            
            return {
                output_image: processed_img,
                original_display: image,
                processed_display: processed_img,
                fake_prob: prob,
                result_text: f"{emoji} {conclusion} (概率: {prob:.2f})"
            }
        except Exception as e:
            print(f"Error in processing: {e}")
            return {
                output_image: None,
                original_display: image,
                processed_display: None,
                fake_prob: 0.0,
                result_text: f"❌ 处理错误: {str(e)}"
            }
        finally:
            if os.path.exists(image_path):
                os.unlink(image_path)

    def process_example(example_data, selected_idx, threshold_value):
        """处理示例图像"""
        if not example_data or selected_idx is None:
            return {
                image_input: None,
                output_image: None,
                original_display: None,
                processed_display: None,
                fake_prob: 0.0,
                result_text: "⚠️ 请先选择示例图片",
                threshold: threshold_value
            }
            
        try:
            selected_idx = int(selected_idx)
            image_info = example_data[selected_idx]
            
            # 处理不同的数据格式
            if isinstance(image_info, (tuple, list)):
                image_path = image_info[0]  # (path, caption)格式
            elif isinstance(image_info, dict):
                image_path = image_info.get("name", image_info.get("path"))
            else:
                image_path = image_info
                
            print(f"Processing selected image (index {selected_idx}): {image_path}")  # 调试日志
            
            # 处理图像
            processed_img, cls_probs = process_single_image(image_path)
            prob = safe_extract_prob(cls_probs)
            original_img = Image.open(image_path)
            
            # 根据概率生成炫彩结论
            if prob > threshold_value + 0.2:
                conclusion = "🚨 高度疑似伪造"
                emoji = "🔴"
            elif prob > threshold_value:
                conclusion = "⚠️ 可能伪造"
                emoji = "🟡"
            else:
                conclusion = "✅ 未检测到伪造"
                emoji = "🟢"
            
            return {
                image_input: original_img,
                output_image: processed_img,
                original_display: original_img,
                processed_display: processed_img,
                fake_prob: prob,
                result_text: f"{emoji} {conclusion} (概率: {prob:.2f})",
                threshold: threshold_value
            }
        except Exception as e:
            print(f"Error in processing example: {e}")
            return {
                image_input: None,
                output_image: None,
                original_display: None,
                processed_display: None,
                fake_prob: 0.0,
                result_text: f"❌ 示例处理错误: {str(e)}",
                threshold: threshold_value
            }

    def clear_all():
        """清除所有输入输出"""
        return {
            image_input: None,
            output_image: None,
            original_display: None,
            processed_display: None,
            fake_prob: 0.0,
            result_text: "🧹 已清除所有数据"
        }

    def update_selected_index(evt: gr.SelectData):
        """更新选中的图片索引"""
        return evt.index

    # 交互逻辑
    upload_button.click(
        process_image,
        [image_input, threshold],
        [output_image, original_display, processed_display, fake_prob, result_text]
    )
    
    # 示例图片选择事件
    example_gallery.select(
        update_selected_index,
        None,
        selected_index
    )
    
    # 示例图片检测按钮点击事件
    example_button.click(
        process_example,
        [example_gallery, selected_index, threshold],
        [image_input, output_image, original_display, processed_display, fake_prob, result_text, threshold]
    )
    
    save_button.click(
        lambda img: (img.save("result.jpg"), "💾 结果已保存为 result.jpg")[1] if img else "❌ 没有图像可保存",
        [output_image],
        None,
        api_name="save_result"
    )
    
    clear_button.click(
        clear_all,
        [],
        [image_input, output_image, original_display, processed_display, fake_prob, result_text]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7864,
        favicon_path="./favicon.ico" if os.path.exists("./favicon.ico") else None
    )