import os
import gradio as gr
from utils import get_models, get_categories, render, get_single_model_result_html

models = get_models()
cats   = get_categories()
default_models = ["GPT-4o", "Gemini 2.0", "Doubao"]

def update_models_incremental(cat, idx, selected_models, current_models_state):
    """Update the model results incrementally based on selected models."""
    if not selected_models:
        return "<p>No models selected</p>", []
    
    selected_set = set(selected_models)
    current_set = set(current_models_state) if current_models_state else set()
    
    models_to_add = selected_set - current_set
    models_to_remove = current_set - selected_set
    
    if not models_to_add and not models_to_remove:
        # no change
        current_html = get_current_html(cat, idx, selected_models)
        return current_html, selected_models
    
    new_html = render_models_html(cat, idx, selected_models)
    return new_html, selected_models

def render_models_html(cat, idx, selected_models):
    if not selected_models:
        return "<p>No models selected</p>"
    
    HTML_HEAD = '''
    <style>
    .model-results-container {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
        gap: 15px;
        margin: 20px 0;
        padding: 10px;
    }
    .single-model-result {
        border: 2px solid #333;
        border-radius: 12px;
        padding: 15px;
        text-align: center;
        background: linear-gradient(145deg, #f9f9f9, #e6e6e6);
        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
        transition: transform 0.2s ease, box-shadow 0.2s ease;
    }
    .single-model-result:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 12px rgba(0,0,0,0.15);
    }
    .model-name {
        font-weight: bold;
        font-size: 18px;
        color: #333;
        margin-bottom: 15px;
        padding: 8px;
        background-color: rgba(255,255,255,0.7);
        border-radius: 6px;
    }
    .model-image-container {
        height: 280px;
        display: flex;
        align-items: center;
        justify-content: center;
        overflow: hidden;
        border-radius: 8px;
        background-color: #fff;
    }
    .model-image {
        max-height: 100%;
        max-width: 100%;
        object-fit: contain;
        border-radius: 4px;
    }
    .no-image {
        color: #666;
        font-style: italic;
        background-color: #f0f0f0;
        padding: 20px;
        border-radius: 4px;
    }
    .loading-placeholder {
        color: #999;
        font-style: italic;
        background: linear-gradient(90deg, #f0f0f0 25%, #e0e0e0 50%, #f0f0f0 75%);
        background-size: 200% 100%;
        animation: loading 2s infinite;
        padding: 20px;
        border-radius: 4px;
    }
    @keyframes loading {
        0% { background-position: 200% 0; }
        100% { background-position: -200% 0; }
    }
    </style>
    <div class="model-results-container">
    '''
    HTML_TAIL = '</div>'
    
    html = HTML_HEAD
    for model in selected_models:
        html += get_single_model_result_html(model, cat, idx)
    html += HTML_TAIL
    
    return html

def get_current_html(cat, idx, selected_models):
    return render_models_html(cat, idx, selected_models)

# gradio demo
with gr.Blocks(title="KRIS Bench Results Viewer") as demo:
    gr.Markdown("# Image Editing Results on KRIS-Bench")
    
    current_idx = gr.State(0)
    current_models_state = gr.State([])
    
    with gr.Row():
        category = gr.Dropdown(
            label="Task Category", 
            choices=cats, 
            value=cats[0] if cats else None
        )
    
    with gr.Row():
        model_checkboxes = gr.CheckboxGroup(
            label="Select Models to Display",
            choices=models,
            value=default_models,
            container=True
        )
    
    with gr.Row():
        with gr.Column(scale=3):
            orig_img = gr.Image(label="Input Image", type="pil", interactive=False, height=300)
            
            with gr.Row():
                prev_btn = gr.Button("PREV", size="sm", scale=2)
                next_btn = gr.Button("NEXT", size="sm", scale=2)
                jump_input = gr.Number(
                    label="Jump to", 
                    precision=0, 
                    minimum=1, 
                    scale=1,
                    container=False
                )
                jump_btn = gr.Button("JUMP", size="sm", scale=1)
                image_counter = gr.HTML("<div style='text-align: center; font-size: 14px; font-weight: bold; margin: 0; padding: 8px; white-space: nowrap;'>0/0</div>")
        
        with gr.Column(scale=2):
            index_text = gr.Textbox(label="Index", interactive=False, lines=1)
            instruction_text = gr.Textbox(label="Instruction", interactive=False, lines=3)
            explanation_text = gr.Textbox(label="Explanation", interactive=False, lines=2)
    
    model_results_html = gr.HTML()

    # Event handlers
    def create_update_function(cat, idx, selected_models):
        if not selected_models:
            selected_models = []
        
        results = render(cat, idx, selected_models)
        
        return results

    def update_models_only(cat, idx, selected_models, current_models):
        html, new_models_state = update_models_incremental(cat, idx, selected_models, current_models)
        return html, new_models_state

    # Define outputs
    full_outputs = [orig_img, index_text, instruction_text, explanation_text, model_results_html, current_idx, image_counter, current_models_state]
    model_only_outputs = [model_results_html, current_models_state]
    
    # Event bindings
    category.change(
        fn=lambda c, m: create_update_function(c, 0, m) + (m,), 
        inputs=[category, model_checkboxes], 
        outputs=full_outputs
    )
    
    model_checkboxes.change(
        fn=update_models_only,
        inputs=[category, current_idx, model_checkboxes, current_models_state],
        outputs=model_only_outputs
    )
    
    prev_btn.click(
        fn=lambda c, i, m: create_update_function(c, i-1, m) + (m,), 
        inputs=[category, current_idx, model_checkboxes], 
        outputs=full_outputs
    )
    
    next_btn.click(
        fn=lambda c, i, m: create_update_function(c, i+1, m) + (m,), 
        inputs=[category, current_idx, model_checkboxes], 
        outputs=full_outputs
    )
    
    jump_btn.click(
        fn=lambda c, j, m: create_update_function(c, max(0, int(j) - 1) if j and str(j).isdigit() else 0, m) + (m,),
        inputs=[category, jump_input, model_checkboxes], 
        outputs=full_outputs
    )
    
    # Initialize
    if cats:
        init_results = create_update_function(cats[0], 0, default_models)
        demo.load(
            fn=lambda: init_results + (default_models,),
            outputs=full_outputs
        )

# simple CSS
demo.css = """
.gradio-container {
    max-width: 1800px !important;
    margin: 0 auto;
}

/* Navigation button styling - smaller size */
.gr-button[size="sm"] {
    font-size: 12px !important;
    font-weight: bold !important;
    padding: 6px 12px !important;
    margin: 2px !important;
    border-radius: 4px !important;
    min-height: 32px !important;
}

/* Jump input styling - smaller and more compact */
.gr-number input {
    font-size: 12px !important;
    padding: 4px 6px !important;
    text-align: center !important;
    border-radius: 4px !important;
    min-height: 32px !important;
    width: 60px !important;
}

/* Counter styling - inline with navigation */
.gr-html {
    display: flex;
    align-items: center;
    justify-content: center;
    min-height: 32px;
}

/* Checkbox group styling - full width */
.gr-checkbox-group {
    max-height: 150px;
    overflow-y: auto;
    border: 1px solid #ddd;
    border-radius: 8px;
    padding: 15px;
    background-color: #fafafa;
    width: 100%;
}

/* Information text boxes */
.gr-textbox textarea {
    font-size: 14px !important;
    line-height: 1.4 !important;
}

/* Row spacing */
.gr-row {
    margin: 8px 0 !important;
}

/* Column spacing */
.gr-column {
    margin: 3px !important;
}

/* Make images responsive */
.gr-image {
    border-radius: 8px !important;
    border: 1px solid #ddd !important;
}

/* HTML content styling */
.model-results-table {
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    border-radius: 8px;
    overflow: hidden;
}

.model-results-table td {
    background-color: #fff;
    transition: background-color 0.2s;
}

.model-results-table td:hover {
    background-color: #f5f5f5;
}

.model-name {
    background-color: transparent !important;
    color: black !important;
    font-size: 20px !important;
    font-weight: bold !important;
    height: 60px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
}

.model-image-container {
    height: 280px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    overflow: hidden !important;
}

.model-image {
    max-height: 100% !important;
    max-width: 100% !important;
    object-fit: contain !important;
}

/* Compact navigation row */
.gr-row .gr-column:has(.gr-button[size="sm"]) {
    margin: 1px !important;
}

/* Navigation container spacing */
.gr-row:has(.gr-button[size="sm"]) {
    margin: 5px 0 !important;
    gap: 5px !important;
}
"""

if __name__ == "__main__":
    demo.launch(server_name='0.0.0.0')