import os
import random
import numpy as np
import torch
import librosa
import soundfile as sf
import gradio as gr

try:
    import spaces
    GPU_DECORATOR = spaces.GPU(duration=60)
except ImportError:
    def GPU_DECORATOR(func):
        return func

from chatterbox.tts import ChatterboxTTS
from chatterbox.vc import ChatterboxVC

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Initializing Chatterbox on {DEVICE}...")

tts_model = ChatterboxTTS.from_pretrained("cpu")
vc_model = ChatterboxVC.from_pretrained("cpu")

if hasattr(tts_model, "t3") and hasattr(tts_model.t3, "tfmr"):
    if hasattr(tts_model.t3.tfmr, "config"):
        tts_model.t3.tfmr.config._attn_implementation = "eager"

def sanitize_reference_audio(audio_path, max_duration=20.0):
    if not audio_path:
        return None
    wav, sr = librosa.load(audio_path, sr=24000, mono=True)
    max_samples = int(max_duration * sr)
    if len(wav) > max_samples:
        wav = wav[:max_samples]
    max_val = np.max(np.abs(wav))
    if max_val > 0:
        wav = (wav / max_val) * 0.95
    clean_path = "/tmp/clean_ref_prompt.wav"
    sf.write(clean_path, wav, sr, subtype="PCM_16")
    return clean_path

@GPU_DECORATOR
def generate_tts(text, audio_prompt_path, exaggeration, temperature, seed_num, cfgw, min_p, top_p, repetition_penalty):
    text = text.strip()
    if not text:
        raise gr.Error("Please enter text to synthesize.")
    if len(text) > 450:
        text = text[:450]

    target_dev = "cuda" if torch.cuda.is_available() else "cpu"
    tts_model.device = target_dev
    tts_model.t3.to(target_dev)
    tts_model.s3gen.to(target_dev)
    tts_model.ve.to(target_dev)
    if tts_model.conds is not None:
        tts_model.conds.to(target_dev)

    if int(seed_num) != 0:
        torch.manual_seed(int(seed_num))
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(int(seed_num))
        random.seed(int(seed_num))
        np.random.seed(int(seed_num))

    clean_prompt = sanitize_reference_audio(audio_prompt_path)

    with torch.inference_mode():
        wav = tts_model.generate(
            text,
            audio_prompt_path=clean_prompt,
            exaggeration=float(exaggeration),
            temperature=float(temperature),
            cfg_weight=float(cfgw),
            min_p=float(min_p),
            top_p=float(top_p),
            repetition_penalty=float(repetition_penalty),
        )
    return (tts_model.sr, wav.squeeze(0).cpu().numpy())

@GPU_DECORATOR
def convert_voice(input_audio, target_voice_audio):
    if input_audio is None or target_voice_audio is None:
        raise gr.Error("Please upload both source audio and target voice audio.")

    target_dev = "cuda" if torch.cuda.is_available() else "cpu"
    vc_model.device = target_dev
    vc_model.s3gen.to(target_dev)

    clean_src = sanitize_reference_audio(input_audio, max_duration=30.0)
    clean_tgt = sanitize_reference_audio(target_voice_audio, max_duration=20.0)

    with torch.inference_mode():
        wav = vc_model.generate(
            audio=clean_src,
            target_voice_path=clean_tgt,
        )
    return (vc_model.sr, wav.squeeze(0).cpu().numpy())

with gr.Blocks(title="Chatterbox Web Studio") as demo:
    gr.Markdown("# 🎙️ Chatterbox Standalone Web Studio")
    with gr.Tabs():
        with gr.TabItem("🗣️ Text-to-Speech (TTS)"):
            with gr.Row():
                with gr.Column():
                    text = gr.Textbox(label="Text to Synthesize", lines=5, value="Welcome to Chatterbox Web Studio!")
                    ref_wav = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Reference Voice Audio (Optional - 10 to 20s WAV)", value=None)
                    with gr.Row():
                        exaggeration = gr.Slider(0.25, 2.0, step=0.05, value=0.5, label="Exaggeration")
                        cfg_weight = gr.Slider(0.0, 1.0, step=0.05, value=0.5, label="CFG / Pace Control")
                    with gr.Accordion("⚙️ Advanced Sampler Settings", open=False):
                        seed_num = gr.Number(value=0, label="Random Seed", precision=0)
                        temp = gr.Slider(0.05, 2.0, step=0.05, value=0.8, label="Temperature")
                        min_p = gr.Slider(0.00, 1.00, step=0.01, value=0.05, label="min_p")
                        top_p = gr.Slider(0.00, 1.00, step=0.01, value=1.00, label="top_p")
                        rep_penalty = gr.Slider(1.00, 2.00, step=0.05, value=1.20, label="Repetition Penalty")
                    tts_btn = gr.Button("Generate Speech ⚡", variant="primary")
                with gr.Column():
                    tts_output = gr.Audio(label="Generated Output", interactive=False)
            tts_btn.click(fn=generate_tts, inputs=[text, ref_wav, exaggeration, temp, seed_num, cfg_weight, min_p, top_p, rep_penalty], outputs=tts_output)

        with gr.TabItem("🔄 Voice Conversion (VC)"):
            with gr.Row():
                with gr.Column():
                    src_audio = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Source Audio")
                    tgt_audio = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Target Voice Audio")
                    vc_btn = gr.Button("Convert Voice 🔄", variant="primary")
                with gr.Column():
                    vc_output = gr.Audio(label="Converted Audio Output", interactive=False)
            vc_btn.click(fn=convert_voice, inputs=[src_audio, tgt_audio], outputs=vc_output)

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