import gradio as gr
import os
import json
from serve.gradio_web_t2s import *
from serve.gradio_web_i2s import *
from serve.leaderboard import build_leaderboard_tab, build_empty_leaderboard_tab
from model.model_manager import ModelManager
from pathlib import Path
from constants import SERVER_PORT, ROOT_PATH, ELO_RESULTS_DIR


def build_combine_demo(models, elo_results_file, leaderboard_table_file):
    with gr.Blocks(
        title="Play with Open 3D Generative Models",
        theme=gr.themes.Default(),
        css=block_css,
    ) as demo:
        with gr.Tabs() as tabs_combine:
            with gr.Tab("Text-to-3D Generation", id=0):
                with gr.Tabs() as tabs_ig:
                    with gr.Tab("Text-to-3D Arena (battle)", id=0):
                        build_t2s_ui_side_by_side_anony(models)
                    with gr.Tab("Text-to-3D Arena (side-by-side)", id=1):
                        build_t2s_ui_side_by_side_named(models)
                    with gr.Tab("Text-to-3D Direct Chat", id=2):
                        build_t2s_ui_single_model(models)
                    if elo_results_file:
                        with gr.Tab("Text-to-3D Leaderboard", id=3):
                            build_leaderboard_tab(elo_results_file['text2shape'], leaderboard_table_file['text2shape'])
                    else:
                        with gr.Tab("Text-to-3D Leaderboard", id=3):
                            build_empty_leaderboard_tab()
                    with gr.Tab("About Us", id=4):
                        build_about()
            
            with gr.Tab("Image-to-3D Generation", id=5):
                with gr.Tabs() as tabs_ie:
                    with gr.Tab("Image-to-3D Arena (battle)", id=5):
                        build_i2s_ui_side_by_side_anony(models)
                    with gr.Tab("Image-to-3D Arena (side-by-side)", id=6):
                        build_i2s_ui_side_by_side_named(models)
                    with gr.Tab("Image-to-3D Direct Chat", id=7):
                        build_i2s_ui_single_model(models)
                    if elo_results_file:
                        with gr.Tab("Image-to-3D Leaderboard", id=8):
                            build_leaderboard_tab(elo_results_file['image2shape'], leaderboard_table_file['image2shape'])
                    else:
                        with gr.Tab("Image-to-3D Leaderboard", id=8):
                            build_empty_leaderboard_tab()
                    with gr.Tab("About Us", id=9):
                        build_about()

    return demo


def load_elo_results(elo_results_dir):
    from collections import defaultdict
    elo_results_file = defaultdict(lambda: None)
    leaderboard_table_file = defaultdict(lambda: None)
    if elo_results_dir is not None:
        elo_results_dir = Path(elo_results_dir)
        elo_results_file = {}
        leaderboard_table_file = {}
        for file in elo_results_dir.glob('elo_results_*.pkl'):
            if 'text2shape' in file.name:
                elo_results_file['text2shape'] = file
            elif 'image2shape' in file.name:
                elo_results_file['image2shape'] = file
            else:
                raise ValueError(f"Unknown file name: {file.name}")
        for file in elo_results_dir.glob('*_leaderboard.csv'):
            if 'text2shape' in file.name:
                leaderboard_table_file['text2shape'] = file
            elif 'image2shape' in file.name:
                leaderboard_table_file['image2shape'] = file
            else:
                raise ValueError(f"Unknown file name: {file.name}")
            
    return elo_results_file, leaderboard_table_file
    
if __name__ == "__main__":
    server_port = int(SERVER_PORT)
    root_path = ROOT_PATH
    elo_results_dir = ELO_RESULTS_DIR
    models = ModelManager()
    
    elo_results_file, leaderboard_table_file = load_elo_results(elo_results_dir)
    # elo_results_file, leaderboard_table_file = None, None
    demo = build_combine_demo(models, elo_results_file, leaderboard_table_file)
    demo.queue(max_size=20).launch(server_port=server_port, root_path=ROOT_PATH, debug=True)