import gradio as gr
import os

from app_utils import (
    update_scene_display,
    update_log_display,
    run_simulation,
    update_history_display,
    record_access,
    cleanup_session,
    TMP_ROOT,
    MODEL_CHOICES,
    SCENE_CONFIGS,
)

header_html = """
<div style="display: flex; justify-content: space-between; align-items: center; width: 100%; margin-bottom: 20px; padding: 20px; background: linear-gradient(135deg, #528bdb 0%, #a7b5d0 100%); border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);">
    <div style="display: flex; align-items: center;">
        <img src="https://www.shlab.org.cn/static/img/index_14.685f6559.png" alt="Institution Logo" style="height: 60px; margin-right: 20px;">
        <div>
            <h1 style="margin: 0; color: #2c3e50; font-weight: 600;">🤖 InternManip Model Inference Demo</h1>
            <p style="margin: 4px 0 0 0; color: #5d6d7e; font-size: 0.9em;">Model trained on InternManip framework</p>
        </div>
    </div>
    <div style="display: flex; gap: 15px; align-items: center;">
        <a href="https://github.com/InternRobotics" target="_blank" style="text-decoration: none; transition: transform 0.2s;" onmouseover="this.style.transform='scale(1.1)'" onmouseout="this.style.transform='scale(1)'">
            <img src="https://github.githubassets.com/images/modules/logos_page/GitHub-Mark.png" alt="GitHub" style="height: 30px;">
        </a>
        <a href="https://huggingface.co/InternRobotics" target="_blank" style="text-decoration: none; transition: transform 0.2s;" onmouseover="this.style.transform='scale(1.1)'" onmouseout="this.style.transform='scale(1)'">
            <img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="HuggingFace" style="height: 30px;">
        </a>
        <a href="https://huggingface.co/spaces/OpenRobotLab/InternNav-eval-demo" target="_blank">
            <button style="padding: 8px 15px; background: #3498db; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: 500; transition: all 0.2s;" 
                    onmouseover="this.style.backgroundColor='#2980b9'; this.style.transform='scale(1.05)'" 
                    onmouseout="this.style.backgroundColor='#3498db'; this.style.transform='scale(1)'">
                Go to InternNav Demo
            </button>
        </a>
    </div>
</div>
"""


# CSS style
custom_css = """
#simulation-panel {
    border-radius: 8px;
    padding: 20px;
    background: #f9f9f9;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
#result-panel {
    border-radius: 8px;
    padding: 20px;
    background: #f0f8ff;
}
.dark #simulation-panel { background: #2a2a2a; }
.dark #result-panel { background: #1a2a3a; }

.history-container {
    max-height: 600px;
    overflow-y: auto;
    margin-top: 20px;
}

.history-accordion {
    margin-bottom: 10px;
}

.logs-container {
    max-height: 500px;
    overflow-y: auto;
    margin-top: 20px;
    padding: 15px;
    background: #f5f5f5;
    border-radius: 8px;
}

.dark .logs-container {
    background: #2a2a2a;
}

.log-table {
    width: 100%;
    border-collapse: collapse;
}

.log-table th, .log-table td {
    padding: 8px 12px;
    border: 1px solid #ddd;
    text-align: left;
}

.dark .log-table th, .dark .log-table td {
    border-color: #444;
}
"""

def start_session(req: gr.Request):
    user_dir = os.path.join(TMP_ROOT, str(req.session_hash))
    os.makedirs(user_dir, exist_ok=True)
    
# Gradio UI
with gr.Blocks(title="InternManip Model Inference Demo", css=custom_css) as demo:
    gr.HTML(header_html)
    
    history_state = gr.State([])

    with gr.Row():
        with gr.Column(elem_id="simulation-panel"):
            gr.Markdown("### Simulation Settings")
            
            scene_dropdown = gr.Dropdown(
                label="Choose a scene",
                choices=list(SCENE_CONFIGS.keys()),
                value="scene_1",
                interactive=True
            )
            
            scene_description = gr.Markdown("")
            scene_preview = gr.Image(
                label="Scene Preview",
                elem_classes=["scene-preview"],
                interactive=False
            )
            
            scene_dropdown.change(
                update_scene_display,
                inputs=scene_dropdown,
                outputs=[scene_description, scene_preview]
            )
            
            prompt_input = gr.Textbox(
                label="Manipulation Prompt",
                value="Move the milk carton to the top of the ceramic bowl.",
                placeholder="Example: 'Move the milk carton to the top of the ceramic bowl.'",
                lines=2,
                max_lines=4
            )
            
            model_dropdown = gr.Dropdown(
                label="Chose a pretrained model",
                choices=MODEL_CHOICES,
                value=MODEL_CHOICES[0]
            )

            with gr.Accordion("Advance Settings", open=False):
                max_steps = gr.Slider(
                    minimum=50,
                    maximum=500,
                    value=300,
                    step=10,
                    label="Max Steps"
                )
            
            submit_btn = gr.Button("Apply and Start Simulation", variant="primary")
        
        with gr.Column(elem_id="result-panel"):
            gr.Markdown("### Result")

            video_output = gr.Video(
                label="Live",
                interactive=False,
                format="mp4",
                autoplay=True,
                streaming=True
            )
            
            with gr.Column() as history_container:
                gr.Markdown("### History")
                gr.Markdown("#### History will be reset after refresh")
                
                # Precreate 10 history slot
                history_slots = []
                for i in range(10):
                    with gr.Column(visible=False) as slot:
                        with gr.Accordion(visible=False, open=False) as accordion:
                            video = gr.Video(interactive=False) 
                            detail_md = gr.Markdown()  # display detail info
                    history_slots.append((slot, accordion, video, detail_md)) 
    
    # display log infomation (DEV ONLY)
    with gr.Accordion("查看系统访问日志(DEV ONLY)", visible=False, open=False):
        logs_display = gr.Markdown()
        refresh_logs_btn = gr.Button("刷新日志", variant="secondary")
        
        refresh_logs_btn.click(
            update_log_display,
            outputs=logs_display
        )
    
    # Examples
    gr.Examples(
        examples=[
            ["scene_1", "Move the milk carton to the top of the ceramic bowl.", "gr1", 300],
        ],
        inputs=[scene_dropdown, prompt_input, model_dropdown, max_steps],
        label="Examples"
    )
    
    # 提交处理
    submit_btn.click(
        fn=run_simulation,
        inputs=[scene_dropdown, prompt_input, model_dropdown, max_steps, history_state],
        outputs=[video_output, history_state],
        queue=True
    ).then(
        fn=update_history_display,
        inputs=history_state,
        outputs=[comp for slot in history_slots for comp in slot],
        queue=True
    ).then(
        fn=update_log_display,
        outputs=logs_display
    )
    
    demo.load(
        start_session
    ).then(
        fn=lambda: update_scene_display("scene_1"),
        outputs=[scene_description, scene_preview]
    ).then(
        fn=record_access,
        inputs=None,
        outputs=logs_display,
        queue=False
    ).then(
        fn=update_log_display,
        outputs=logs_display
    )
    
    demo.queue(default_concurrency_limit=8)

    demo.unload(fn=cleanup_session)


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