import os
import random
import numpy as np
import torch
import gradio as gr  
import spaces
from typing import Optional, Tuple
from pathlib import Path
import tempfile
import soundfile as sf
import time
from datetime import datetime
from soe_vinorm import SoeNormalizer

from assets.ui_strings import UI_STRINGS
from assets.text_samples import TARGET_TEXT_SAMPLES


def log(msg: str):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {msg}")


def setup_cache_env():
    """
    Setup cache environment variables.
    Must be called in GPU worker context as well.
    """
    _cache_home = os.path.join(os.path.expanduser("~"), ".cache")
    
    # HuggingFace cache
    os.environ["HF_HOME"] = os.path.join(_cache_home, "huggingface")
    os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.join(_cache_home, "huggingface", "hub")
    
    # ModelScope cache (for FunASR SenseVoice)
    os.environ["MODELSCOPE_CACHE"] = os.path.join(_cache_home, "modelscope")
    
    # Torch Hub cache (for some audio models like ZipEnhancer)
    os.environ["TORCH_HOME"] = os.path.join(_cache_home, "torch")
    
    # Create cache directories
    for d in [os.environ["HF_HOME"], os.environ["MODELSCOPE_CACHE"], os.environ["TORCH_HOME"]]:
        os.makedirs(d, exist_ok=True)


# Setup cache in main process BEFORE any imports
setup_cache_env()

# Limit thread count to avoid OpenBLAS resource errors in ZeroGPU
os.environ["OPENBLAS_NUM_THREADS"] = "4"
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
if os.environ.get("HF_REPO_ID", "").strip() == "":
    os.environ["HF_REPO_ID"] = "openbmb/VoxCPM1.5"

# Global model cache for ZeroGPU
_voxcpm_model = None

# Global Vietnamese text normalizer (soe-vinorm)
_vn_text_normalizer = SoeNormalizer()

# Fixed local path for VoxCPM model (to avoid repeated downloads in GPU workers)
VOXCPM_LOCAL_DIR = "./models/VoxCPM1.5"


def predownload_models():
    """
    Pre-download models at startup (runs in main process, not GPU worker).
    Download to fixed local directories so GPU workers can reuse them.
    """
    print("=" * 50)
    print("Pre-downloading VoxCPM model to local directory...")
    print("=" * 50)
    
    # Pre-download VoxCPM model to fixed local directory
    if not os.path.isdir(VOXCPM_LOCAL_DIR) or not os.path.exists(os.path.join(VOXCPM_LOCAL_DIR, "model.safetensors")):
        try:
            from huggingface_hub import snapshot_download
            voxcpm_model_id = os.environ.get("HF_REPO_ID", "JayLL13/VoxCPM-1.5-VN-4")
            print(f"Pre-downloading VoxCPM model: {voxcpm_model_id} -> {VOXCPM_LOCAL_DIR}")
            os.makedirs(VOXCPM_LOCAL_DIR, exist_ok=True)
            snapshot_download(
                repo_id=voxcpm_model_id,
                local_dir=VOXCPM_LOCAL_DIR,
                token=os.environ.get("HF_TOKEN"),
            )
            print(f"VoxCPM model downloaded to: {VOXCPM_LOCAL_DIR}")
        except Exception as e:
            print(f"Warning: Failed to pre-download VoxCPM model: {e}")
    else:
        print(f"VoxCPM model already exists at: {VOXCPM_LOCAL_DIR}")
    
    print("=" * 50)
    print("Model pre-download complete!")
    print("=" * 50)


# Run pre-download at startup
predownload_models()


def get_voxcpm_model():
    """Lazy load VoxCPM model (without denoiser)."""
    global _voxcpm_model
    if _voxcpm_model is None:
        import voxcpm
        log("=" * 50)
        log(f"Loading VoxCPM model from: {VOXCPM_LOCAL_DIR}")
        start_time = time.time()
        _voxcpm_model = voxcpm.VoxCPM(
            voxcpm_model_path=VOXCPM_LOCAL_DIR, 
            optimize=False,
            enable_denoiser=False,  # Disable denoiser to avoid ZipEnhancer download
        )
        load_time = time.time() - start_time
        log(f"VoxCPM model loaded. (耗时: {load_time:.2f}s)")
        log("=" * 50)
    return _voxcpm_model


@spaces.GPU(duration=120)
def generate_tts_audio_gpu(
    text_input: str,
    prompt_wav_data: Optional[Tuple[np.ndarray, int]] = None,
    prompt_text_input: Optional[str] = None,
    cfg_value_input: float = 2.0,
    inference_timesteps_input: int = 10,
    do_normalize: bool = True,
    lang: str = "vi",
) -> Tuple[int, np.ndarray]:
    """
    GPU function: Generate speech from text using VoxCPM.
    prompt_wav_data is (audio_array, sample_rate) tuple.
    """
    voxcpm_model = get_voxcpm_model()

    text = (text_input or "").strip()
    if len(text) == 0:
        _lang = "vi" if lang not in UI_STRINGS else lang
        raise ValueError(UI_STRINGS[_lang]["error_no_text"])

    # Optional Vietnamese text normalization using soe-vinorm
    if do_normalize:
        text = _vn_text_normalizer.normalize(text)

    prompt_text = prompt_text_input if prompt_text_input else None
    prompt_wav_path = None

    # If prompt audio data provided, write to temp file for voxcpm
    if prompt_wav_data is not None:
        audio_array, sr = prompt_wav_data
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            sf.write(f.name, audio_array, sr)
            prompt_wav_path = f.name

    try:
        log("=" * 50)
        log("[TTS] 开始语音合成...")
        log(f"[TTS] 目标文本: {text}")
        start_time = time.time()
        wav = voxcpm_model.generate(
            text=text,
            prompt_text=prompt_text,
            prompt_wav_path=prompt_wav_path,
            cfg_value=float(cfg_value_input),
            inference_timesteps=int(inference_timesteps_input),
            normalize=False,  # Normalization handled externally by soe-vinorm
            denoise=False,  # Denoiser disabled
        )
        inference_time = time.time() - start_time
        audio_duration = len(wav) / voxcpm_model.tts_model.sample_rate
        rtf = inference_time / audio_duration if audio_duration > 0 else 0
        log(f"[TTS] 推理耗时: {inference_time:.2f}s | 音频时长: {audio_duration:.2f}s | RTF: {rtf:.3f}")
        log("=" * 50)
        return (voxcpm_model.tts_model.sample_rate, wav)
    finally:
        # Cleanup temp file
        if prompt_wav_path and os.path.exists(prompt_wav_path):
            try:
                os.unlink(prompt_wav_path)
            except Exception:
                pass


def generate_tts_audio(
    text_input: str,
    prompt_wav_path_input: Optional[str] = None,
    prompt_text_input: Optional[str] = None,
    cfg_value_input: float = 2.0,
    inference_timesteps_input: int = 10,
    do_normalize: bool = True,
    lang: str = "vi",
) -> Tuple[int, np.ndarray]:
    """
    Wrapper: Read audio file in CPU, then call GPU function.
    """
    # Ensure prompt_wav_path and prompt_text follow VoxCPM requirement:
    # they must both be provided together or both be None.
    # If user does not provide a reference audio, fall back to default examples/audio.wav.
    if not prompt_wav_path_input:
        default_prompt_wav = "./examples/audio.wav"
        if os.path.exists(default_prompt_wav):
            prompt_wav_path_input = default_prompt_wav

    # If we still don't have any prompt audio, force prompt_text to None.
    # Otherwise, ensure we always have some non-empty prompt text when using a reference audio.
    if not prompt_wav_path_input:
        prompt_text_input = None
    else:
        if not prompt_text_input or not str(prompt_text_input).strip():
            prompt_text_input = "Giọng tham chiếu tiếng Việt mặc định."

    prompt_wav_data = None
    
    # Read audio file before entering GPU context
    if prompt_wav_path_input and os.path.exists(prompt_wav_path_input):
        try:
            audio_array, sr = sf.read(prompt_wav_path_input, dtype='float32')
            prompt_wav_data = (audio_array, sr)
            print(f"Loaded prompt audio: {audio_array.shape}, sr={sr}")
        except Exception as e:
            print(f"Warning: Failed to load prompt audio: {e}")
            prompt_wav_data = None
    
    return generate_tts_audio_gpu(
        text_input=text_input,
        prompt_wav_data=prompt_wav_data,
        prompt_text_input=prompt_text_input,
        cfg_value_input=cfg_value_input,
        inference_timesteps_input=inference_timesteps_input,
        do_normalize=do_normalize,
        lang=lang,
    )


# ---------- UI Builders ----------

def create_demo_interface():
    """Build the Gradio UI for VoxCPM demo."""
    # static assets (logo path)
    try:
        gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"])
    except Exception:
        pass

    with gr.Blocks(
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="gray",
            neutral_hue="slate",
            font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"]
        ),
        css="""
        .logo-container {
            text-align: center;
            margin: 0.5rem 0 1rem 0;
        }
        .logo-container img {
            height: 80px;
            width: auto;
            max-width: 200px;
            display: inline-block;
        }
        /* Bold accordion labels */
        #acc_quick details > summary,
        #acc_tips details > summary {
            font-weight: 600 !important;
            font-size: 1.1em !important;
        }
        /* Bold labels for specific checkboxes */
        #chk_denoise label,
        #chk_denoise span,
        #chk_normalize label,
        #chk_normalize span {
            font-weight: 600;
        }
        .lang-toggle { margin-left: auto; }
        """
    ) as interface:
        s0 = UI_STRINGS["vi"]
        # Header: logo + author (cập nhật theo lang) + nút chuyển ngôn ngữ góc phải
        with gr.Row(elem_id="header_row"):
            with gr.Column(scale=1):
                gr.HTML(
                    '<div class="logo-container"><img src="/gradio_api/file=assets/voxcpm-logo.png" alt="VoxCPM Logo"><img src="/gradio_api/file=assets/vn-flag.gif" alt="Vietnam Flag" style="height:40px; width:auto; margin-left:8px;"></div>'
                )
                author_html = gr.HTML(value=s0["author_line"])
            with gr.Column(scale=0, min_width=120):
                lang_radio = gr.Radio(
                    choices=[("🇻🇳", "vi"), ("🇪🇳", "en")],
                    value="vi",
                    label="Ngôn ngữ",
                    show_label=False,
                    elem_classes=["lang-toggle"],
                )
        # Quick Start
        acc_quick = gr.Accordion(s0["acc_quick"], open=False, elem_id="acc_quick")
        with acc_quick:
            quick_start_md = gr.Markdown(value=s0["quick_md"])

        # Pro Tips
        acc_tips = gr.Accordion(s0["acc_tips"], open=False, elem_id="acc_tips")
        with acc_tips:
            tips_md = gr.Markdown(value=s0["tips_md"])

        with gr.Row():
            with gr.Column():
                prompt_text = gr.Textbox(
                    value="Em chỉ mong được làm một người vợ ngoan và sinh nở được mẹ tròn con vuông.",
                    label=s0["prompt_label"],
                    placeholder=s0["prompt_placeholder"],
                )
                with gr.Row():
                    text = gr.Textbox(
                        value="Trời trưa trong trẻo, trên triền núi trập trùng, lũ trẻ trâu trẻ trung tranh nhau thổi sáo trúc trong trẻo giữa trưa hè đầy trắc trở.",
                        label=s0["target_label"],
                        placeholder=s0["target_placeholder"],
                        lines=4,
                        scale=20,
                    )
                    sample_random_btn = gr.Button("🎲", variant="secondary", scale=1, min_width=52)
                run_btn = gr.Button(s0["btn_generate"], variant="primary")

            with gr.Column():
                reference_audio = gr.Audio(
                    value="./examples/audio.wav" if os.path.exists("./examples/audio.wav") else None,
                    sources=["upload"],
                    type="filepath",
                    label=s0["ref_audio_label"],
                    interactive=True,
                )
                audio_output = gr.Audio(label=s0["output_audio_label"])

        # Advanced config
        acc_advanced = gr.Accordion(s0["acc_advanced"], open=False)
        with acc_advanced:
            cfg_value = gr.Slider(
                minimum=1.0,
                maximum=3.0,
                value=2.0,
                step=0.1,
                label=s0["cfg_label"],
                info=s0["cfg_info"],
            )
            inference_timesteps = gr.Slider(
                minimum=4,
                maximum=30,
                value=10,
                step=1,
                label=s0["timesteps_label"],
                info=s0["timesteps_info"],
            )
            DoNormalizeText = gr.Checkbox(
                value=True,
                label=s0["normalize_label"],
                elem_id="chk_normalize",
                info=s0["normalize_info"],
            )

        def random_sample():
            return random.choice(TARGET_TEXT_SAMPLES)

        sample_random_btn.click(fn=random_sample, inputs=[], outputs=[text])

        # Đổi ngôn ngữ UI khi chọn Tiếng Việt / English
        def apply_lang(lang):
            if lang not in UI_STRINGS:
                lang = "vi"
            s = UI_STRINGS[lang]
            return (
                gr.update(value=s["author_line"]),
                gr.update(label=s["acc_quick"]),
                gr.update(value=s["quick_md"]),
                gr.update(label=s["acc_tips"]),
                gr.update(value=s["tips_md"]),
                gr.update(label=s["prompt_label"], placeholder=s["prompt_placeholder"]),
                gr.update(label=s["target_label"], placeholder=s["target_placeholder"]),
                gr.update(value=s["btn_generate"]),
                gr.update(label=s["ref_audio_label"]),
                gr.update(label=s["output_audio_label"]),
                gr.update(label=s["acc_advanced"]),
                gr.update(label=s["cfg_label"], info=s["cfg_info"]),
                gr.update(label=s["timesteps_label"], info=s["timesteps_info"]),
                gr.update(label=s["normalize_label"], info=s["normalize_info"]),
            )

        lang_radio.change(
            fn=apply_lang,
            inputs=[lang_radio],
            outputs=[
                author_html,
                acc_quick,
                quick_start_md,
                acc_tips,
                tips_md,
                prompt_text,
                text,
                run_btn,
                reference_audio,
                audio_output,
                acc_advanced,
                cfg_value,
                inference_timesteps,
                DoNormalizeText,
            ],
        )

        # Wiring
        run_btn.click(
            fn=generate_tts_audio,
            inputs=[text, reference_audio, prompt_text, cfg_value, inference_timesteps, DoNormalizeText, lang_radio],
            outputs=[audio_output],
            show_progress=True,
            api_name="generate",
        )

    return interface


def run_demo(server_name: str = "0.0.0.0", server_port: int = 7860, show_error: bool = True):
    interface = create_demo_interface()
    # Recommended to enable queue on Spaces for better throughput
    interface.queue(max_size=10).launch(server_name=server_name, server_port=server_port, show_error=show_error)


if __name__ == "__main__":
    run_demo()