import gradio as gr
import spaces
import io
import re
import threading
import numpy as np
import soundfile as sf
import torch
import hashlib
import random
import tempfile
import os
import uuid
from huggingface_hub import snapshot_download
from voxcpm import VoxCPM
import whisper

# ----------------------------------------------------
# 1. Download and Auto-Convert Safetensors
# ----------------------------------------------------
try:
    snapshot_dir = snapshot_download(repo_id="openbmb/VoxCPM2")
    bin_path = os.path.join(snapshot_dir, "pytorch_model.bin")
    safetensors_path = os.path.join(snapshot_dir, "model.safetensors")
    if not os.path.exists(bin_path) and os.path.exists(safetensors_path):
        from safetensors.torch import load_file
        state_dict = load_file(safetensors_path)
        torch.save({"state_dict": state_dict}, bin_path)
except Exception as e:
    pass

vox_model = None
whisper_model = None

MAX_CHUNK_CHARS = 180
lock = threading.Lock()

def split_text(text, max_chars=MAX_CHUNK_CHARS):
    text = re.sub(r"\s+", " ", text.strip())
    parts = re.split(r"(?<=[\u17d4.!?])\s*", text)
    chunks = []
    current = ""
    for part in parts:
        if not part: continue
        if len(current) + len(part) + 1 <= max_chars:
            current = (current + " " + part).strip()
        else:
            if current: chunks.append(current)
            current = part
    if current: chunks.append(current)
    return chunks

# ----------------------------------------------------
# 2. ZeroGPU Intercept Functions
# ----------------------------------------------------

# Added duration=60 to bypass the 90s default request limit
@spaces.GPU(duration=60)
def generate_voice(text, voice_id, gender, ref_audio_path, cfg_value, inference_timesteps):
    global vox_model
    if vox_model is None:
        vox_model = VoxCPM.from_pretrained("openbmb/VoxCPM2", load_denoiser=False, optimize=False)

    if gender and gender != "Auto":
        voice_id = f"{gender}_{voice_id}"

    if voice_id:
        seed = int(hashlib.md5(voice_id.encode('utf-8')).hexdigest()[:8], 16)
    else:
        seed = 42

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        
    chunks = split_text(text)
    wavs = []
    
    with lock:
        for chunk in chunks:
            final_text = chunk if ref_audio_path else (f"({voice_id}){chunk}" if voice_id else chunk)
            wav = vox_model.generate(
                text=final_text,
                reference_wav_path=ref_audio_path,
                cfg_value=float(cfg_value),
                inference_timesteps=int(inference_timesteps),
                normalize=True,
                denoise=False,
                retry_badcase=True,
            )
            wavs.append(np.asarray(wav, dtype=np.float32))

    full_wav = np.concatenate(wavs)
    temp_result = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    sf.write(temp_result.name, full_wav, vox_model.tts_model.sample_rate, format="WAV")
    temp_result.close()
    
    return temp_result.name


# Added duration=30 for transcription
@spaces.GPU(duration=30)
def transcribe_audio(audio_path):
    global whisper_model
    if whisper_model is None:
        whisper_model = whisper.load_model("base")
    result = whisper_model.transcribe(audio_path)
    return result["text"].strip()

def check_connection():
    return '{"ok": true, "model": "VoxCPM2", "mode": "ZeroGPU"}'

# ----------------------------------------------------
# 3. Gradio Web Interface (For ZeroGPU detection)
# ----------------------------------------------------

with gr.Blocks() as demo:
    gr.Markdown("# VoxCPM2 ZeroGPU Backend API")
    
    with gr.Tab("TTS"):
        text = gr.Textbox(label="Text")
        voice_id = gr.Textbox(label="Voice ID")
        gender = gr.Textbox(label="Gender")
        ref_audio = gr.Audio(type="filepath", label="Reference Audio")
        cfg_value = gr.Number(value=1.8, label="CFG Value")
        timesteps = gr.Number(value=8, label="Timesteps")
        tts_out = gr.Audio(type="filepath", label="Output")
        tts_btn = gr.Button("Generate")
        
        tts_btn.click(generate_voice, 
                      inputs=[text, voice_id, gender, ref_audio, cfg_value, timesteps], 
                      outputs=[tts_out], 
                      api_name="tts")

    with gr.Tab("Transcribe"):
        trans_in = gr.Audio(type="filepath")
        trans_out = gr.Textbox()
        trans_btn = gr.Button("Transcribe")
        
        trans_btn.click(transcribe_audio, 
                        inputs=[trans_in], 
                        outputs=[trans_out], 
                        api_name="transcribe")
        
    with gr.Tab("Status"):
        status_out = gr.Textbox()
        status_btn = gr.Button("Check")
        
        status_btn.click(check_connection, 
                         inputs=[], 
                         outputs=[status_out], 
                         api_name="status")

demo.launch()