# ruff: noqa: E402

import os
import tempfile
from importlib.resources import files

import gradio as gr
import numpy as np
import soundfile as sf
import torch
import torchaudio
from cached_path import cached_path

try:
    import spaces

    USING_SPACES = True
except ImportError:
    USING_SPACES = False


def gpu_decorator(func):
    if USING_SPACES:
        return spaces.GPU(func)
    else:
        return func


from f5_tts.infer.utils_infer import (
    load_model,
    load_vocoder,
    preprocess_ref_audio_text,
    remove_silence_for_generated_wav,
    tempfile_kwargs,
)
from f5_tts.model import DiT

from habibi_tts.infer.utils_infer import infer_process
from habibi_tts.model.utils import dialect_id_map


vocoder = load_vocoder()

# Model config
v1_base_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)

# SAU Specialized model (default)
sau_vocab_path = str(cached_path("hf://SWivid/Habibi-TTS/Specialized/SAU/vocab.txt"))
sau_model_path = str(cached_path("hf://SWivid/Habibi-TTS/Specialized/SAU/model_200000.safetensors"))

# Unified model paths (lazy loaded if needed)
uni_vocab_path = str(cached_path("hf://SWivid/Habibi-TTS/Unified/vocab.txt"))
uni_model_path = str(cached_path("hf://SWivid/Habibi-TTS/Unified/model_200000.safetensors"))

# Model cache
current_model = None
current_model_type = None


def get_model(model_type="Specialized"):
    global current_model, current_model_type
    if current_model is not None and current_model_type == model_type:
        return current_model

    if model_type == "Unified":
        current_model = load_model(DiT, v1_base_cfg, uni_model_path, vocab_file=uni_vocab_path)
    else:
        current_model = load_model(DiT, v1_base_cfg, sau_model_path, vocab_file=sau_vocab_path)
    current_model_type = model_type
    return current_model


# Reference audio paths
_assets_path = str(files("habibi_tts").joinpath("assets"))
najdi_audio = str(files("habibi_tts").joinpath("assets/Najdi.wav"))
hijazi_audio = str(files("habibi_tts").joinpath("assets/Hijazi.wav"))
gulf_audio = str(files("habibi_tts").joinpath("assets/Gulf.wav"))

najdi_text = "تكفى طمني انا اليوم ماني بنايم ولا هو بداخل عيني النوم الين اتطمن عليه."
hijazi_text = "ابغاك تحقق معاه بس بشكل ودي لانه سلطان يمر بظروف صعبة شوية."
gulf_text = "وين تو الناس متى تصحى ومتى تفطر وتغير يبيلك ساعة يعني بالله تروح الشغل الساعة عشره."


@gpu_decorator
def infer_speech(
    model_choice,
    ref_audio_orig,
    ref_text,
    gen_text,
    remove_silence,
    speed,
):
    if not ref_audio_orig:
        gr.Warning("⚠️ الرجاء رفع صوت مرجعي أو اختيار مثال.")
        return None, ref_text

    if not gen_text or not gen_text.strip():
        gr.Warning("⚠️ الرجاء كتابة النص المراد تحويله.")
        return None, ref_text

    if not ref_text or not ref_text.strip():
        ref_text = ""

    seed = np.random.randint(0, 2**31 - 1)
    torch.manual_seed(seed)

    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=gr.Info)

    model = get_model(model_choice)

    if model_choice == "Specialized":
        dialect_id = None
    else:
        dialect_id = dialect_id_map["SAU"]

    final_wave, final_sample_rate, _ = infer_process(
        ref_audio,
        ref_text,
        gen_text,
        model,
        vocoder,
        cross_fade_duration=0.15,
        nfe_step=32,
        speed=speed,
        show_info=gr.Info,
        progress=gr.Progress(),
        dialect_id=dialect_id,
    )

    if remove_silence:
        with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f:
            temp_path = f.name
        try:
            sf.write(temp_path, final_wave, final_sample_rate)
            remove_silence_for_generated_wav(temp_path)
            final_wave, _ = torchaudio.load(temp_path)
        finally:
            os.unlink(temp_path)
        final_wave = final_wave.squeeze().cpu().numpy()

    return (final_sample_rate, final_wave), ref_text


def load_example(audio_path, text):
    """Load a reference example into the input fields."""
    return audio_path, text


# Custom CSS
custom_css = """
.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}
.main-header {
    text-align: center;
    padding: 20px 0;
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%);
    border-radius: 16px;
    margin-bottom: 24px;
    border: 1px solid rgba(255,255,255,0.1);
}
.main-header h1 {
    color: #e2e8f0 !important;
    font-size: 2em !important;
    margin-bottom: 8px !important;
}
.main-header p {
    color: #94a3b8 !important;
    font-size: 1.05em !important;
}
.example-btn {
    min-height: 70px !important;
    border: 2px solid rgba(99, 102, 241, 0.3) !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    background: rgba(99, 102, 241, 0.05) !important;
}
.example-btn:hover {
    border-color: rgba(99, 102, 241, 0.7) !important;
    background: rgba(99, 102, 241, 0.15) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 12px rgba(99, 102, 241, 0.2) !important;
}
.generate-btn {
    min-height: 56px !important;
    font-size: 1.2em !important;
    border-radius: 12px !important;
    background: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%) !important;
    border: none !important;
    transition: all 0.3s ease !important;
}
.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(99, 102, 241, 0.4) !important;
}
.output-audio {
    border: 2px solid rgba(99, 102, 241, 0.2) !important;
    border-radius: 12px !important;
    padding: 16px !important;
}
.settings-row {
    padding: 8px 0;
}
footer { display: none !important; }
"""

# Build Gradio UI
with gr.Blocks(
    title="Habibi TTS - اللهجة السعودية",
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="purple",
        neutral_hue="slate",
    ),
    css=custom_css,
) as app:

    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🎙️ Habibi TTS</h1>
        <p>تحويل النص إلى كلام باللهجة السعودية (نجدي · حجازي · خليجي)</p>
    </div>
    """)

    with gr.Row(equal_height=False):
        # Left Column - Inputs
        with gr.Column(scale=1):
            gr.Markdown("### 🎯 اختر مثال صوتي")
            gr.Markdown("اضغط على أحد الأمثلة لتعبئة الصوت المرجعي والنص تلقائياً:")

            with gr.Row():
                btn_najdi = gr.Button("🏜️ نجدي", elem_classes=["example-btn"], size="lg")
                btn_hijazi = gr.Button("🕌 حجازي", elem_classes=["example-btn"], size="lg")
                btn_gulf = gr.Button("🌊 خليجي", elem_classes=["example-btn"], size="lg")

            gr.Markdown("---")

            ref_wav_input = gr.Audio(
                label="🎧 الصوت المرجعي",
                type="filepath",
            )
            ref_txt_input = gr.Textbox(
                label="📝 النص المرجعي (نص الصوت المرجعي)",
                placeholder="يتم تعبئته تلقائياً عند اختيار مثال، أو اكتب نص الصوت المرجعي هنا...",
                lines=3,
                rtl=True,
            )

        # Right Column - Output & Generate
        with gr.Column(scale=1):
            gr.Markdown("### ✍️ أدخل النص المراد تحويله")
            gen_txt_input = gr.Textbox(
                label="النص المراد تحويله إلى كلام",
                placeholder="اكتب هنا النص الذي تريد تحويله إلى صوت باللهجة السعودية...",
                lines=5,
                rtl=True,
            )

            with gr.Row(elem_classes=["settings-row"]):
                model_choice = gr.Radio(
                    choices=["Specialized", "Unified"],
                    label="نوع النموذج",
                    value="Specialized",
                    info="Specialized = متخصص للسعودية (أدق) | Unified = موحد لكل اللهجات",
                )

            with gr.Row(elem_classes=["settings-row"]):
                speed_slider = gr.Slider(
                    label="⚡ السرعة",
                    minimum=0.5,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                )
                remove_silence_check = gr.Checkbox(
                    label="✂️ إزالة الصمت",
                    value=False,
                )

            generate_btn = gr.Button(
                "🎵 توليد الصوت",
                variant="primary",
                size="lg",
                elem_classes=["generate-btn"],
            )

            audio_output = gr.Audio(
                label="🔊 الصوت المُولَّد",
                type="numpy",
                elem_classes=["output-audio"],
            )

    # Example button click handlers
    btn_najdi.click(
        fn=lambda: (najdi_audio, najdi_text),
        outputs=[ref_wav_input, ref_txt_input],
    )
    btn_hijazi.click(
        fn=lambda: (hijazi_audio, hijazi_text),
        outputs=[ref_wav_input, ref_txt_input],
    )
    btn_gulf.click(
        fn=lambda: (gulf_audio, gulf_text),
        outputs=[ref_wav_input, ref_txt_input],
    )

    # Generate button handler
    generate_btn.click(
        infer_speech,
        inputs=[
            model_choice,
            ref_wav_input,
            ref_txt_input,
            gen_txt_input,
            remove_silence_check,
            speed_slider,
        ],
        outputs=[audio_output, ref_txt_input],
    )

    gr.Markdown(
        """
        ---
        <center>
        <small>
        مبني على <a href="https://github.com/SWivid/Habibi-TTS">Habibi-TTS</a>
        · استخدم صوت مرجعي أقل من 12 ثانية · تأكد أن النص المرجعي يتطابق مع الصوت
        </small>
        </center>
        """
    )

app.queue().launch(show_error=True, allowed_paths=[_assets_path])