import gradio as gr
from gradio_rerun import Rerun
from data.loader import load_simulation_data
from visualization.visualizer import visualize_simulation
from visualization.et_visualizer import visualize_et_data


def update_simulation_dropdown(file):
    simulations, descriptions = load_simulation_data(file)
    return gr.Dropdown(
        choices=descriptions if descriptions else [],
        value=None,
        allow_custom_value=False
    )


def create_app():
    with gr.Blocks() as demo:
        with gr.Tabs() as tabs:
            with gr.Tab("Camera Simulation"):
                gr.Markdown("""
                # Camera Simulation Visualizer
                Upload a JSON file containing camera simulation data and select a simulation to visualize.
                """)

                with gr.Row():
                    file_input = gr.File(
                        label="Upload Simulation JSON",
                        file_types=[".json"]
                    )
                    simulation_dropdown = gr.Dropdown(
                        label="Select Simulation",
                        choices=[],
                        type="index",
                        scale=2
                    )

                with gr.Row():
                    viewer = Rerun(streaming=False)

                file_input.change(
                    update_simulation_dropdown,
                    inputs=[file_input],
                    outputs=[simulation_dropdown]
                )

                simulation_dropdown.change(
                    visualize_simulation,
                    inputs=[file_input, simulation_dropdown],
                    outputs=[viewer]
                )

            with gr.Tab("E.T. Dataset"):
                gr.Markdown("""
                # E.T. Dataset Visualizer
                Upload trajectory (.txt) and character (.npy) files to visualize the E.T. dataset.
                """)

                with gr.Row():
                    traj_file = gr.File(
                        label="Trajectory File (.txt)",
                        file_types=[".txt"]
                    )
                    char_file = gr.File(
                        label="Character File (.npy)",
                        file_types=[".npy"]
                    )

                with gr.Row():
                    et_viewer = Rerun(streaming=False)

                def process_et_files(traj_file, char_file):
                    if traj_file is None or char_file is None:
                        return None
                    return visualize_et_data(traj_file.name, char_file.name)

                with gr.Row():
                    visualize_btn = gr.Button("Visualize")
                    visualize_btn.click(
                        process_et_files,
                        inputs=[traj_file, char_file],
                        outputs=[et_viewer]
                    )

    return demo


if __name__ == "__main__":
    demo = create_app()
    demo.queue().launch(share=False)