import os
import zipfile
import shutil
import gdown
import gradio as gr
import spaces
import torch

from rvc_engine.infer import voice_convert


MODEL_DIR = "models"

os.makedirs(
    MODEL_DIR,
    exist_ok=True
)


print("CUDA:", torch.cuda.is_available())

if torch.cuda.is_available():
    print(
        "GPU:",
        torch.cuda.get_device_name(0)
    )


def download_model(drive_url):

    zip_path = "model.zip"

    if os.path.exists(zip_path):
        os.remove(zip_path)


    try:
        gdown.download(
            drive_url,
            zip_path,
            quiet=False
        )

    except Exception as e:
        return None, None, f"خطأ تحميل الموديل: {str(e)}"


    if not os.path.exists(zip_path):
        return None, None, "فشل تحميل ملف ZIP"


    extract_dir = os.path.join(
        MODEL_DIR,
        "current"
    )


    if os.path.exists(extract_dir):
        shutil.rmtree(extract_dir)


    os.makedirs(
        extract_dir,
        exist_ok=True
    )


    try:

        with zipfile.ZipFile(zip_path, "r") as z:
            z.extractall(extract_dir)

    except Exception as e:

        return None, None, f"خطأ فك الضغط: {str(e)}"



    pth = None
    index = None


    for root, dirs, files in os.walk(extract_dir):

        for file in files:

            if file.endswith(".pth"):
                pth = os.path.join(
                    root,
                    file
                )


            if file.endswith(".index"):
                index = os.path.join(
                    root,
                    file
                )



    if not pth or not index:

        return None, None, (
            "لم يتم العثور على ملفات "
            ".pth و .index"
        )


    return (
        pth,
        index,
        "تم تحميل الموديل بنجاح"
    )



@spaces.GPU
def convert(audio, model_url):


    if not audio:

        return None, "يرجى رفع ملف صوتي"



    if not model_url:

        return None, "ضع رابط Google Drive للموديل"



    pth, index, status = download_model(
        model_url
    )


    if not pth:

        return None, status



    output = "output.wav"



    try:

        result = voice_convert(
            audio,
            pth,
            index,
            output
        )


        return (
            result,
            "تم التحويل بنجاح"
        )


    except Exception as e:

        return (
            None,
            f"خطأ في محرك RVC: {str(e)}"
        )





with gr.Blocks() as demo:


    gr.Markdown(
        "# 🎙️ RVC Voice Converter"
    )


    audio = gr.Audio(
        type="filepath",
        label="الصوت الأصلي"
    )


    model_url = gr.Textbox(
        label="رابط Google Drive للموديل ZIP",
        placeholder="ملف ZIP يحتوي pth و index"
    )


    button = gr.Button(
        "تحويل الصوت"
    )


    output = gr.Audio(
        label="الصوت الناتج"
    )


    info = gr.Textbox(
        label="الحالة"
    )


    button.click(
        convert,
        inputs=[
            audio,
            model_url
        ],
        outputs=[
            output,
            info
        ]
    )



demo.launch()