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 time
import pandas as pd
# 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"
    }
}
# Lưu lịch sử (global list)
history = []

@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", None
        
        # 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ự", None
        
        # 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", None
        
        # 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
        
        # Lưu vào lịch sử
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        history.append({
            "Thời gian": timestamp,
            "Văn bản": text,
            "Giọng": voice_choice if voice_choice else "Tùy chỉnh",
            "Audio": output_path
        })
        
        print("✅ Hoàn thành!")
        return output_path, f"✅ Tổng hợp thành công trên {device.upper()}!", update_history()
        
    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)}", None

# Hàm cập nhật lịch sử cho Dataframe
def update_history():
    if not history:
        return pd.DataFrame(columns=["Thời gian", "Văn bản", "Giọng", "Audio"])
    df = pd.DataFrame(history)
    df = df[["Thời gian", "Văn bản", "Giọng"]]  # Không hiển thị path audio đầy đủ trong df
    return df

# Hàm tải về lịch sử dưới dạng CSV
def download_history():
    if not history:
        return None
    df = pd.DataFrame(history)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp_file:
        df.to_csv(tmp_file.name, index=False, encoding="utf-8-sig")
        return tmp_file.name

# 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 Nguyễn 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)
    
    # Lịch sử
    gr.Markdown("### 📜 Lịch sử tổng hợp")
    history_df = gr.Dataframe(
        value=update_history(),
        headers=["Thời gian", "Văn bản", "Giọng"],
        interactive=False
    )
    download_btn = gr.Button("📥 Tải về lịch sử (CSV)")
    download_file = gr.File(label="Tệp tải về", interactive=False)
    
    # Examples
    gr.Markdown("### 💡 Ví dụ nhanh")
    gr.Examples(
        examples=examples,
        inputs=[text_input, voice_select],
        outputs=[audio_output, status_output, history_df],
        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 for synthesize
    submit_btn.click(
        fn=synthesize_speech,
        inputs=[text_input, voice_select, custom_audio, custom_text],
        outputs=[audio_output, status_output, history_df]
    )
    
    # Event handler for download
    download_btn.click(
        fn=download_history,
        inputs=[],
        outputs=[download_file]
    )
    
    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
    )