import gradio as gr
from infer.modules.train.extract.extract_f0_rmvpe import FeatureInput
from infer.modules.train.extract_feature_print import HubertFeatureExtractor
from zero import zero


@zero(duration=300)
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

    hfe = HubertFeatureExtractor(exp_dir)
    try:
        hfe.run()
    except Exception as e:
        err = e

    hfe.logfile.seek(0)
    log += hfe.logfile.read()

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

    return log


class FeatureExtractionTab:
    def __init__(self):
        pass

    def ui(self):
        gr.Markdown("# Feature Extraction")
        gr.Markdown(
            "Before training, you need to extract features from the audio files. "
            "This process may take a while, depending on the number of audio files. "
            "Under the hood, this process extracts speech features using HuBERT and extracts F0 by RMVPE."
        )

        with gr.Row():
            self.extract_features_btn = gr.Button(
                value="Extract features", variant="primary"
            )
        with gr.Row():
            self.extract_features_log = gr.Textbox(
                label="Feature extraction log", lines=10
            )

    def build(self, exp_dir: gr.Textbox):
        self.extract_features_btn.click(
            fn=extract_features,
            inputs=[exp_dir],
            outputs=[self.extract_features_log],
        )