import gradio as gr

from modules.Enhancer.ResembleEnhance import unload_enhancer
from modules.models import unload_chat_tts
from modules.speaker import speaker_mgr
from modules.webui import webui_config
from modules.webui.webui_utils import get_speaker_names

from .ft_ui_utils import get_datasets_listfile, run_speaker_ft
from .ProcessMonitor import ProcessMonitor


class SpeakerFt:
    def __init__(self):
        self.process_monitor = ProcessMonitor()
        self.status_str = "idle"

    def unload_main_thread_models(self):
        unload_chat_tts()
        unload_enhancer()

    def run(
        self,
        batch_size: int,
        epochs: int,
        lr: str,
        train_text: bool,
        data_path: str,
        select_speaker: str = "",
    ):
        if self.process_monitor.process:
            return
        self.unload_main_thread_models()
        spk_path = None
        if select_speaker != "" and select_speaker != "none":
            select_speaker = select_speaker.split(" : ")[1].strip()
            spk = speaker_mgr.get_speaker(select_speaker)
            if spk is None:
                return ["Speaker not found"]
            spk_filename = speaker_mgr.get_speaker_filename(spk.id)
            spk_path = f"./data/speakers/{spk_filename}"

        command = ["python3", "-m", "modules.finetune.train_speaker"]
        command += [
            f"--batch_size={batch_size}",
            f"--epochs={epochs}",
            f"--data_path={data_path}",
        ]
        if train_text:
            command.append("--train_text")
        if spk_path:
            command.append(f"--init_speaker={spk_path}")

        self.status("Training process starting")

        self.process_monitor.start_process(command)

        self.status("Training started")

    def status(self, text: str):
        self.status_str = text

    def flush(self):
        stdout, stderr = self.process_monitor.get_output()
        return f"{self.status_str}\n{stdout}\n{stderr}"

    def clear(self):
        self.process_monitor.stdout = ""
        self.process_monitor.stderr = ""
        self.status("Logs cleared")

    def stop(self):
        self.process_monitor.stop_process()
        self.status("Training stopped")


def create_speaker_ft_tab(demo: gr.Blocks):
    spk_ft = SpeakerFt()
    speakers, speaker_names = get_speaker_names()
    speaker_names = ["none"] + speaker_names

    with gr.Row():
        with gr.Column(scale=2):
            with gr.Group():
                gr.Markdown("🎛️hparams")
                dataset_input = gr.Dropdown(
                    label="Dataset", choices=get_datasets_listfile()
                )
                lr_input = gr.Textbox(label="Learning Rate", value="1e-2")
                epochs_input = gr.Slider(
                    label="Epochs", value=10, minimum=1, maximum=100, step=1
                )
                batch_size_input = gr.Slider(
                    label="Batch Size", value=4, minimum=1, maximum=64, step=1
                )
                train_text_checkbox = gr.Checkbox(label="Train text_loss", value=True)
                init_spk_dropdown = gr.Dropdown(
                    label="Initial Speaker",
                    choices=speaker_names,
                    value="none",
                )

            with gr.Group():
                start_train_btn = gr.Button("Start Training")
                stop_train_btn = gr.Button("Stop Training")
                clear_train_btn = gr.Button("Clear logs")
        with gr.Column(scale=5):
            with gr.Group():
                # log
                gr.Markdown("📜logs")
                log_output = gr.Textbox(
                    show_label=False, label="Log", value="", lines=20, interactive=True
                )

    start_train_btn.click(
        spk_ft.run,
        inputs=[
            batch_size_input,
            epochs_input,
            lr_input,
            train_text_checkbox,
            dataset_input,
            init_spk_dropdown,
        ],
        outputs=[],
    )
    stop_train_btn.click(spk_ft.stop)
    clear_train_btn.click(spk_ft.clear)

    if webui_config.experimental:
        demo.load(spk_ft.flush, every=1, outputs=[log_output])