import os
import torch
import torchaudio as ta
import gradio as gr
import numpy as np
import spaces
from huggingface_hub import snapshot_download
from safetensors.torch import load_file as load_safetensors
from chatterbox import mtl_tts

# 1. Configuration & Model Loading
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
REPO_ID = "NAMAA-Space/NAMAA-Egyptian-TTS"

print(f"🚀 Loading model assets...")

# Download model assets
ckpt_dir = snapshot_download(
    repo_id=REPO_ID,
    repo_type="model",
    revision="main"
)

# Initialize the base Chatterbox model
# We load the base model on CPU first, then move components to GPU inside the decorated function if needed,
# but for ZeroGPU, we can initialize here and the @spaces.GPU will handle the device mapping.
model = mtl_tts.ChatterboxMultilingualTTS.from_pretrained(device=DEVICE)

# Load the specific Egyptian TTS weights (t3_mtl23ls_v2.safetensors)
t3_state = load_safetensors(
    os.path.join(ckpt_dir, "t3_mtl23ls_v2.safetensors"),
    device=DEVICE
)
model.t3.load_state_dict(t3_state)
model.t3.to(DEVICE).eval()

print("✅ Model assets loaded successfully!")

@spaces.GPU
def inference(text):
    """
    Performs TTS inference using the NAMAA Egyptian model with ZeroGPU support.
    """
    if not text.strip():
        return None
    
    try:
        # Ensure model is on the correct device assigned by ZeroGPU
        # model.to(DEVICE) # spaces.GPU handles this usually, but safe to ensure
        
        # Generate waveform
        # language_id="ar" is required for Egyptian Arabic processing
        wav = model.generate(text, language_id="ar")
        
        # Save to a temporary file
        output_path = "output.wav"
        ta.save(output_path, wav, model.sr)
        
        return output_path
    except Exception as e:
        print(f"Error during inference: {e}")
        return None

# 2. Gradio Interface Design
title = "NAMAA Egyptian TTS 🇪🇬"
description = """
<center>
هذا النموذج مقدم من مجتمع <b>NAMAA</b> لتحويل النص المكتوب باللهجة المصرية إلى صوت طبيعي.
<br>
يستخدم هذا الـ Space تقنية <b>ZeroGPU</b> لتوفير معالجة سريعة ومجانية.
<br>
This model is developed by <b>NAMAA Community</b> for Egyptian Arabic Text-to-Speech.
</center>
"""

examples = [
    ["انا سبت الشغل و راجع دلوقتي علي طول."],
    ["ازيك يا صاحبي، عامل ايه النهاردة؟"],
    ["مصر بلد جميلة جداً وفيها حاجات كتير تتحب."],
    ["ممكن تشرحلي الموضوع ده ببساطة؟"]
]

demo = gr.Interface(
    fn=inference,
    inputs=gr.Textbox(
        label="أدخل النص باللهجة المصرية (Input Text)", 
        placeholder="اكتب هنا...", 
        lines=3
    ),
    outputs=gr.Audio(label="الصوت الناتج (Generated Speech)", type="filepath"),
    title=title,
    description=description,
    examples=examples,
    cache_examples=False,
    theme=gr.themes.Soft()
)

# 3. Launch
if __name__ == "__main__":
    demo.launch()