import os
import sys
import traceback
import uuid
from types import ModuleType

# Keep compatibility with dependencies that import lzma in this environment.
class MockLZMA(ModuleType):
    LZMAError = Exception
    FORMAT_AUTO = 0
    FORMAT_XZ = 1
    FORMAT_ALONE = 2
    FORMAT_RAW = 3
    CHECK_NONE = 0
    CHECK_CRC32 = 1
    CHECK_CRC64 = 4
    CHECK_SHA256 = 12

    class LZMAFile:
        def __init__(self, *args, **kwargs):
            raise NotImplementedError("lzma is not supported in this environment")

        def read(self, *args, **kwargs):
            pass

        def write(self, *args, **kwargs):
            pass

        def seek(self, *args, **kwargs):
            pass

        def tell(self, *args, **kwargs):
            pass

    class LZMACompressor:
        def __init__(self, *args, **kwargs):
            raise NotImplementedError("lzma is not supported in this environment")

    class LZMADecompressor:
        def __init__(self, *args, **kwargs):
            raise NotImplementedError("lzma is not supported in this environment")

    def open(self, *args, **kwargs):
        raise NotImplementedError("lzma is not supported in this environment")

    def compress(self, *args, **kwargs):
        raise NotImplementedError("lzma is not supported in this environment")

    def decompress(self, *args, **kwargs):
        raise NotImplementedError("lzma is not supported in this environment")

    def is_check_supported(self, *args, **kwargs):
        return False


sys.modules["lzma"] = MockLZMA("lzma")

import gradio as gr

# Try importing spaces for Hugging Face ZeroGPU.
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=None):
            if fn is None:
                return lambda wrapped: wrapped
            return fn

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import torch

if hasattr(torch, "backends") and hasattr(torch.backends, "mps"):
    torch.backends.mps.is_available = lambda: False

from rvc_engine import RVCEngine

print("Initializing RVC engine...")
rvc_engine = RVCEngine()

TEMP_DIR = "temp_audio"
os.makedirs(TEMP_DIR, exist_ok=True)

VOICE_PRESETS = {
    "Adam": ("GOTHMOMMY.pth", -10),
    "MM": ("BetterMommyGirlfriend.pth", 3),
}


@spaces.GPU(duration=300)
def convert_voice_gpu(input_audio, output_path, model_name, pitch_shift):
    try:
        rvc_engine.convert(input_audio, output_path, model_name, pitch_shift, in_process=True)
    except Exception as e:
        traceback.print_exc()
        return f"{type(e).__name__}: {e}"
    return None


def rvc_handler(input_audio, voice_preset):
    if not input_audio:
        raise gr.Error("Please upload an audio file to convert.")
    if voice_preset not in VOICE_PRESETS:
        raise gr.Error("Please select a voice.")

    model_name, pitch_shift = VOICE_PRESETS[voice_preset]
    if model_name not in rvc_engine.get_available_models():
        raise gr.Error(f"The model for {voice_preset} is not available.")

    output_path = os.path.join(TEMP_DIR, f"rvc_output_{uuid.uuid4().hex}.wav")
    try:
        conversion_error = convert_voice_gpu(
            input_audio, output_path, model_name, pitch_shift
        )
        if conversion_error:
            raise gr.Error(f"RVC conversion failed: {conversion_error}")
        return output_path
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"RVC conversion failed: {e}") from e


custom_css = """
.container {
    max-width: 1000px;
    margin: 0 auto;
    padding: 20px;
}
.header {
    text-align: center;
    margin-bottom: 30px;
    background: linear-gradient(135deg, #4f46e5, #0891b2);
    padding: 30px;
    border-radius: 12px;
    color: white;
    box-shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1), 0 2px 4px -2px rgb(0 0 0 / 0.1);
}
.header h1 {
    font-size: 2.5rem;
    font-weight: 800;
    margin-bottom: 8px;
    color: white !important;
}
.header p {
    font-size: 1.1rem;
    opacity: 0.95;
    margin: 0;
}
"""

theme = gr.themes.Soft(
    primary_hue="stone",
    secondary_hue="yellow",
    font=[gr.themes.GoogleFont("Outfit"), "ui-sans-serif", "system-ui", "sans-serif"],
)

with gr.Blocks(title="The Voice of Mr. Adam", theme=theme, css=custom_css) as demo:
    with gr.Column(elem_classes="container"):
        with gr.Column(elem_classes="header"):
            gr.Markdown("# ☢️ The Voice of Mr. Adam")
            gr.Markdown(
                "<p>If I sent you here from my resume, it is to prove the work is my own despite obvious differences in presentation.</p>"
            )

        gr.Markdown("### Choose an input audio file and a voice.")
        with gr.Row():
            with gr.Column(scale=1):
                rvc_audio = gr.Audio(
                    sources=["upload", "microphone"],
                    type="filepath",
                    label="Input Voice Audio",
                )
                rvc_voice = gr.Radio(
                    choices=["Adam", "MM"],
                    value="Adam",
                    label="Voice",
                )
                rvc_btn = gr.Button("Convert Voice", variant="primary")
            with gr.Column(scale=1):
                rvc_output = gr.Audio(label="Converted Audio Output", type="filepath")

        rvc_btn.click(
            fn=rvc_handler,
            inputs=[rvc_audio, rvc_voice],
            outputs=[rvc_output],
        )

demo.queue()

if __name__ == "__main__":
    demo.launch()