import gradio as gr
import zipfile
import os
import tempfile
import shutil
from infer.modules.train.preprocess import PreProcess, preprocess_trainset
from infer.modules.train.extract.extract_f0_rmvpe import FeatureInput
from zero import zero


def extract_audio_files(zip_file: str, target_dir: str) -> list[str]:
    with zipfile.ZipFile(zip_file, "r") as zip_ref:
        zip_ref.extractall(target_dir)

    audio_files = [
        os.path.join(target_dir, f)
        for f in os.listdir(target_dir)
        if f.endswith((".wav", ".mp3", ".ogg"))
    ]
    if not audio_files:
        raise gr.Error("No audio files found at the top level of the zip file")

    return audio_files


def train_rvc_model(audio_files: list[str]) -> str:
    return "model_path"


def preprocess(zip_file: str) -> str:
    temp_dir = tempfile.mkdtemp()
    print(f"Using exp dir: {temp_dir}")

    data_dir = os.path.join(temp_dir, "_data")
    os.makedirs(data_dir)
    audio_files = extract_audio_files(zip_file, data_dir)

    pp = PreProcess(48000, temp_dir, 3.0, False)
    pp.pipeline_mp_inp_dir(data_dir, 4)

    pp.logfile.seek(0)
    log = pp.logfile.read()

    return temp_dir, f"Preprocessed {len(audio_files)} audio files.\n{log}"


def download_expdir(exp_dir: str) -> str:
    shutil.make_archive(exp_dir, "zip", exp_dir)
    return f"{exp_dir}.zip"


@zero(duration=120)
def extract_features(exp_dir: str) -> str:
    err = None
    fi = FeatureInput(exp_dir)
    try:
        fi.run()
    except Exception as e:
        err = e

    fi.logfile.seek(0)
    log = fi.logfile.read()

    if err:
        log = f"Error: {err}\n{log}"

    return log


with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            zip_file = gr.File(
                label="Upload a zip file containing audio files for training",
                file_types=["zip"],
            )
            exp_dir = gr.Textbox(label="Experiment directory", visible=True)
            preprocess_btn = gr.Button(value="Preprocess", variant="primary")
        with gr.Column():
            preprocess_output = gr.Textbox(label="Preprocessing output", lines=5)

    with gr.Row():
        with gr.Column():
            extract_features_btn = gr.Button(
                value="Extract features", variant="primary"
            )
        with gr.Column():
            extract_features_output = gr.Textbox(
                label="Feature extraction output", lines=5
            )

    with gr.Row():
        with gr.Column():
            download_expdir_btn = gr.Button(
                value="Download experiment directory", variant="primary"
            )
        with gr.Column():
            download_expdir_output = gr.File(label="Download experiment directory")

    preprocess_btn.click(
        fn=preprocess,
        inputs=[zip_file],
        outputs=[exp_dir, preprocess_output],
    )

    extract_features_btn.click(
        fn=extract_features,
        inputs=[exp_dir],
        outputs=[extract_features_output],
    )

    download_expdir_btn.click(
        fn=download_expdir,
        inputs=[exp_dir],
        outputs=[download_expdir_output],
    )

    app.launch()