import gc
import os

import gradio as gr
import spaces
import torch
from huggingface_hub import login, snapshot_download

HF_TOKEN = os.environ.get("HF_TOKEN")
MODEL_ID = os.environ.get("MODEL_ID", "sleeper371/sparktts")

if HF_TOKEN:
    login(token=HF_TOKEN)

_codec = None
_tokenizer = None
_model = None
_normalizer = None

# ---------------------------------------------------------------------------
# Audio prompts
# ---------------------------------------------------------------------------

AUDIO_PROMPTS_DIR = os.path.join(os.path.dirname(__file__), "audio_prompts")
_AUDIO_EXTS = {".wav", ".mp3", ".ogg", ".flac", ".m4a"}


def _scan_prompts():
    """Return list of (display_name, filepath) sorted by filename."""
    prompts = []
    if os.path.isdir(AUDIO_PROMPTS_DIR):
        for fname in sorted(os.listdir(AUDIO_PROMPTS_DIR)):
            stem, ext = os.path.splitext(fname)
            if ext.lower() in _AUDIO_EXTS:
                label = stem.replace("_", " ").replace("-", " ").title()
                prompts.append((label, os.path.join(AUDIO_PROMPTS_DIR, fname)))
    return prompts


PRESETS = _scan_prompts()
PRESET_NAMES = [name for name, _ in PRESETS]
PRESET_MAP = {name: path for name, path in PRESETS}

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------


def _load_models():
    global _codec, _tokenizer, _model

    if _codec is None:
        from ncodec.codec import TTSCodec
        from ncodec.decoder.model import AudioDecoder
        from ncodec.encoder.model import AudioEncoder

        try:
            model_path = snapshot_download(MODEL_ID, token=HF_TOKEN)
            decoder_path = os.path.join(model_path, "decoders")
            if not os.path.isdir(decoder_path):
                raise FileNotFoundError
        except Exception:
            mira_path = snapshot_download("YatharthS/MiraTTS")
            decoder_path = os.path.join(mira_path, "decoders")

        _codec = TTSCodec.__new__(TTSCodec)
        _codec.audio_encoder = AudioEncoder(decoder_path)
        _codec.audio_decoder = AudioDecoder(decoder_path)

    if _tokenizer is None:
        from transformers import AutoTokenizer

        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)

    if _model is None:
        from transformers import AutoModelForCausalLM

        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            token=HF_TOKEN,
        )
        _model.eval()

    return _codec, _tokenizer, _model


def _get_eos_token_id(tokenizer):
    for tok in ("<|prompt_speech_end|>", "<|end_of_speech|>", "<|im_end|>"):
        if tok in tokenizer.get_vocab():
            return tokenizer.convert_tokens_to_ids(tok)
    return tokenizer.eos_token_id


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


def _get_normalizer():
    global _normalizer
    if _normalizer is None:
        from soe_vinorm import SoeNormalizer
        _normalizer = SoeNormalizer()
    return _normalizer


@spaces.GPU(duration=120)
def generate(text, ref_audio, temperature, top_p, top_k):
    if not text or not text.strip():
        gr.Warning("Please enter some text.")
        return None
    if ref_audio is None:
        gr.Warning("Please select a preset voice or upload a reference audio file.")
        return None

    codec, tokenizer, model = _load_models()
    normalizer = _get_normalizer()

    normalized_text = normalizer.normalize(text.strip())
    context_tokens = codec.encode(ref_audio)
    prompt = codec.format_prompt(normalized_text, context_tokens, None)
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    eos_id = _get_eos_token_id(tokenizer)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=2048,
            do_sample=True,
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
            repetition_penalty=1.2,
            eos_token_id=eos_id,
        )

    new_ids = output_ids[0][inputs.input_ids.shape[1] :]
    generated_text = tokenizer.decode(new_ids, skip_special_tokens=False)
    audio = codec.decode(generated_text, context_tokens)

    if isinstance(audio, torch.Tensor):
        audio = audio.cpu().float().numpy()

    gc.collect()
    torch.cuda.empty_cache()
    return (48000, audio)


# ---------------------------------------------------------------------------
# UI helpers
# ---------------------------------------------------------------------------


def select_preset(name):
    """Return the filepath for a chosen preset so ref_audio updates."""
    return PRESET_MAP.get(name)


# ---------------------------------------------------------------------------
# Gradio app
# ---------------------------------------------------------------------------

_default_preset = (
    "Mai" if "Mai" in PRESET_MAP else (PRESET_NAMES[0] if PRESET_NAMES else None)
)
_default_ref = PRESET_MAP[_default_preset] if _default_preset else None

TEXT_EXAMPLES = [
    "Mùa xuân sẽ có hoa mơ, hoa mận, trong khi mùa hè mang đến cảnh quan xanh mướt, mát mẻ"
]

with gr.Blocks(title="Vietnamese TTS Demo") as demo:
    gr.Markdown(
        """
        # 🎙️ Vietnamese TTS Demo
        A Vietnamese text-to-speech model based on **SparkTTS**, fine-tuned on Vietnamese speech data.
        Pick a preset voice or upload your own reference clip (3–10 s), enter text, and hit **Generate**.
        """
    )

    with gr.Row():
        # ── Left column ────────────────────────────────────────────────────
        with gr.Column(scale=1):
            # Preset voice picker (only shown when prompts exist)
            if PRESET_NAMES:
                preset_dropdown = gr.Dropdown(
                    choices=PRESET_NAMES,
                    value=_default_preset,
                    label="Preset Voice",
                    interactive=True,
                )

            gr.Markdown("### Reference Audio")
            ref_audio = gr.Audio(
                label="Reference audio (select a preset above or upload your own)",
                value=_default_ref,
                type="filepath",
                interactive=True,
            )

            text_input = gr.Textbox(
                label="Vietnamese Text",
                placeholder="Enter Vietnamese text here...",
                lines=4,
            )

            with gr.Accordion("Advanced Settings", open=False):
                temperature = gr.Slider(
                    0.1, 1.5, value=0.8, step=0.05, label="Temperature"
                )
                top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
                top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")

            generate_btn = gr.Button("🔊 Generate Speech", variant="primary", size="lg")

        # ── Right column ───────────────────────────────────────────────────
        with gr.Column(scale=1):
            output_audio = gr.Audio(label="Generated Speech", type="numpy")

    # Text quick-pick examples
    gr.Examples(
        examples=[[t] for t in TEXT_EXAMPLES],
        inputs=[text_input],
        label="Example Texts",
    )

    gr.Markdown(
        """
        ---
        **Tip:** For best results, use a reference clip that is 3–10 seconds long, clear, and low-noise.
        """
    )

    # Wire preset selection → ref_audio
    if PRESET_NAMES:
        preset_dropdown.change(
            fn=select_preset,
            inputs=preset_dropdown,
            outputs=ref_audio,
        )

    generate_btn.click(
        fn=generate,
        inputs=[text_input, ref_audio, temperature, top_p, top_k],
        outputs=output_audio,
    )

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