import gradio as gr
from cached_path import cached_path
import tempfile
import os

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,
)

print("🧠 Memuat Vocoder di CPU...")
vocoder = load_vocoder()

print("🧠 Memuat Model F5-TTS INDO V2 di CPU...")
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://Eempostor/F5-TTS-INDO-FINETUNE-V2/f5_tts_indo_v2.pt")
    ),
    vocab_file=str(
        cached_path("hf://Eempostor/F5-TTS-INDO-FINETUNE-V2/vocab.txt")
    ),
    device="cpu" # 👈 Paksa masuk CPU
)
print("✅ Model berhasil dimuat tanpa error!")

# ❌ Dekorator @spaces.GPU SUDAH DIHAPUS KARENA PAKE CPU

def infer(ref_audio_orig: str, ref_text: str, gen_text: str, speed: float = 1.0):
    if not ref_audio_orig:
        raise gr.Error("Bro, upload dulu file audio referensinya (Voice)!")

    if not gen_text or gen_text.strip() == "":
        raise gr.Error("Teks yang mau di-generate belum diisi nih.")

    try:
        print(f"🚀 Memproses TTS di CPU: {gen_text}")
        
        ref_audio, processed_ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text)
        
        # ⚡ NFE step diturunkan ke 16 agar CPU gratisan nggak ngos-ngosan
        final_wave, final_sample_rate, combined_spectrogram = infer_process(
            ref_audio,
            processed_ref_text,
            gen_text,
            model,
            vocoder,
            cross_fade_duration=0.15,
            nfe_step=16, 
            speed=speed,
        )

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

        print("✅ Generate selesai!")
        return (final_sample_rate, final_wave), spectrogram_path
    
    except Exception as e:
        raise gr.Error(f"Terjadi kesalahan saat inferensi: {e}")

# ==============================================================================
# INTERFACE GRADIO
# ==============================================================================
with gr.Blocks(title="F5-TTS Indonesia Ready") as app:
    gr.Markdown("# 🇮🇩 F5-TTS Indonesia Finetune V2 (CPU Mode)")
    gr.Markdown("Berjalan di mode CPU. Waktu generate akan memakan waktu lebih lama (sekitar 30-60 detik).")
    
    with gr.Row():
        with gr.Column():
            voice = gr.Audio(type="filepath", label="1. Upload Ref Audio (Maks 10-15 dtk)")
            ref_text = gr.Textbox(
                label="2. Transkrip dari Audio Ref (Opsional)", 
                lines=2, 
                placeholder="Bisa dikosongkan. Kalau kosong, sistem otomatis pakai Whisper buat deteksi teksnya!"
            )
            speed = gr.Slider(
                label="Kecepatan Bicara (Speed)",
                minimum=0.3,
                maximum=2.0,
                value=1.0,
                step=0.1,
            )
        with gr.Column():
            text = gr.Textbox(label="3. Teks yang ingin di-Generate", lines=6, placeholder="Ketik kalimat barunya di sini...")
            
    btn = gr.Button("Generate Audio", variant="primary")
    
    with gr.Row():
        audio_out = gr.Audio(type="numpy", label="Output Audio")
        spectrogram_out = gr.Image(type="filepath", label="Spectrogram (Opsional)")
    
    btn.click(
        fn=infer, 
        inputs=[voice, ref_text, text, speed], 
        outputs=[audio_out, spectrogram_out]
    )

if __name__ == "__main__":
    # Paksa jalan di Docker host lokal
    app.queue().launch(server_name="0.0.0.0", server_port=7860)