import os
import gradio as gr
import numpy as np
from PIL import Image

from inference.seg import process_image_or_video
from config import SAPIENS_LITE_MODELS_PATH

def update_model_choices(task):
    model_choices = list(SAPIENS_LITE_MODELS_PATH[task.lower()].keys())
    return gr.Dropdown(choices=model_choices, value=model_choices[0] if model_choices else None)

def gradio_wrapper(input_image, task, version):
    if isinstance(input_image, np.ndarray):
        input_image = Image.fromarray(input_image)
    
    result = process_image_or_video(input_image, task=task.lower(), version=version)
    
    return result

with gr.Blocks() as demo:
    gr.Markdown("# Sapiens Arena 🤸🏽‍♂️ - WIP devmode- Not yet available")
    with gr.Tabs():
        with gr.TabItem('Image'):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(label="Input Image", type="pil")
                    select_task = gr.Radio(
                        ["seg", "pose", "depth", "normal"], 
                        label="Task", 
                        info="Choose the task to perform",
                        value="seg"
                    )
                    model_name = gr.Dropdown(
                        label="Model Version",
                        choices=list(SAPIENS_LITE_MODELS_PATH["seg"].keys()),
                        value="sapiens_0.3b",
                    )
                with gr.Column():
                    result_image = gr.Image(label="Result")
                    run_button = gr.Button("Run")
        
        with gr.TabItem('Video'):
            gr.Markdown("In construction")

    select_task.change(fn=update_model_choices, inputs=select_task, outputs=model_name)

    run_button.click(
        fn=gradio_wrapper,
        inputs=[input_image, select_task, model_name],
        outputs=[result_image],
    )

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