import gradio as gr
import subprocess
import os

# Não precisamos mais de nenhuma função de setup, clone ou pip install aqui!
# O Dockerfile já cuidou de tudo.

def run_inference_app(video_file, seed_num):
    if video_file is None:
        return None, "Por favor, envie um arquivo de vídeo."
    
    # O app roda de dentro do diretório /app/SeedVR, então os caminhos são relativos
    input_folder = "inputs"
    os.makedirs(input_folder, exist_ok=True)
    
    # O Gradio nos dá um caminho temporário, vamos movê-lo para um local conhecido
    input_video_path = os.path.join(input_folder, os.path.basename(video_file.name))
    os.rename(video_file.name, input_video_path)
        
    output_folder = "outputs"
    os.makedirs(output_folder, exist_ok=True)
    
    # Comando de inferência. O ambiente conda já está "ativado" graças ao Dockerfile.
    command = [
        "torchrun", "--nproc-per-node=4",
        "projects/inference_seedvr2_3b.py",
        "--video_path", input_folder,
        "--output_dir", output_folder,
        "--seed", str(seed_num),
        "--res_h", "320",
        "--res_w", "512",
    ]

    log_output = "Iniciando a inferência...\n" + ' '.join(command) + "\n\n"
    
    try:
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding='utf-8'
        )

        while True:
            line = process.stdout.readline()
            if not line:
                break
            log_output += line
            print(line.strip())
            yield None, log_output

        process.wait()

        if process.returncode != 0:
            raise RuntimeError("O script de inferência falhou. Verifique os logs.")
        
        result_files = [os.path.join(output_folder, f) for f in os.listdir(output_folder) if f.endswith('.mp4')]
        if not result_files:
            return None, log_output + "\n\nERRO: Nenhum vídeo foi gerado."
            
        return result_files, log_output + "\n\nInferência concluída com sucesso!"

    except Exception as e:
        error_message = f"{log_output}\n\nOcorreu um erro: {str(e)}"
        return None, error_message

with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Inferência SeedVR2 com Ambiente Conda")
    gr.Markdown("Este ambiente foi construído com Conda usando um Dockerfile para máxima estabilidade.")

    with gr.Row():
        with gr.Column(scale=1):
            video_input = gr.File(label="Vídeo de Entrada")
            seed_input = gr.Number(label="Seed", value=123)
            run_button = gr.Button("Gerar Vídeo", variant="primary")
        with gr.Column(scale=2):
            gallery_output = gr.Gallery(label="Vídeo de Saída", show_label=True)
            log_display = gr.Textbox(label="Logs de Execução", lines=15, interactive=False, autoscroll=True)
            
    run_button.click(
        fn=run_inference_app,
        inputs=[video_input, seed_input],
        outputs=[gallery_output, log_display]
    )

demo.launch()