import os
import uuid
from pathlib import Path
import numpy as np

import gradio as gr
import soundfile as sf
import spaces
from voxcpm import VoxCPM

MODEL_ID = "openbmb/VoxCPM2"
OUTPUT_DIR = Path("outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

_model = None


def get_model():
    global _model
    if _model is None:
        try:
            _model = VoxCPM.from_pretrained(
                MODEL_ID,
                load_denoiser=False,
                optimize=False,
            )
        except TypeError:
            _model = VoxCPM.from_pretrained(
                MODEL_ID,
                load_denoiser=False,
            )
    return _model


def clean_text(text: str) -> str:
    return (
        (text or "")
        .replace("\u200b", "")
        .replace("\u200c", "")
        .replace("\u200d", "")
        .replace("\ufeff", "")
        .replace("\r\n", "\n")
        .replace("\r", "\n")
        .strip()
    )


@spaces.GPU(duration=120)
def synthesize(target_text: str, reference_audio: str, reference_text: str):
    target_text = clean_text(target_text)
    reference_text = clean_text(reference_text)

    if not target_text:
        raise gr.Error("Target text is required.")
    if not reference_audio:
        raise gr.Error("Reference audio is required.")
    if not reference_text:
        raise gr.Error("Reference transcript is required.")

    model = get_model()
    sample_rate = model.tts_model.sample_rate

    # VoxCPM2 Core Generation
    wav = model.generate(
        text=target_text,
        prompt_wav_path=reference_audio,
        prompt_text=reference_text,
        reference_wav_path=reference_audio,
        cfg_value=2.0,
        inference_timesteps=10,
        normalize=True,
        denoise=False,
    )

    # Safe convert to numpy array
    if hasattr(wav, "detach"):
        wav = wav.detach().cpu().numpy()
    wav = np.asarray(wav, dtype=np.float32).squeeze()

    # 🔥 အစဦးပိုင်း ကလစ်သံ (Click/Pop Artifact) သီးသန့်ကို စစ်ထုတ်ဖျက်ဆီးမည့် Logic အသစ်
    search_limit = int(0.35 * sample_rate)  # ထိပ်ဆုံး 350ms အတွင်း ရှာဖွေပါမည်
    if len(wav) > search_limit:
        abs_wav = np.abs(wav[:search_limit])
        max_val = np.max(abs_wav) if len(abs_wav) > 0 else 1.0
        threshold = max(0.012, max_val * 0.05)
        
        # ၁၀ မီလီစက္ကန့်စာ Frame ကွက်များဖြင့် စကားသံ အစစ်အမှန်ကို ရှာဖွေခြင်း
        frame_len = int(0.01 * sample_rate)  
        start_idx = 0
        
        for i in range(0, search_limit - (frame_len * 2), frame_len):
            current_frame = abs_wav[i : i + frame_len]
            # တစ်ချက်တည်းထွက်တဲ့ ကလစ်သံမဟုတ်ကြောင်း သေချာစေရန် နောက်ထပ် Frame ကိုပါ ပူးတွဲစစ်ဆေးပါသည်
            next_frame = abs_wav[i + frame_len : i + (frame_len * 2)]
            
            if np.mean(current_frame) > threshold and np.mean(next_frame) > threshold:
                start_idx = i
                break
                
        if start_idx > 0:
            # စကားသံ အစဦးပိုင်း ပြတ်မသွားစေရန် 5ms မျှ Cushion ချန်လှပ်၍ ညှပ်ထုတ်ပါသည်
            cushion = int(0.005 * sample_rate)
            start_idx = max(0, start_idx - cushion)
            print(f"[Info] Successfully stripped initial pop/click artifact at index: {start_idx}")
            wav = wav[start_idx:]

    out_path = OUTPUT_DIR / f"voxcpm2_{uuid.uuid4().hex}.wav"
    sf.write(str(out_path), wav, sample_rate)

    return str(out_path)


with gr.Blocks(title="Myanmar VoxCPM2 Voice Clone") as demo:
    gr.Markdown("## Myanmar VoxCPM2 Voice Clone")

    target_text = gr.Textbox(
        label="Target text",
        lines=5,
        placeholder="ဒီမှာ ပြောစေချင်တဲ့ မြန်မာစာသားကို ရိုက်ထည့်ပါ။",
    )

    reference_audio = gr.Audio(
        label="Reference audio",
        sources=["upload", "microphone"],
        type="filepath",
    )

    reference_text = gr.Textbox(
        label="Reference transcript",
        lines=3,
        placeholder="Reference audio ထဲမှာ ပြောထားတဲ့ စကားကို တိတိကျကျ ရိုက်ထည့်ပါ။",
    )

    generate_btn = gr.Button("Generate", variant="primary")

    output_audio = gr.Audio(
        label="Generated voice",
        type="filepath",
    )

    generate_btn.click(
        fn=synthesize,
        inputs=[target_text, reference_audio, reference_text],
        outputs=output_audio,
        api_name="synthesize",
    )


if __name__ == "__main__":
    demo.queue(max_size=10).launch()