import os
import torch

# --- CRITICAL FIX FOR PYTORCH 2.6+ ---
# Coqui TTS relies on pickling custom classes, which breaks with the new secure default.
# We must allow unsafe loading for this library to work.
torch.serialization.add_safe_globals = lambda *args, **kwargs: None 
_original_load = torch.load
def safe_load(*args, **kwargs):
    if 'weights_only' not in kwargs:
        kwargs['weights_only'] = False
    return _original_load(*args, **kwargs)
torch.load = safe_load
# -------------------------------------

from TTS.api import TTS
import gradio as gr

# --- Configuration ---
REFERENCE_AUDIO = "assets/voice_samples/reference.wav"
OUTPUT_FOLDER = "output"
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# --- Check Hardware ---
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Hardware: {device}")

# --- Load Model (Once) ---
print("⏳ Loading XTTS-v2 model... (This may take a minute)")
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
print("✅ Model Loaded!")


def generate_voice(text, audio_upload):
    if not text.strip():
        return None, "❌ Error: Please enter text."
    
    # Use uploaded audio if provided, otherwise fallback to default
    ref_audio_path = audio_upload if audio_upload else REFERENCE_AUDIO
    
    if not ref_audio_path or not os.path.exists(ref_audio_path):
        return None, f"❌ Error: No reference audio provided or file not found: {ref_audio_path}"

    print(f"🎙️ Generating audio for: {text[:50]}...")
    print(f"🎧 Using reference audio: {ref_audio_path}")
    
    output_path = os.path.join(OUTPUT_FOLDER, "test_output.wav")
    
    try:
        tts.tts_to_file(
            text=text,
            file_path=output_path,
            speaker_wav=ref_audio_path,
            language="hi"
        )
        return output_path, "✅ Generation Complete!"
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# --- Gradio UI ---
with gr.Blocks(title="Adhyatmik Voice Pro (XTTS-v2)") as app:
    gr.Markdown("## 🕉️ Adhyatmik Voice Pro (XTTS-v2 Testing)")
    gr.Markdown(f"**Status:** Running on **{device.upper()}**")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Enter Hindi Text (Adhyatmik Script)", 
                lines=5, 
                value="नमस्ते, यह एक परीक्षण है। अध्यात्म का अर्थ है अपनी आत्मा को जानना।"
            )
            # Make reference audio interactive
            ref_audio_input = gr.Audio(
                label="Reference Voice (Upload .wav, .mp3, etc.)", 
                type="filepath", 
                interactive=True,
                value=REFERENCE_AUDIO if os.path.exists(REFERENCE_AUDIO) else None
            )
            
            generate_btn = gr.Button("🚀 Generate Voice", variant="primary")
            status_msg = gr.Textbox(label="Status", interactive=False)
            
        with gr.Column():
            audio_player = gr.Audio(label="Generated Audio", type="filepath")

    generate_btn.click(
        fn=generate_voice, 
        inputs=[input_text, ref_audio_input], 
        outputs=[audio_player, status_msg]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)