import gradio as gr
import os, sys
from colorama import Fore
from rvcsynch import download_from_url
now_dir = os.getcwd()
sys.path.append(now_dir)


def show_available(filepath,format=None):
    if format:
        print(f"Format: {format}")
        files = []
        for file in os.listdir(filepath):
            if file.endswith(format):
                print(f"Matches format: {file}")
                files.append(file)
            else:
                print(f"Does not match format: {file}")
        print(f"Matches: {files}")
        if len(files) < 1:
            return ['']
        return files
    if len(os.listdir(filepath)) < 1:
        return ['']
    return os.listdir(filepath)

# Function to detect the .pth and .index files
def detect_files(model_name):
    model_dir = f"{now_dir}/assets/weights/{model_name}"
    index_dir = f"{now_dir}/logs/{model_name}"

    # Detect .pth file
    model_pth_file = None
    for file in os.listdir(model_dir):
        if file.endswith(".pth"):
            model_pth_file = os.path.join(model_dir, file)
            break

    # Detect .index file
    index_file = None
    for file in os.listdir(index_dir):
        if file.endswith(".index"):
            index_file = os.path.join(index_dir, file)
            break

    if model_pth_file and index_file:
        return f"Model .pth file: {model_pth_file}\nIndex file: {index_file}"
    else:
        return "Model .pth or index file not found."

# Function to process the audio using the detected files
def process_audio(model_name, pitch, input_path, f0_method, save_as, index_rate, volume_normalization, consonant_protection):
    model_dir = f"{now_dir}/assets/weights/{model_name}"
    index_dir = f"{now_dir}/logs/{model_name}"

    # Detect files
    model_pth_file = None
    index_file = None
    for file in os.listdir(model_dir):
        if file.endswith(".pth"):
            model_pth_file = os.path.join(model_dir, file)
            break

    for file in os.listdir(index_dir):
        if file.endswith(".index"):
            index_file = os.path.join(index_dir, file)
            break

    if not model_pth_file or not index_file:
        return "Model .pth or index file not found.", None

    if not os.path.exists(input_path):
        return f"{input_path} was not found in your RVC folder.", None

    # Set environment variables for paths
    os.environ['index_root'] = os.path.dirname(index_file)
    index_path = os.path.basename(index_file)
    os.environ['weight_root'] = os.path.dirname(model_pth_file)

    # Remove any previous output
    if os.path.exists(save_as):
        os.remove(save_as)

    # Execute the CLI command
    os.system(f"python {now_dir}/tools/infer_cli.py --f0up_key {pitch} --input_path {input_path} --index_path {index_path} --f0method {f0_method} --opt_path {save_as} --model_name {model_name} --index_rate {index_rate} --device 'cuda:0' --is_half True --filter_radius 3 --resample_sr 0 --rms_mix_rate {volume_normalization} --protect {consonant_protection}")

    if os.path.exists(save_as):
        return "Processing complete. Here is your output audio:", save_as
    else:
        return "Error in processing audio.", None

# Gradio interface
with gr.Blocks(theme="gradio/soft") as demo:
    gr.Markdown("# 🔊 ** RVC GUI**")

    with gr.Tabs():
        model_name = gr.Textbox(label="Model Name", value="Ren")
        pitch = gr.Slider(minimum=-12, maximum=12, step=1, label="Pitch", value=0)
        with gr.Tab("Infernece"):
            input_path = gr.Dropdown(label="",choices=show_available('audios'),value='',interactive=True)
            f0_method = gr.Radio(choices=["rmvpe", "pm", "crepe"], label="F0 Method", value="rmvpe")
            save_as = gr.Textbox(label="Save As", value="/content/RVC/audios/cli_output.wav")
            index_rate = gr.Slider(minimum=0, maximum=1, step=0.01, label="Index Rate", value=0.5)
            volume_normalization = gr.Slider(minimum=0, maximum=1, step=0.01, label="Volume Normalization", value=0)
            consonant_protection = gr.Slider(minimum=0, maximum=1, step=0.01, label="Consonant Protection", value=0.5)

        output_text = gr.Textbox(label="Output")
        output_audio = gr.Audio(label="Processed Audio")
        # Button to detect files
        detect_btn = gr.Button("Detect Files")
        detect_btn.click(fn=detect_files, inputs=[model_name], outputs=output_text)
        # Button to process the audio and return audio output
        submit_btn = gr.Button("Submit")
        submit_btn.click(fn=process_audio, 
                         inputs=[model_name, pitch, input_path, f0_method, save_as, index_rate, volume_normalization, consonant_protection], 
                         outputs=[output_text, output_audio])
        
    with gr.Tab("Download a model"):
        url = gr.Textbox(label="url models...")
        url_name = gr.Textbox(label="models names...")
        download_fm = gr.Button("download model")
        download_fm.click(fn=download_from_url, inputs=[url, url_name], outputs=url_name)
            

    
        
# Launch the app
demo.launch()