import gradio as gr
import os
import shutil
import subprocess

# ---- Upload audio function ----
def upload_audio(audio_files, model_name):
    if not audio_files:
        return "⚠️ Please upload at least one audio file."
    if not model_name:
        return "⚠️ Model name cannot be empty."
    
    save_dir = os.path.join("dataset", model_name)
    os.makedirs(save_dir, exist_ok=True)

    for i, file in enumerate(audio_files):
        shutil.copy(file.name, os.path.join(save_dir, f"{model_name}_{i}.wav"))
    
    return f"✅ {len(audio_files)} audio files saved to '{save_dir}' folder!"

# ---- Start training function ----
def train_rvc(model_name, sample_rate, epochs, batch_size):
    dataset_dir = os.path.join("dataset", model_name)
    if not os.path.exists(dataset_dir):
        return f"⚠️ {dataset_dir} folder not found. Upload audio files first."

    try:
        cmd = [
            "python3", "train.py",
            "--model_name", model_name,
            "--dataset", dataset_dir,
            "--sample_rate", str(sample_rate),
            "--epochs", str(epochs),
            "--batch_size", str(batch_size)
        ]
        subprocess.Popen(cmd)
        return f"🚀 RVC v2 Training started!\nModel: {model_name}\nEpochs: {epochs}\nBatch: {batch_size}\nSample Rate: {sample_rate}Hz\n\nCheck console for progress!"
    except Exception as e:
        return f"❌ Training failed to start: {e}"

# ---- Interface ----
with gr.Blocks(title="RVC v2 Training") as demo:
    gr.Markdown("# 🎙️ RVC v2 Voice Model Training Tool")
    gr.Markdown("1️⃣ Upload audio → 2️⃣ Enter model name → 3️⃣ Start training")

    with gr.Tab("Upload Audio"):
        with gr.Row():
            audio_files = gr.File(file_count="multiple", label="🎧 Upload Audio Files (.wav)")
            model_name = gr.Textbox(label="Model Name", placeholder="e.g.: zeynep_rvc")
        output_upload = gr.Textbox(label="Status", lines=3)
        upload_button = gr.Button("📦 Upload Audio", variant="primary")
        upload_button.click(upload_audio, inputs=[audio_files, model_name], outputs=output_upload)

    with gr.Tab("Start Training"):
        sample_rate = gr.Dropdown(choices=[32000, 40000, 48000], value=40000, label="Sample Rate (Hz)")
        epochs = gr.Slider(50, 1000, value=200, step=50, label="Number of Epochs")
        batch_size = gr.Slider(4, 16, value=8, step=4, label="Batch Size")
        output_train = gr.Textbox(label="Training Status", lines=5)
        train_button = gr.Button("🚀 Start RVC v2 Training", variant="primary")
        train_button.click(train_rvc, inputs=[model_name, sample_rate, epochs, batch_size], outputs=output_train)


demo.launch()