import spaces  # PHẢI import TRƯỚC mọi thứ!
import os
os.environ['SPACES_ZERO_GPU'] = '1'  # Set environment variable

import gradio as gr
import numpy as np
import soundfile as sf
import tempfile
import torch

# Import vieneutts SAU khi đã setup spaces
from vieneutts import VieNeuTTS

# Khởi tạo model trên CPU trước
print("📦 Đang tải model...")
tts = VieNeuTTS(
    backbone_repo="pnnbao-ump/VieNeu-TTS",
    backbone_device="cpu",  # Load trên CPU trước
    codec_repo="neuphonic/neucodec",
    codec_device="cpu"
)
print("✅ Model đã tải xong!")

# Danh sách giọng mẫu
VOICE_SAMPLES = {
    "Nam miền Nam": {
        "audio": "./sample/id_0001.wav",
        "text": "./sample/id_0001.txt"
    },
    "Nữ miền Nam": {
        "audio": "./sample/id_0002.wav",
        "text": "./sample/id_0002.txt"
    }
}

@spaces.GPU(duration=120)  # Giữ GPU trong 120 giây
def synthesize_speech(text, voice_choice, custom_audio=None, custom_text=None):
    """
    Tổng hợp giọng nói từ văn bản - Chạy trên GPU
    """
    try:
        # Kiểm tra text input
        if not text or text.strip() == "":
            return None, "❌ Vui lòng nhập văn bản cần tổng hợp"
        
        # Giới hạn độ dài text
        if len(text) > 500:
            return None, "❌ Văn bản quá dài! Vui lòng nhập tối đa 500 ký tự"
        
        # Xác định reference audio và text
        if custom_audio is not None and custom_text:
            ref_audio_path = custom_audio
            ref_text = custom_text
        elif voice_choice in VOICE_SAMPLES:
            ref_audio_path = VOICE_SAMPLES[voice_choice]["audio"]
            ref_text_path = VOICE_SAMPLES[voice_choice]["text"]
            with open(ref_text_path, "r", encoding="utf-8") as f:
                ref_text = f.read()
        else:
            return None, "❌ Vui lòng chọn giọng hoặc tải lên audio tùy chỉnh"
        
        # Di chuyển model lên GPU
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if device == "cuda":
            print("🚀 Đang chuyển model lên GPU...")
            tts.backbone = tts.backbone.to("cuda")
            tts.codec = tts.codec.to("cuda")
        
        # Encode reference audio
        print(f"📝 Đang xử lý: {text[:50]}...")
        ref_codes = tts.encode_reference(ref_audio_path)
        
        # Tổng hợp giọng nói
        print(f"🎵 Đang tổng hợp giọng nói trên {device.upper()}...")
        wav = tts.infer(text, ref_codes, ref_text)
        
        # Di chuyển model về CPU
        if device == "cuda":
            print("💾 Đang giải phóng GPU...")
            tts.backbone = tts.backbone.to("cpu")
            tts.codec = tts.codec.to("cpu")
            torch.cuda.empty_cache()
        
        # Lưu file tạm
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
            sf.write(tmp_file.name, wav, 24000)
            output_path = tmp_file.name
        
        print("✅ Hoàn thành!")
        return output_path, f"✅ Tổng hợp thành công trên {device.upper()}!"
        
    except Exception as e:
        print(f"❌ Lỗi: {str(e)}")
        import traceback
        traceback.print_exc()
        
        # Giải phóng GPU khi có lỗi
        try:
            if torch.cuda.is_available():
                tts.backbone = tts.backbone.to("cpu")
                tts.codec = tts.codec.to("cpu")
                torch.cuda.empty_cache()
        except:
            pass
            
        return None, f"❌ Lỗi: {str(e)}"

# Các ví dụ mẫu
examples = [
    ["Legacy là một bộ phim đột phá về mặt âm nhạc, quay phim, hiệu ứng đặc biệt, và tôi rất mừng vì cuối cùng nó cũng được cả giới phê bình lẫn người hâm mộ đánh giá lại. Chúng ta đã quá bất công với bộ phim này vào năm 2010.", "Nam miền Nam"],
    ["Từ nhiều nguồn tài liệu lịch sử, có thể thấy nuôi con theo phong cách Do Thái không chỉ tốt cho đứa trẻ mà còn tốt cho cả các bậc cha mẹ.", "Nữ miền Nam"],
    ["Các bác sĩ đang nghiên cứu một loại vaccine mới chống lại virus cúm mùa. Thí nghiệm lâm sàng cho thấy phản ứng miễn dịch mạnh mẽ và ít tác dụng phụ, mở ra hy vọng phòng chống dịch bệnh hiệu quả hơn trong tương lai.", "Nam miền Nam"],
]

# Custom CSS
custom_css = """
.gradio-container {
    max-width: 900px !important;
}
#warning {
    background-color: #fff3cd;
    border: 1px solid #ffc107;
    border-radius: 5px;
    padding: 10px;
    margin: 10px 0;
}
#info {
    background-color: #d1ecf1;
    border: 1px solid #17a2b8;
    border-radius: 5px;
    padding: 10px;
    margin: 10px 0;
}
"""

# Tạo giao diện Gradio
with gr.Blocks(title="VieNeu-TTS", css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎙️ VieNeu-TTS: Vietnamese Text-to-Speech
    
    Hệ thống tổng hợp tiếng nói tiếng Việt được **finetune từ NeuTTS-Air** - một mô hình TTS tiên tiến sử dụng Large Language Model và Neural Codec.
    
    Tác giả: [Phạm Ngọc Bảo](https://github.com/pnnbao97)
    Model: [VieNeu-TTS](https://huggingface.co/pnnbao-ump/VieNeu-TTS)
    Code: [GitHub](https://github.com/pnnbao97/VieNeu-TTS)
    Demo: [Hugging Face](https://huggingface.co/spaces/pnnbao-ump/VieNeu-TTS)
    """)
    
    with gr.Row():
        with gr.Column():
            # Input text
            text_input = gr.Textbox(
                label="📝 Văn bản đầu vào (tối đa 500 ký tự)",
                placeholder="Nhập văn bản tiếng Việt...",
                lines=4,
                max_lines=6,
                value="Legacy là một bộ phim đột phá về mặt âm nhạc, quay phim, hiệu ứng đặc biệt, và tôi rất mừng vì cuối cùng nó cũng được cả giới phê bình lẫn người hâm mộ đánh giá lại. Chúng ta đã quá bất công với bộ phim này vào năm 2010."  # Example 1 làm mặc định
            )
            
            # Character counter
            char_count = gr.Markdown("209 / 500 ký tự")  # Update số ký tự của example 1
            
            # Voice selection
            voice_select = gr.Radio(
                choices=list(VOICE_SAMPLES.keys()),
                label="🎤 Chọn giọng",
                value="Nam miền Nam"
            )
            
            # Custom voice option
            with gr.Accordion("🎨 Hoặc sử dụng giọng tùy chỉnh", open=False):
                gr.Markdown("*Upload file audio (.wav) và nội dung text tương ứng*")
                custom_audio = gr.Audio(
                    label="File audio mẫu",
                    type="filepath"
                )
                custom_text = gr.Textbox(
                    label="Nội dung của audio mẫu",
                    placeholder="Nhập chính xác nội dung...",
                    lines=2
                )
            
            # Submit button
            submit_btn = gr.Button("🎵 Tổng hợp giọng nói", variant="primary", size="lg")
        
        with gr.Column():
            # Output
            audio_output = gr.Audio(label="🔊 Kết quả")
            status_output = gr.Textbox(label="📊 Trạng thái", interactive=False)
    
    # Examples
    gr.Markdown("### 💡 Ví dụ nhanh")
    gr.Examples(
        examples=examples,
        inputs=[text_input, voice_select],
        outputs=[audio_output, status_output],
        fn=synthesize_speech,
        cache_examples=False
    )
    
    # Update character count
    def update_char_count(text):
        count = len(text) if text else 0
        color = "red" if count > 500 else "green"
        return f"<span style='color: {color}'>{count} / 500 ký tự</span>"
    
    text_input.change(
        fn=update_char_count,
        inputs=[text_input],
        outputs=[char_count]
    )
    
    # Event handler
    submit_btn.click(
        fn=synthesize_speech,
        inputs=[text_input, voice_select, custom_audio, custom_text],
        outputs=[audio_output, status_output]
    )
    
    gr.Markdown("""
    ---
    ### 📌 Thông tin
    
    **Liên kết:**
    - [GitHub Repository](https://github.com/pnnbao97/VieNeu-TTS)
    - [Model Card](https://huggingface.co/pnnbao-ump/VieNeu-TTS)
    
    <sub>Powered by VieNeu-TTS | Built with ❤️ for Vietnamese TTS</sub>
    """)

# Launch
if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )