import gradio as gr
from rvc_inferpy import infer_audio
import os
import zipfile
import shutil
import urllib.request
import gdown

main_dir = os.getcwd()
models_dir = os.path.join(main_dir, "models")
os.makedirs(models_dir, exist_ok=True)

audio_files=[]
for filename in os.listdir("./audios"):
    if filename.endswith(('.wav','.mp3','.ogg','.flac','.m4a','.aac','.mp4')):
        audio_files.append(os.path.join('./audios',filename).replace('\\', '/'))

def extract_zip(extraction_folder, zip_name):
    os.makedirs(extraction_folder, exist_ok=True)
    with zipfile.ZipFile(zip_name, 'r') as zip_ref:
        zip_ref.extractall(extraction_folder)
    os.remove(zip_name)

    index_filepath, model_filepath = None, None
    for root, dirs, files in os.walk(extraction_folder):
        for name in files:
            if name.endswith('.index') and os.stat(os.path.join(root, name)).st_size > 1024 * 100:
                index_filepath = os.path.join(root, name)
            if name.endswith('.pth') and os.stat(os.path.join(root, name)).st_size > 1024 * 1024 * 40:
                model_filepath = os.path.join(root, name)

    if not model_filepath:
        raise Exception(f'No .pth model file was found in the extracted zip. Please check {extraction_folder}.')

    os.rename(model_filepath, os.path.join(extraction_folder, os.path.basename(model_filepath)))
    if index_filepath:
        os.rename(index_filepath, os.path.join(extraction_folder, os.path.basename(index_filepath)))
    
    for filepath in os.listdir(extraction_folder):
        if os.path.isdir(os.path.join(extraction_folder, filepath)):
            shutil.rmtree(os.path.join(extraction_folder, filepath))

def download_online_model(url, dir_name):
    try:
        print(f'[~] Downloading voice model with name {dir_name}...')
        zip_name = dir_name + ".zip"
        extraction_folder = os.path.join(models_dir, dir_name)
        if os.path.exists(extraction_folder):
            raise Exception(f'Voice model directory {dir_name} already exists! Choose a different name.')

        if 'drive.google.com' in url:
            gdown.download(url, output=zip_name, use_cookies=True, quiet=True, fuzzy=True)
        else:
            urllib.request.urlretrieve(url, zip_name)

        print('[~] Extracting zip file...')
        extract_zip(extraction_folder, zip_name)
        print(f'[+] {dir_name} Model successfully downloaded!')

    except Exception as e:
        raise Exception(str(e))

def process_audio(model_name, sound_path, f0_change, f0_method, output_format):
    try:
        print(model_name, sound_path, f0_change, f0_method, output_format) #
        # https://github.com/TheNeodev/rvc_inferpy/tree/main?tab=readme-ov-file#as-a-dependency-in-a-python-project <= wrong
        # https://github.com/TheNeodev/rvc_inferpy/blob/main/rvc_inferpy/infer.py#L117 <= correct
        inferred_audio = infer_audio(
            model_name=model_name,
            audio_path=sound_path,
            f0_change=f0_change,
            f0_method=f0_method
        )
        # You might need additional logic for handling output_format here
        return inferred_audio
    except Exception as e:
        print(e)
        raise gr.Error(e)

def save_to_wav2(dropbox):
    file_path=dropbox
    shutil.move(file_path,'./audios')
    return os.path.join('./audios',os.path.basename(file_path))
    
def get_name():
    if len(audio_files) > 0:
        return sorted(audio_files)[0]
    else:
        return ''

def change_choices2():
    audio_files=[]
    for filename in os.listdir("./audios"):
        if filename.endswith(('.wav','.mp3','.ogg','.flac','.m4a','.aac','.mp4')):
            audio_files.append(os.path.join('./audios',filename).replace('\\', '/'))
    return {"choices": sorted(audio_files), "__type__": "update"}, {"__type__": "update"}
    
def get_model_files():
    model_files=[]
    for root, dirs, files in os.walk(models_dir):
        for name in files:
            if filename.endswith(('.pth')):
               model_files.append(os.path.join(root, name).replace('\\', '/'))
    if len(model_files) == 0: model_files = [""]
    return model_files

def update_model_name():
    return gr.update(choices=get_model_files())

with gr.Blocks(theme=gr.themes.Base(), title=" rvc inferpy") as demo:
    gr.Markdown("<h1><center>rvc inferpy (Neo RVC Fork)</center></h1>")
    gr.Markdown("most simplest RVC inference")
    with gr.Tabs():
        with gr.TabItem("Inference"):
            model_name_input = gr.Dropdown(label="Model Name", choices=get_model_files(), value=get_model_files()[0], allow_custom_value=True)
            with gr.Row():
                dropbox = gr.Audio(label="Upload Audio for inference")
            with gr.Row():
                input_audio0 = gr.Dropdown(
                    label="Choose your audio.",
                    value="",
                    choices=audio_files
                )
            dropbox.upload(fn=save_to_wav2, inputs=[dropbox], outputs=[input_audio0])\
            .then(fn=change_choices2, inputs=None, outputs=[input_audio0])\
            .then(fn=update_model_name, inputs=None, outputs=[model_name_input])
            with gr.Row():
                f0_change_input = gr.Number(label="F0 Change", value=0)
            f0_method_input = gr.Dropdown(label="F0 Method", choices=["crepe", "dio", "harvest", "rmvpe", "fcpe", "hybrid[fcpe+rmvpe]"], value="crepe")
            output_format_input = gr.Dropdown(label="Output Format", choices=["wav", "mp3"], value="wav")
            submit_button = gr.Button("Run Inference")
            output_audio = gr.Audio(label="Inferred Audio", type="filepath")
            
            submit_button.click(
                fn=process_audio,
                inputs=[model_name_input, input_audio0, f0_change_input, f0_method_input, output_format_input],
                outputs=output_audio
            )
        
        with gr.TabItem("Download Models"):
            with gr.Row():
                url_input = gr.Textbox(label="Model URL")
                url_name_input = gr.Textbox(label="Model Name")
            download_button = gr.Button("Download")
            url_output = gr.Textbox(label="Output")
            
            download_button.click(
                fn=download_online_model,
                inputs=[url_input, url_name_input],
                outputs=url_output
            )

demo.launch()