import os
import gradio as gr
import spaces
from infer_rvc_python import BaseLoader
import random
import logging
import time
import soundfile as sf
from infer_rvc_python.main import download_manager
import zipfile
import edge_tts
import asyncio
import librosa
import traceback
import soundfile as sf
from pedalboard import Pedalboard, Reverb, Compressor, HighpassFilter
from pedalboard.io import AudioFile
from pydub import AudioSegment
import noisereduce as nr
import numpy as np
import urllib.request
import shutil
import threading

logging.getLogger("infer_rvc_python").setLevel(logging.ERROR)

# Ensure the correct path to the models directory
model_dir = os.path.join(os.path.dirname(__file__), "models")

converter = BaseLoader(only_cpu=False, hubert_path=None, rmvpe_path=None)

title = "<center><strong><font size='7'>Vodex AI</font></strong></center>"
theme = "aliabid94/new-theme"

def find_files(directory):
    file_paths = []
    for filename in os.listdir(directory):
        if filename.endswith('.pth') or filename.endswith('.zip') or filename.endswith('.index'):
            file_paths.append(os.path.join(directory, filename))
    return file_paths

def unzip_in_folder(my_zip, my_dir):
    with zipfile.ZipFile(my_zip) as zip:
        for zip_info in zip.infolist():
            if zip_info.is_dir():
                continue
            zip_info.filename = os.path.basename(zip_info.filename)
            zip.extract(zip_info, my_dir)

def find_my_model(a_, b_):
    if a_ is None or a_.endswith(".pth"):
        return a_, b_

    txt_files = []
    for base_file in [a_, b_]:
        if base_file is not None and base_file.endswith(".txt"):
            txt_files.append(base_file)

    directory = os.path.dirname(a_)

    for txt in txt_files:
        with open(txt, 'r') as file:
            first_line = file.readline()

        download_manager(
            url=first_line.strip(),
            path=directory,
            extension="",
        )

    for f in find_files(directory):
        if f.endswith(".zip"):
            unzip_in_folder(f, directory)

    model = None
    index = None
    end_files = find_files(directory)

    for ff in end_files:
        if ff.endswith(".pth"):
            model = os.path.join(directory, ff)
            gr.Info(f"Model found: {ff}")
        if ff.endswith(".index"):
            index = os.path.join(directory, ff)
            gr.Info(f"Index found: {ff}")

    if not model:
        gr.Error(f"Model not found in: {end_files}")

    if not index:
        gr.Warning("Index not found")

    return model, index

def get_file_size(url):
    if "huggingface" not in url:
        raise ValueError("Only downloads from Hugging Face are allowed")

    try:
        with urllib.request.urlopen(url) as response:
            info = response.info()
            content_length = info.get("Content-Length")

        file_size = int(content_length)
        if file_size > 500000000:
            raise ValueError("The file is too large. You can only download files up to 500 MB in size.")

    except Exception as e:
        raise e

def clear_files(directory):
    time.sleep(15)
    print(f"Clearing files: {directory}.")
    shutil.rmtree(directory)

def get_my_model(url_data):
    if not url_data:
        return None, None

    if "," in url_data:
        a_, b_ = url_data.split()
        a_, b_ = a_.strip().replace("/blob/", "/resolve/"), b_.strip().replace("/blob/", "/resolve/")
    else:
        a_, b_ = url_data.strip().replace("/blob/", "/resolve/"), None

    out_dir = "downloads"
    folder_download = str(random.randint(1000, 9999))
    directory = os.path.join(out_dir, folder_download)
    os.makedirs(directory, exist_ok=True)

    try:
        get_file_size(a_)
        if b_:
            get_file_size(b_)

        valid_url = [a_] if not b_ else [a_, b_]
        for link in valid_url:
            download_manager(
                url=link,
                path=directory,
                extension="",
            )

        for f in find_files(directory):
            if f.endswith(".zip"):
                unzip_in_folder(f, directory)

        model = None
        index = None
        end_files = find_files(directory)

        for ff in end_files:
            if ff.endswith(".pth"):
                model = ff
                gr.Info(f"Model found: {ff}")
            if ff.endswith(".index"):
                index = ff
                gr.Info(f"Index found: {ff}")

        if not model:
            raise ValueError(f"Model not found in: {end_files}")

        if not index:
            gr.Warning("Index not found")
        else:
            index = os.path.abspath(index)

        return os.path.abspath(model), index

    except Exception as e:
        raise e
    finally:
        t = threading.Thread(target=clear_files, args=(directory,))
        t.start()

def convert_now(audio_files, random_tag, converter):
    return converter(
        audio_files,
        random_tag,
        overwrite=False,
        parallel_workers=8
    )

def apply_noisereduce(audio_list):
    print("Applying noise reduction")

    result = []
    for audio_path in audio_list:
        out_path = f'{os.path.splitext(audio_path)[0]}_noisereduce.wav'

        try:
            # Load audio file
            audio = AudioSegment.from_file(audio_path)

            # Convert audio to numpy array
            samples = np.array(audio.get_array_of_samples())

            reduced_noise = nr.reduce_noise(samples, sr=audio.frame_rate, prop_decrease=0.6)


            reduced_audio = AudioSegment(
                reduced_noise.tobytes(),
                frame_rate=audio.frame_rate,
                sample_width=audio.sample_width,
                channels=audio.channels
            )


            reduced_audio.export(out_path, format="wav")
            result.append(out_path)

        except Exception as e:
            traceback.print_exc()
            print(f"Error in noise reduction: {str(e)}")
            result.append(audio_path)

    return result

def run(audio_files, file_m, file_index):
    if not audio_files:
        raise ValueError("Please provide an audio file.")

    if isinstance(audio_files, str):
        audio_files = [audio_files]

    try:
        duration_base = librosa.get_duration(filename=audio_files[0])
        print("Duration:", duration_base)
    except Exception as e:
        print(e)

    file_m = os.path.join(model_dir, file_m)
    file_index = os.path.join(model_dir, file_index) if file_index else None

    random_tag = "USER_" + str(random.randint(10000000, 99999999))

    converter.apply_conf(
        tag=random_tag,
        file_model=file_m,
        pitch_algo="rmvpe+",
        pitch_lvl=0,
        file_index=file_index,
        index_influence=0.75,
        respiration_median_filtering=3,
        envelope_ratio=0.25,
        consonant_breath_protection=0.5,
        resample_sr=44100 if audio_files[0].endswith('.mp3') else 0,
    )
    time.sleep(0.1)

    result = convert_now(audio_files, random_tag, converter)
    result = apply_noisereduce(result)

    return result, result[0]  # Assuming result is a list of file paths


def process_audio(audio_files, file_m, file_index):
    result, first_audio = run(audio_files, file_m, file_index)
    
    # Ensure first_audio is correctly formatted as a valid file path
    if not os.path.exists(first_audio):
        raise ValueError(f"Processed audio file not found at: {first_audio}")
    
    return result, gr.update(value=first_audio, visible=True)


def model_conf():
    model_files = [f for f in os.listdir(model_dir) if f.endswith(".pth")]
    return gr.Dropdown(
        label="Select Model File",
        choices=model_files,
        value=model_files[0] if model_files else None,
        interactive=True,
    )

def index_conf():
    index_files = [f for f in os.listdir(model_dir) if f.endswith(".index")]
    return gr.Dropdown(
        label="Select Index File",
        choices=index_files,
        value=index_files[0] if index_files else None,
        interactive=True,
    )

def audio_conf():
    return gr.File(
        label="Audio files",
        file_count="multiple",
        type="filepath",
        container=True,
    )

def button_conf():
    return gr.Button(
        "Inference",
        variant="primary",
    )

def output_conf():
    return gr.File(label="Result", file_count="multiple", interactive=False), gr.Audio(label="Play Result",visible=False,show_share_button=False)

def get_gui(theme):
    with gr.Blocks(theme=theme, delete_cache=(3200, 3200)) as app:
        gr.Markdown(title)

        aud = audio_conf()
        model = model_conf()
        indx = index_conf()
        button_base = button_conf()
        output_file, output_audio = output_conf()

        button_base.click(
            process_audio,
            inputs=[aud, model, indx],
            outputs=[output_file, output_audio],
        )

        gr.Examples(
            examples=[
                ["./test.ogg", "./model.pth", "./model.index"],
                ["./example2/test2.ogg", "./example2/model.pth", "./example2/model.index"],
            ],
            fn=process_audio,
            inputs=[aud, model, indx],
            outputs=[output_file, output_audio],
            cache_examples=False,
        )

    return app

if __name__ == "__main__":
    app = get_gui(theme)
    app.queue(default_concurrency_limit=40)
    app.launch(
        max_threads=40,
        share=False,
        show_error=True,
        quiet=False,
        debug=False,
        allowed_paths=["./downloads/"],
    )