import os
import subprocess
import gradio as gr
import time

# Variabile di controllo per il setup
SETUP_COMPLETED = False

def perform_setup():
    global SETUP_COMPLETED
    if SETUP_COMPLETED:
        return "Setup già completato", True
    
    try:
        # Clona il repository
        if not os.path.exists("Open-Sora"):
            os.system("git clone https://github.com/hpcaitech/Open-Sora")
        
        # Installa il pacchetto in modalità sviluppatore
        os.system("cd Open-Sora && pip install -e .")
        
        # Installa xformers con la versione appropriata
        os.system("pip install xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu118")
        
        # Scarica il modello
        if not os.path.exists("ckpts"):
            os.system("mkdir -p ckpts")
            os.system("huggingface-cli download hpcai-tech/Open-Sora-v2 --local-dir ./ckpts")
        
        SETUP_COMPLETED = True
        return "Setup completato con successo!", True
    except Exception as e:
        return f"Errore durante il setup: {str(e)}", False

def generate_video(prompt, setup_first=True):
    # Esegui il setup se necessario
    if setup_first and not SETUP_COMPLETED:
        status, success = perform_setup()
        if not success:
            return status, None
    
    try:
        # Crea directory per gli output
        if not os.path.exists("outputs"):
            os.makedirs("outputs")
        
        # ID univoco per l'output
        timestamp = int(time.time())
        output_dir = f"outputs/video_{timestamp}"
        os.makedirs(output_dir, exist_ok=True)
        
        # Prepara il comando (modifica la configurazione per disabilitare flash-attn)
        cmd = f"cd Open-Sora && python -m scripts.diffusion.inference configs/diffusion/inference/256px.py --prompt \"{prompt}\" --save-dir ../{output_dir} --model.use-flash-attn=False"
        
        print(f"Esecuzione comando: {cmd}")
        result = os.system(cmd)
        
        if result != 0:
            return "Errore durante la generazione del video. Controlla i log.", None
        
        # Cerca il video generato
        video_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if f.endswith(".mp4")]
        if not video_files:
            return "Nessun video generato trovato", None
        
        # Prendi il più recente
        video_path = sorted(video_files, key=lambda x: os.path.getmtime(x), reverse=True)[0]
        
        return "Video generato con successo!", video_path
    except Exception as e:
        return f"Errore: {str(e)}", None

# Interfaccia Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Generatore di Video con Open-Sora")
    
    with gr.Row():
        with gr.Column():
            setup_button = gr.Button("Setup iniziale (esegui una volta)")
            setup_output = gr.Textbox(label="Stato Setup", value="Non eseguito")
            
            gr.Markdown("---")
            
            prompt_input = gr.Textbox(
                label="Descrivi il video che vuoi generare",
                placeholder="Es. Un panda che mangia bambù in una foresta, stile cinematografico"
            )
            generate_btn = gr.Button("Genera Video")
        
        with gr.Column():
            status_output = gr.Textbox(label="Stato Generazione")
            video_output = gr.Video(label="Video Generato")
    
    # Collega i pulsanti alle funzioni
    setup_button.click(
        fn=perform_setup,
        outputs=[setup_output]
    )
    
    generate_btn.click(
        fn=generate_video,
        inputs=[prompt_input],
        outputs=[status_output, video_output]
    )

# Avvia l'interfaccia
demo.launch()