import os, sys
import gradio as gr
import shutil
import requests
from zipfile import ZipFile
from tempfile import TemporaryDirectory

now_dir = os.getcwd()
sys.path.append(now_dir)

from assets.i18n.i18n import I18nAuto
from core import run_model_blender_script

i18n = I18nAuto()

def download_and_extract_model(url):
    temp_dir = TemporaryDirectory()
    zip_path = os.path.join(temp_dir.name, "model.zip")

    with requests.get(url, stream=True) as r:
        with open(zip_path, 'wb') as f:
            shutil.copyfileobj(r.raw, f)

    with ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(temp_dir.name)
    
    pth_files = [os.path.join(temp_dir.name, f) for f in os.listdir(temp_dir.name) if f.endswith('.pth')]

    if pth_files:
        return pth_files[0]
    else:
        raise FileNotFoundError("No se encontró un archivo .pth en el zip descargado.")

def download_and_set_paths(link_a, link_b):
    path_a = download_and_extract_model(link_a)
    path_b = download_and_extract_model(link_b)
    return path_a, path_b

def voice_blender_tab():
    gr.Markdown(i18n("## Voice Blender"))
    gr.Markdown(
        i18n(
            "Select two voice models, set your desired blend percentage, and blend them into an entirely new voice."
        )
    )
    with gr.Column():
        model_fusion_name = gr.Textbox(
            label=i18n("Model Name"),
            info=i18n("Name of the new model."),
            value="",
            max_lines=1,
            interactive=True,
            placeholder=i18n("Enter model name"),
        )
        with gr.Row():
            with gr.Column():
                model_fusion_a_link = gr.Textbox(
                    label=i18n("Hugging Face Link for Model A"),
                    value="",
                    interactive=True,
                    placeholder=i18n("Enter Hugging Face link for model A"),
                )
                model_fusion_a_path = gr.Textbox(
                    label=i18n("Path to Model A"),
                    value="",
                    interactive=False,
                    placeholder=i18n("Path will be set automatically after download"),
                )
            with gr.Column():
                model_fusion_b_link = gr.Textbox(
                    label=i18n("Hugging Face Link for Model B"),
                    value="",
                    interactive=True,
                    placeholder=i18n("Enter Hugging Face link for model B"),
                )
                model_fusion_b_path = gr.Textbox(
                    label=i18n("Path to Model B"),
                    value="",
                    interactive=False,
                    placeholder=i18n("Path will be set automatically after download"),
                )
        alpha_a = gr.Slider(
            minimum=0,
            maximum=1,
            label=i18n("Blend Ratio"),
            value=0.5,
            interactive=True,
            info=i18n(
                "Adjusting the position more towards one side or the other will make the model more similar to the first or second."
            ),
        )
        model_fusion_button = gr.Button(i18n("Fusion"), variant="primary")
        with gr.Row():
            model_fusion_output_info = gr.Textbox(
                label=i18n("Output Information"),
                info=i18n("The output information will be displayed here."),
                value="",
            )
            model_fusion_pth_output = gr.File(
                label=i18n("Download Model"), type="filepath", interactive=False
            )

    model_fusion_button.click(
        fn=download_and_set_paths,
        inputs=[model_fusion_a_link, model_fusion_b_link],
        outputs=[model_fusion_a_path, model_fusion_b_path],
    )

    model_fusion_button.click(
        fn=run_model_blender_script,
        inputs=[
            model_fusion_name,
            model_fusion_a_path,
            model_fusion_b_path,
            alpha_a,
        ],
        outputs=[model_fusion_output_info, model_fusion_pth_output],
    )