import spaces
import os
import numpy as np
from huggingface_hub import login
import gradio as gr
from cached_path import cached_path
import tempfile
from vinorm import TTSnorm

from f5_tts.model import DiT
from f5_tts.infer.utils_infer import (
    preprocess_ref_audio_text,
    load_vocoder,
    load_model,
    infer_process,
    save_spectrogram,
)

# Retrieve token from secrets
hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Log in to Hugging Face
if hf_token:
    login(token=hf_token)

def post_process(text: str):
    """
    Chuẩn hóa text trước khi synthesize.
    """
    text = " " + text + " "
    text = text.replace(" . . ", " . ")
    text = text.replace(" .. ", " . ")
    text = text.replace(" , , ", " , ")
    text = text.replace(" ,, ", " , ")
    text = text.replace('"', "")
    return " ".join(text.split()).strip()

def synthesize_with_pauses(ref_audio, ref_text, text, model, vocoder, speed=1.0, volume=1.0, pause_duration=1.0):
    """
    Chia text theo dấu chấm, synthesize từng câu và ghép lại với khoảng im lặng.
    """
    processed_text = post_process(TTSnorm(text)).lower()
    sentences = [s.strip() for s in processed_text.split(".") if s.strip()]

    all_waves = []
    sr = 22050  # sample rate mặc định (cập nhật sau từ infer_process)

    for idx, sentence in enumerate(sentences):
        wave, sr, _ = infer_process(ref_audio, ref_text.lower(), sentence, model, vocoder, speed=speed)
        wave = np.clip(wave * volume, -1.0, 1.0)
        all_waves.append(wave)

        # Thêm im lặng giữa các câu (trừ câu cuối)
        if idx < len(sentences) - 1:
            silence = np.zeros(int(sr * pause_duration), dtype=np.float32)
            all_waves.append(silence)

    if all_waves:
        final_wave = np.concatenate(all_waves)
    else:
        final_wave = np.array([], dtype=np.float32)
    return final_wave, sr

# Load models
vocoder = load_vocoder()
model = load_model(
    DiT,
    dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
    ckpt_path=str(cached_path("hf://hynt/F5-TTS-Vietnamese-ViVoice/model_last.pt")),
    vocab_file=str(cached_path("hf://hynt/F5-TTS-Vietnamese-ViVoice/config.json")),
)

@spaces.GPU
def infer_tts(ref_audio_orig: str, gen_text: str, speed: float = 1.0, volume: float = 1.0, pause: float = 1.0, request: gr.Request = None):
    if not ref_audio_orig:
        raise gr.Error("Please upload a sample audio file.")
    if not gen_text.strip():
        raise gr.Error("Please enter the text content to generate voice.")
    if len(gen_text.split()) > 1000:
        raise gr.Error("Please enter text content with less than 1000 words.")

    try:
        # Tiền xử lý sample voice
        ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, "")

        # Synthesize với ngắt nghỉ
        final_wave, final_sample_rate = synthesize_with_pauses(
            ref_audio, ref_text, gen_text, model, vocoder, speed=speed, volume=volume, pause_duration=pause
        )

        # Tạo spectrogram (dùng đoạn text đầy đủ để hiển thị, nhưng không tái synthesize)
        _, _, spectrogram = infer_process(ref_audio, ref_text.lower(), post_process(TTSnorm(gen_text)).lower(), model, vocoder, speed=speed)

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
            spectrogram_path = tmp_spectrogram.name
            save_spectrogram(spectrogram, spectrogram_path)

        return (final_sample_rate, final_wave), spectrogram_path
    except Exception as e:
        raise gr.Error(f"Error generating voice: {e}")

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎤 Chương trình chuyển đổi text thành giọng nói.
    """)
    
    with gr.Row():
        ref_audio = gr.Audio(label="🔊 Sample Voice", type="filepath")
        gen_text = gr.Textbox(label="📝 Text", placeholder="Enter the text to generate voice...", lines=3)
    
    speed = gr.Slider(0.3, 2.0, value=0.95, step=0.01, label="⚡ Speed")
    volume = gr.Slider(0.1, 2.0, value=1.0, step=0.01, label="🔊 Volume")
    pause = gr.Slider(0.0, 3.0, value=1.1, step=0.01, label="⏸ Pause between sentences (seconds)")
    
    btn_synthesize = gr.Button("🔥 Generate Voice")
    
    with gr.Row():
        output_audio = gr.Audio(label="🎧 Generated Audio", type="numpy")
        output_spectrogram = gr.Image(label="📊 Spectrogram")
    
    model_limitations = gr.Textbox(
        value="""1. This model may not perform well with numerical characters, dates, special characters, etc. => A text normalization module is needed.
2. The rhythm of some generated audios may be inconsistent or choppy => It is recommended to select clearly pronounced sample audios with minimal pauses for better synthesis quality.
3. Default, reference audio text uses the pho-whisper-medium model, which may not always accurately recognize Vietnamese, resulting in poor voice synthesis quality.
4. Inference with overly long paragraphs may produce poor results.""", 
        label="❗ Model Limitations",
        lines=4,
        interactive=False
    )

    btn_synthesize.click(
        infer_tts,
        inputs=[ref_audio, gen_text, speed, volume, pause],
        outputs=[output_audio, output_spectrogram]
    )

# Run Gradio
demo.queue().launch()