import gradio as gr
import os

# 配置视频文件路径和场景信息
SCENES = {
    "basketball": {
        "ego_query": "videos/basketball/basketball_ego_original.mp4",
        "exo_target": "videos/basketball/basketball_exo_original.mp4",
        "objects": ["ball", "hoop"]
    },
    "cooking": {
        "ego_query": "videos/cooking/kitchen_ego_original.mp4",
        "exo_target": "videos/cooking/kitchen_exo_original.mp4",
        "objects": ["food", "knife_block"]
    },
    "health": {
        "ego_query": "videos/health/health_ego_original.mp4",
        "exo_target": "videos/health/health_exo_original.mp4",
        "objects": ["instruction_manual", "transport_box", "test_strip"]
    }
}

def update_scene(scene):
    """更新场景时的回调函数"""
    ego_video = SCENES[scene]["ego_query"]
    exo_video = SCENES[scene]["exo_target"]
    objects = SCENES[scene]["objects"]
    
    # 更新物体选择下拉菜单
    object_dropdown = gr.Dropdown(
        choices=objects,
        value=None,
        label="Select Object",
        interactive=True
    )
    
    return ego_video, exo_video, object_dropdown

def update_videos_with_object(scene, selected_object):
    """选择物体时更新视频"""
    if selected_object is None:
        # 如果没有选择物体，返回原始视频
        ego_video = SCENES[scene]["ego_query"]
        exo_video = SCENES[scene]["exo_target"]
    else:
        # 这里假设有mask视频的命名规则
        # 例如：ball_ego_query_mask.mp4, ball_exo_target_mask.mp4
        ego_video = f"videos/{scene}/{selected_object}_ego_query_mask.mp4"
        exo_video = f"videos/{scene}/{selected_object}_exo_target_mask.mp4"
        
        # 如果mask视频不存在，回退到原始视频
        if not os.path.exists(ego_video):
            ego_video = SCENES[scene]["ego_query"]
        if not os.path.exists(exo_video):
            exo_video = SCENES[scene]["exo_target"]
    
    # 返回视频路径
    return ego_video, exo_video

def sync_video_playback():
    """JavaScript代码：检测两个视频都加载完毕后同步播放"""
    return """
    () => {
        const egoVideo = document.querySelector('#ego_video video');
        const exoVideo = document.querySelector('#exo_video video');
        
        if (!egoVideo || !exoVideo) {
            setTimeout(() => {
                const egoVideo = document.querySelector('#ego_video video');
                const exoVideo = document.querySelector('#exo_video video');
                if (egoVideo && exoVideo) {
                    let egoLoaded = false;
                    let exoLoaded = false;
                    
                    const checkBothLoaded = () => {
                        if (egoLoaded && exoLoaded) {
                            egoVideo.currentTime = 0;
                            exoVideo.currentTime = 0;
                            Promise.all([egoVideo.play(), exoVideo.play()]).then(() => {
                                console.log('Both videos started playing synchronously');
                            }).catch(error => console.error('Error playing videos:', error));
                        }
                    };
                    
                    const onEgoLoaded = () => { egoLoaded = true; checkBothLoaded(); };
                    const onExoLoaded = () => { exoLoaded = true; checkBothLoaded(); };
                    
                    egoVideo.addEventListener('loadeddata', onEgoLoaded, { once: true });
                    exoVideo.addEventListener('loadeddata', onExoLoaded, { once: true });
                    
                    if (egoVideo.readyState >= 2) onEgoLoaded();
                    if (exoVideo.readyState >= 2) onExoLoaded();
                }
            }, 100);
            return;
        }
        
        let egoLoaded = false;
        let exoLoaded = false;
        
        const checkBothLoaded = () => {
            if (egoLoaded && exoLoaded) {
                egoVideo.currentTime = 0;
                exoVideo.currentTime = 0;
                Promise.all([egoVideo.play(), exoVideo.play()]).then(() => {
                    console.log('Both videos started playing synchronously');
                }).catch(error => console.error('Error playing videos:', error));
            }
        };
        
        const onEgoLoaded = () => { egoLoaded = true; checkBothLoaded(); };
        const onExoLoaded = () => { exoLoaded = true; checkBothLoaded(); };
        
        egoVideo.addEventListener('loadeddata', onEgoLoaded, { once: true });
        exoVideo.addEventListener('loadeddata', onExoLoaded, { once: true });
        
        if (egoVideo.readyState >= 2) onEgoLoaded();
        if (exoVideo.readyState >= 2) onExoLoaded();
    }
    """

def clear_loading_message():
    """清除加载消息"""
    return ""

# 创建 Gradio 界面
with gr.Blocks(title="ObjectRelator", theme=gr.themes.Soft()) as demo:
    # 标题
    gr.Markdown("""
    <div style="text-align: center; margin-bottom: 30px;">
        <h1 style="font-size: 3em; color: #2c3e50; margin: 0; text-shadow: 2px 2px 4px rgba(0,0,0,0.1);">
            ✨ ObjectRelator ✨
        </h1>
        <h2 style="font-size: 1.2em; color: #7f8c8d; margin: 10px 0; font-weight: 300;">
            👤 Enabling Cross-View Object Relation Understanding Across Ego-Centric and Exo-Centric Perspectives 📹
        </h2>
    </div>
    """, elem_id="title")
    

    # 控制面板
    with gr.Row(elem_id="control_panel"):
        with gr.Column(scale=1):
            scene_dropdown = gr.Dropdown(
                choices=list(SCENES.keys()),
                value="basketball",
                label="Select Scene",
                interactive=True,
                elem_id="scene_dropdown"
            )
        
        with gr.Column(scale=1):
            object_dropdown = gr.Dropdown(
                choices=SCENES["basketball"]["objects"],
                value=None,
                label="Select Object",
                interactive=True,
                elem_id="object_dropdown"
            )
    
    # 视频显示区域
    with gr.Row(elem_id="video_container"):
        with gr.Column():
            gr.Markdown("""
            <div style="text-align: center; margin-bottom: 15px;">
                <h3 style="color: #34495e; margin: 0; display: flex; align-items: center; justify-content: center;">
                    <span style="margin-right: 8px;">👤</span>
                    Ego View
                    <span style="margin-left: 8px;"></span>
                </h3>
            </div>
            """)
            ego_video = gr.Video(
                value=SCENES["basketball"]["ego_query"],
                autoplay=False,
                loop=True,
                show_download_button=False,
                show_share_button=False,
                elem_id="ego_video"
            )
        
        with gr.Column():
            gr.Markdown("""
            <div style="text-align: center; margin-bottom: 15px;">
                <h3 style="color: #34495e; margin: 0; display: flex; align-items: center; justify-content: center;">
                    <span style="margin-right: 8px;">📹</span>
                    Exo View
                    <span style="margin-left: 8px;"></span>
                </h3>
            </div>
            """)
            exo_video = gr.Video(
                value=SCENES["basketball"]["exo_target"],
                autoplay=False,
                loop=True,
                show_download_button=False,
                show_share_button=False,
                elem_id="exo_video"
            )
    
    # 场景选择事件处理
    scene_dropdown.change(
        fn=update_scene,
        inputs=[scene_dropdown],
        outputs=[ego_video, exo_video, object_dropdown]
    ).then(
        fn=None,
        js=sync_video_playback()
    )
    
    # 物体选择事件处理
    object_dropdown.change(
        fn=update_videos_with_object,
        inputs=[scene_dropdown, object_dropdown],
        outputs=[ego_video, exo_video]
    ).then(
        fn=None,
        js=sync_video_playback()
    )
    
    # 页面加载完成后执行同步播放
    demo.load(fn=None, js=sync_video_playback())

# 自定义CSS样式和初始JavaScript
demo.css = """
/* 全局样式 */
.gradio-container {
    max-width: 1400px;
    margin: 0 auto;
    padding: 20px;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    min-height: 100vh;
}

/* 主内容区域 */
.contain {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 20px;
    box-shadow: 0 20px 40px rgba(0,0,0,0.1);
    padding: 30px;
    backdrop-filter: blur(10px);
}

/* 标题样式 */
#title {
    margin-bottom: 30px;
}

/* 控制面板样式 */
#control_panel {
    margin-bottom: 25px;
    padding: 20px;
    background: linear-gradient(145deg, #f8f9fa, #e9ecef);
    border-radius: 15px;
    box-shadow: inset 5px 5px 10px #d1d5db, inset -5px -5px 10px #ffffff;
}

/* 下拉菜单样式 */
#scene_dropdown, #object_dropdown {
    margin: 10px;
}

.gr-dropdown {
    border-radius: 10px !important;
    box-shadow: 0 4px 8px rgba(0,0,0,0.1) !important;
    transition: all 0.3s ease !important;
}

.gr-dropdown:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 12px rgba(0,0,0,0.15) !important;
}

/* 视频容器样式 */
#video_container {
    gap: 30px;
    margin-top: 20px;
}

/* 视频样式 */
video {
    border-radius: 15px !important;
    box-shadow: 0 10px 30px rgba(0,0,0,0.2) !important;
    transition: all 0.3s ease !important;
    border: 3px solid transparent !important;
    background: linear-gradient(white, white) padding-box,
                linear-gradient(45deg, #667eea, #764ba2) border-box !important;
}

video:hover {
    transform: scale(1.02) !important;
    box-shadow: 0 15px 40px rgba(0,0,0,0.3) !important;
}

/* 视频标题样式 */
.gr-markdown h3 {
    background: linear-gradient(45deg, #667eea, #764ba2);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-weight: bold;
    text-shadow: none;
}

/* 响应式设计 */
@media (max-width: 768px) {
    .gradio-container {
        padding: 10px;
    }
    
    .contain {
        padding: 20px;
    }
    
    #video_container {
        flex-direction: column;
    }
    
    video {
        width: 100% !important;
    }
}

/* 加载动画 */
@keyframes pulse {
    0% { opacity: 1; }
    50% { opacity: 0.6; }
    100% { opacity: 1; }
}

.loading {
    animation: pulse 1.5s infinite;
}

/* 按钮和交互元素的美化 */
.gr-button {
    border-radius: 10px !important;
    background: linear-gradient(45deg, #667eea, #764ba2) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    transition: all 0.3s ease !important;
}

.gr-button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 15px rgba(102, 126, 234, 0.3) !important;
}
"""

if __name__ == "__main__":
    demo.launch(share=True)