import gradio as gr
import os
import shutil
import subprocess

UPLOAD_DIR = "training_images"
OUTPUT_DIR = "lora_output"

def train_lora(images, learning_rate, num_epochs, rank):
    if os.path.exists(UPLOAD_DIR):
        shutil.rmtree(UPLOAD_DIR)
    os.makedirs(UPLOAD_DIR, exist_ok=True)

    if os.path.exists(OUTPUT_DIR):
        shutil.rmtree(OUTPUT_DIR)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    for idx, img in enumerate(images):
        img.save(os.path.join(UPLOAD_DIR, f"image_{idx}.png"))

    cmd = [
        "python", "train_lora.py",
        "--images_dir", UPLOAD_DIR,
        "--output_dir", OUTPUT_DIR,
        "--learning_rate", str(learning_rate),
        "--num_epochs", str(num_epochs),
        "--rank", str(rank),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)

    output_file = os.path.join(OUTPUT_DIR, "lora.safetensors")
    if os.path.exists(output_file):
        return f"✅ Treinamento finalizado!\nModelo salvo em: {output_file}\n\nLogs:\n{result.stdout}"
    else:
        return f"❌ Erro no treinamento:\n{result.stderr}"

with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Criador & Treinador de LoRA")
    with gr.Row():
        image_input = gr.File(
            file_types=[".png", ".jpg", ".jpeg"],
            file_types_display="images",
            file_count="multiple",
            label="Envie suas imagens (10–50)"
        )
    with gr.Row():
        learning_rate = gr.Number(value=1e-4, label="Learning Rate")
        num_epochs = gr.Number(value=10, label="Número de Epochs")
        rank = gr.Number(value=4, label="Rank do LoRA")
    with gr.Row():
        train_button = gr.Button("🚀 Treinar LoRA")
    output_text = gr.Textbox(label="Saída", lines=15)

    train_button.click(
        fn=train_lora,
        inputs=[image_input, learning_rate, num_epochs, rank],
        outputs=output_text
    )

# 👇 MUITO IMPORTANTE: apenas expor a variável demo
demo