app.py

import gradio as gr
import soundfile as sf
import tempfile
import os
import time
import numpy as np
import librosa
import re
import json
import shutil
from pathlib import Path
from datetime import datetime
from llama_cpp import Llama
from neucodec import NeuCodecOnnxDecoder
import torch
from utils.phonemize_text import phonemize_with_dict
import threading
from queue import Queue

print("⏳ Đang khởi động VieNeu-TTS...")

# --- CONSTANTS ---
MAX_CHARS_PER_CHUNK = 256
SAMPLE_RATE = 24000
DEVICE_INFO = "Q4 GGUF (llama-cpp) + ONNX Codec"

# Thư mục lưu lịch sử và custom voices
try:
    BASE_DIR = "./vieneu_data"
    HISTORY_DIR = os.path.join(BASE_DIR, "tts_history")
    CUSTOM_VOICES_DIR = os.path.join(BASE_DIR, "custom_voices")
    
    os.makedirs(HISTORY_DIR, exist_ok=True)
    os.makedirs(CUSTOM_VOICES_DIR, exist_ok=True)
    
    # Test write permission
    test_file = os.path.join(BASE_DIR, ".test")
    with open(test_file, 'w') as f:
        f.write("test")
    os.remove(test_file)
except (PermissionError, OSError):
    # Fallback to temp directory
    BASE_DIR = os.path.join(tempfile.gettempdir(), "vieneu_tts_data")
    HISTORY_DIR = os.path.join(BASE_DIR, "tts_history")
    CUSTOM_VOICES_DIR = os.path.join(BASE_DIR, "custom_voices")
    
    os.makedirs(HISTORY_DIR, exist_ok=True)
    os.makedirs(CUSTOM_VOICES_DIR, exist_ok=True)
    print(f"⚠️ Không có quyền ghi, sử dụng thư mục tạm: {BASE_DIR}")

HISTORY_JSON = os.path.join(HISTORY_DIR, "history.json")
CUSTOM_VOICES_JSON = os.path.join(CUSTOM_VOICES_DIR, "voices.json")

# Đường dẫn model
BACKBONE_REPO = "pnnbao-ump/VieNeu-TTS-q8-gguf"
CODEC_REPO = "neuphonic/neucodec-onnx-decoder"

# Giọng mẫu mặc định
DEFAULT_VOICE_SAMPLES = {
    "Tuyên (nam miền Bắc)": {
        "audio": "./sample/Tuyên (nam miền Bắc).wav",
        "text": "./sample/Tuyên (nam miền Bắc).txt",
        "codes": "./sample/Tuyên (nam miền Bắc).pt",
        "is_custom": False
    },
    "Vĩnh (nam miền Nam)": {
        "audio": "./sample/Vĩnh (nam miền Nam).wav",
        "text": "./sample/Vĩnh (nam miền Nam).txt",
        "codes": "./sample/Vĩnh (nam miền Nam).pt",
        "is_custom": False
    },
    "Bình (nam miền Bắc)": {
        "audio": "./sample/Bình (nam miền Bắc).wav",
        "text": "./sample/Bình (nam miền Bắc).txt",
        "codes": "./sample/Bình (nam miền Bắc).pt",
        "is_custom": False
    },
    "Nguyên (nam miền Nam)": {
        "audio": "./sample/Nguyên (nam miền Nam).wav",
        "text": "./sample/Nguyên (nam miền Nam).txt",
        "codes": "./sample/Nguyên (nam miền Nam).pt",
        "is_custom": False
    },
    "Sơn (nam miền Nam)": {
        "audio": "./sample/Sơn (nam miền Nam).wav",
        "text": "./sample/Sơn (nam miền Nam).txt",
        "codes": "./sample/Sơn (nam miền Nam).pt",
        "is_custom": False
    },
    "Đoan (nữ miền Nam)": {
        "audio": "./sample/Đoan (nữ miền Nam).wav",
        "text": "./sample/Đoan (nữ miền Nam).txt",
        "codes": "./sample/Đoan (nữ miền Nam).pt",
        "is_custom": False
    },
    "Ngọc (nữ miền Bắc)": {
        "audio": "./sample/Ngọc (nữ miền Bắc).wav",
        "text": "./sample/Ngọc (nữ miền Bắc).txt",
        "codes": "./sample/Ngọc (nữ miền Bắc).pt",
        "is_custom": False
    },
    "Ly (nữ miền Bắc)": {
        "audio": "./sample/Ly (nữ miền Bắc).wav",
        "text": "./sample/Ly (nữ miền Bắc).txt",
        "codes": "./sample/Ly (nữ miền Bắc).pt",
        "is_custom": False
    },
    "Dung (nữ miền Nam)": {
        "audio": "./sample/Dung (nữ miền Nam).wav",
        "text": "./sample/Dung (nữ miền Nam).txt",
        "codes": "./sample/Dung (nữ miền Nam).pt",
        "is_custom": False
    }
}

VOICE_SAMPLES = DEFAULT_VOICE_SAMPLES.copy()

# --- CUSTOM VOICE MANAGEMENT ---
voices_lock = threading.Lock()

def load_custom_voices():
    """Tải danh sách giọng custom từ file JSON"""
    global VOICE_SAMPLES
    
    with voices_lock:
        if os.path.exists(CUSTOM_VOICES_JSON):
            try:
                with open(CUSTOM_VOICES_JSON, 'r', encoding='utf-8') as f:
                    custom_voices = json.load(f)
                
                # Merge với default voices
                VOICE_SAMPLES = DEFAULT_VOICE_SAMPLES.copy()
                VOICE_SAMPLES.update(custom_voices)
                
                print(f"✅ Đã tải {len(custom_voices)} giọng custom")
                return True
            except Exception as e:
                print(f"⚠️ Lỗi đọc custom voices: {e}")
                return False
        return True

def save_custom_voices():
    """Lưu danh sách giọng custom vào file JSON"""
    with voices_lock:
        try:
            # Chỉ lưu các giọng custom
            custom_only = {k: v for k, v in VOICE_SAMPLES.items() if v.get("is_custom", False)}
            
            with open(CUSTOM_VOICES_JSON, 'w', encoding='utf-8') as f:
                json.dump(custom_only, f, ensure_ascii=False, indent=2)
            
            return True
        except Exception as e:
            print(f"⚠️ Lỗi ghi custom voices: {e}")
            return False

def extract_audio_codes(audio_path, codec):
    """Trích xuất codes từ file audio (placeholder - cần implement thật)"""
    try:
        # Load audio
        audio, sr = librosa.load(audio_path, sr=SAMPLE_RATE)
        
        # Placeholder: Tạo codes giả lập
        # Trong thực tế, bạn cần encoder để tạo codes từ audio
        # Ví dụ: codes = codec.encode(audio)
        
        # Tạm thời sử dụng random codes để demo
        duration_sec = len(audio) / SAMPLE_RATE
        num_codes = int(duration_sec * 50)  # ~50 codes/giây
        codes = np.random.randint(0, 1024, size=num_codes)
        
        return codes
    except Exception as e:
        print(f"❌ Lỗi trích xuất codes: {e}")
        return None

def add_custom_voice(voice_name, audio_file, text_content):
    """Thêm giọng custom mới"""
    global VOICE_SAMPLES
    
    if not voice_name or voice_name.strip() == "":
        return False, "⚠️ Vui lòng nhập tên giọng"
    
    voice_name = voice_name.strip()
    
    # Kiểm tra trùng tên
    if voice_name in VOICE_SAMPLES:
        return False, f"⚠️ Giọng '{voice_name}' đã tồn tại"
    
    if not audio_file:
        return False, "⚠️ Vui lòng upload file audio"
    
    if not text_content or text_content.strip() == "":
        return False, "⚠️ Vui lòng nhập text tương ứng"
    
    try:
        # Tạo ID duy nhất cho giọng
        voice_id = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        voice_dir = os.path.join(CUSTOM_VOICES_DIR, voice_id)
        os.makedirs(voice_dir, exist_ok=True)
        
        # Lưu audio
        audio_dest = os.path.join(voice_dir, "audio.wav")
        shutil.copy2(audio_file, audio_dest)
        
        # Lưu text
        text_dest = os.path.join(voice_dir, "text.txt")
        with open(text_dest, 'w', encoding='utf-8') as f:
            f.write(text_content.strip())
        
        # Trích xuất codes (placeholder)
        print(f"🔄 Đang trích xuất codes cho giọng '{voice_name}'...")
        codes = extract_audio_codes(audio_file, codec)
        
        if codes is None or len(codes) == 0:
            shutil.rmtree(voice_dir)
            return False, "❌ Không thể trích xuất codes từ audio"
        
        # Lưu codes
        codes_dest = os.path.join(voice_dir, "codes.pt")
        torch.save(torch.from_numpy(codes), codes_dest)
        
        # Thêm vào VOICE_SAMPLES
        with voices_lock:
            VOICE_SAMPLES[voice_name] = {
                "audio": audio_dest,
                "text": text_dest,
                "codes": codes_dest,
                "is_custom": True,
                "voice_id": voice_id,
                "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
        
        # Lưu vào JSON
        if save_custom_voices():
            return True, f"✅ Đã thêm giọng '{voice_name}' thành công!"
        else:
            # Rollback nếu lưu thất bại
            del VOICE_SAMPLES[voice_name]
            shutil.rmtree(voice_dir)
            return False, "❌ Lỗi khi lưu thông tin giọng"
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return False, f"❌ Lỗi: {str(e)}"

def delete_custom_voice(voice_name):
    """Xóa giọng custom"""
    global VOICE_SAMPLES
    
    if voice_name not in VOICE_SAMPLES:
        return False, "⚠️ Không tìm thấy giọng"
    
    if not VOICE_SAMPLES[voice_name].get("is_custom", False):
        return False, "⚠️ Không thể xóa giọng mặc định"
    
    try:
        voice_info = VOICE_SAMPLES[voice_name]
        voice_id = voice_info.get("voice_id")
        
        # Xóa thư mục chứa files
        if voice_id:
            voice_dir = os.path.join(CUSTOM_VOICES_DIR, voice_id)
            if os.path.exists(voice_dir):
                shutil.rmtree(voice_dir)
        
        # Xóa khỏi VOICE_SAMPLES
        with voices_lock:
            del VOICE_SAMPLES[voice_name]
        
        # Lưu lại JSON
        save_custom_voices()
        
        return True, f"✅ Đã xóa giọng '{voice_name}'"
        
    except Exception as e:
        return False, f"❌ Lỗi khi xóa: {str(e)}"

def get_voice_list():
    """Lấy danh sách tất cả giọng"""
    return list(VOICE_SAMPLES.keys())

def get_custom_voices_display():
    """Hiển thị danh sách giọng custom dạng HTML"""
    custom_voices = {k: v for k, v in VOICE_SAMPLES.items() if v.get("is_custom", False)}
    
    if not custom_voices:
        return """
        <div style='padding: 20px; text-align: center; color: #64748b;'>
            <p style='font-size: 1.1em;'>🎤 Chưa có giọng tùy chỉnh nào</p>
            <p style='font-size: 0.9em;'>Thêm giọng mới bằng cách upload audio và text tương ứng</p>
        </div>
        """
    
    html_parts = ["<div style='font-family: system-ui; line-height: 1.6;'>"]
    
    for i, (name, info) in enumerate(custom_voices.items(), 1):
        created_at = info.get("created_at", "N/A")
        voice_id = info.get("voice_id", "N/A")
        
        html_parts.append(f"""
        <div style='
            background: white;
            border: 1px solid #e2e8f0;
            border-radius: 8px;
            padding: 15px;
            margin-bottom: 12px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
        '>
            <div style='display: flex; justify-content: space-between; align-items: start; margin-bottom: 8px;'>
                <div style='font-weight: 600; color: #1e293b; font-size: 0.95em;'>
                    <span style='color: #64748b;'>#{i}</span> {name}
                </div>
                <div style='font-size: 0.85em; color: #64748b;'>
                    {created_at}
                </div>
            </div>
            
            <div style='display: flex; gap: 20px; font-size: 0.85em; color: #64748b;'>
                <div>🎤 Giọng tùy chỉnh</div>
                <div style='margin-left: auto; color: #3b82f6;'>
                    ID: {voice_id[:13]}...
                </div>
            </div>
        </div>
        """)
    
    html_parts.append("</div>")
    return "".join(html_parts)

# Load custom voices khi khởi động
load_custom_voices()

# --- HISTORY MANAGEMENT ---
history_lock = threading.Lock()

def load_history():
    """Tải lịch sử từ file JSON - Thread-safe"""
    with history_lock:
        if os.path.exists(HISTORY_JSON):
            try:
                with open(HISTORY_JSON, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                print(f"⚠️ Lỗi đọc history.json: {e}")
                return []
        return []

def save_history(history):
    """Lưu lịch sử vào file JSON - Thread-safe"""
    with history_lock:
        try:
            with open(HISTORY_JSON, 'w', encoding='utf-8') as f:
                json.dump(history, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"⚠️ Lỗi ghi history.json: {e}")

def add_to_history(text, voice, audio_path, duration, status):
    """Thêm một bản ghi vào lịch sử - Thread-safe"""
    try:
        history = load_history()
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = f"tts_{timestamp}.wav"
        permanent_path = os.path.join(HISTORY_DIR, filename)
        
        if audio_path and os.path.exists(audio_path):
            try:
                shutil.copy2(audio_path, permanent_path)
            except Exception as e:
                print(f"⚠️ Không thể copy file audio: {e}")
                permanent_path = audio_path
        else:
            permanent_path = None
        
        record = {
            "id": timestamp,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "text": text[:100] + "..." if len(text) > 100 else text,
            "full_text": text,
            "voice": voice,
            "audio_path": permanent_path,
            "duration": duration,
            "status": status
        }
        
        history.insert(0, record)
        
        if len(history) > 100:
            old_record = history.pop()
            try:
                if old_record.get('audio_path') and os.path.exists(old_record['audio_path']):
                    os.remove(old_record['audio_path'])
            except Exception as e:
                print(f"⚠️ Không thể xóa file cũ: {e}")
        
        save_history(history)
        return permanent_path
    except Exception as e:
        print(f"⚠️ Lỗi khi lưu lịch sử: {e}")
        import traceback
        traceback.print_exc()
        return audio_path if audio_path else None

def get_history_list():
    """Lấy danh sách lịch sử để hiển thị dạng HTML"""
    history = load_history()
    if not history:
        return """
        <div style='padding: 20px; text-align: center; color: #64748b;'>
            <p style='font-size: 1.1em;'>📭 Chưa có lịch sử tổng hợp nào</p>
            <p style='font-size: 0.9em;'>Các bản ghi sẽ xuất hiện ở đây sau khi bạn tổng hợp giọng nói</p>
        </div>
        """
    
    html_parts = ["<div style='font-family: system-ui; line-height: 1.6;'>"]
    
    for i, record in enumerate(history[:50], 1):
        status_color = "#10b981" if record['status'] == "Thành công" else "#ef4444"
        status_icon = "✅" if record['status'] == "Thành công" else "❌"
        
        html_parts.append(f"""
        <div style='
            background: white;
            border: 1px solid #e2e8f0;
            border-radius: 8px;
            padding: 15px;
            margin-bottom: 12px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
        '>
            <div style='display: flex; justify-content: space-between; align-items: start; margin-bottom: 8px;'>
                <div style='font-weight: 600; color: #1e293b; font-size: 0.95em;'>
                    <span style='color: #64748b;'>#{i}</span> {record['voice']}
                </div>
                <div style='font-size: 0.85em; color: #64748b;'>
                    {record['timestamp']}
                </div>
            </div>
            
            <div style='
                background: #f8fafc;
                padding: 10px;
                border-radius: 6px;
                margin-bottom: 8px;
                color: #334155;
                font-size: 0.9em;
                border-left: 3px solid #3b82f6;
            '>
                {record['text']}
            </div>
            
            <div style='display: flex; gap: 20px; font-size: 0.85em; color: #64748b;'>
                <div>⏱️ {record['duration']:.2f}s</div>
                <div style='color: {status_color}; font-weight: 500;'>
                    {status_icon} {record['status']}
                </div>
                <div style='margin-left: auto; color: #3b82f6; cursor: pointer;'>
                    ID: {record['id'][:13]}...
                </div>
            </div>
        </div>
        """)
    
    html_parts.append("</div>")
    return "".join(html_parts)

def clear_all_history():
    """Xóa toàn bộ lịch sử"""
    history = load_history()
    
    for record in history:
        try:
            if record.get('audio_path') and os.path.exists(record['audio_path']):
                os.remove(record['audio_path'])
        except Exception as e:
            print(f"⚠️ Không thể xóa file: {e}")
    
    save_history([])
    return "✅ Đã xóa toàn bộ lịch sử"

# --- BACKGROUND PROCESSING QUEUE ---
processing_queue = Queue()
is_processing = False
processing_stats = {"total": 0, "success": 0, "failed": 0}

def background_processor():
    """Xử lý queue tổng hợp trong background"""
    global is_processing, processing_stats
    
    while True:
        task = processing_queue.get()
        if task is None:
            break
        
        is_processing = True
        text, voice = task
        
        try:
            print(f"[Background] Bắt đầu tổng hợp: {text[:50]}...")
            result = synthesize_speech_internal(text, voice)
            
            if result:
                processing_stats["success"] += 1
                print(f"[Background] ✅ Hoàn thành: {result}")
            else:
                processing_stats["failed"] += 1
                print(f"[Background] ❌ Thất bại")
            
            processing_stats["total"] += 1
            
        except Exception as e:
            processing_stats["failed"] += 1
            processing_stats["total"] += 1
            print(f"[Background] ❌ Lỗi: {e}")
            import traceback
            traceback.print_exc()
        
        is_processing = False
        processing_queue.task_done()

bg_thread = threading.Thread(target=background_processor, daemon=True)
bg_thread.start()

def get_processing_stats():
    """Lấy thống kê xử lý"""
    return f"📊 Tổng: {processing_stats['total']} | ✅ Thành công: {processing_stats['success']} | ❌ Thất bại: {processing_stats['failed']}"

def split_text_into_chunks(text, max_chars=256):
    """Chia text thành chunks"""
    sentences = re.split(r'([.!?,;])', text)
    chunks = []
    current = ""
    
    for i in range(0, len(sentences), 2):
        sentence = sentences[i]
        punct = sentences[i+1] if i+1 < len(sentences) else ""
        segment = sentence + punct
        
        if len(current) + len(segment) <= max_chars:
            current += segment
        else:
            if current:
                chunks.append(current.strip())
            current = segment
    
    if current:
        chunks.append(current.strip())
    
    return chunks if chunks else [text]

def decode_audio(codes_str, codec):
    """Decode speech tokens to audio"""
    speech_ids = [int(num) for num in re.findall(r"<\|speech_(\d+)\|>", codes_str)]
    
    if len(speech_ids) == 0:
        print("Không tìm thấy mã giọng nói hợp lệ.")
        return np.array([], dtype=np.float32)
    
    codes = np.array(speech_ids, dtype=np.int32)[np.newaxis, np.newaxis, :]
    recon = codec.decode_code(codes)
    return recon[0, 0, :]

# --- MODEL LOADING ---
print("📦 Đang tải model Q4 GGUF và Codec ONNX...")
model_loaded = False
backbone = None
codec = None

try:
    backbone = Llama.from_pretrained(
        repo_id=BACKBONE_REPO,
        filename="*.gguf",
        verbose=False,
        n_gpu_layers=-1,
        n_ctx=2048,
        mlock=True,
        flash_attn=True,
    )
    
    codec = NeuCodecOnnxDecoder.from_pretrained(CODEC_REPO)
    
    print("✅ Model đã tải thành công!")
    model_loaded = True
except Exception as e:
    import traceback
    traceback.print_exc()
    print(f"❌ Lỗi khi tải model: {e}")
    model_loaded = False

# --- SYNTHESIS FUNCTION (Internal) ---
def synthesize_speech_internal(text, voice_choice):
    """Internal synthesis function không phụ thuộc UI"""
    global backbone, codec, model_loaded
    
    if not model_loaded:
        print("❌ Model chưa được tải")
        return None
    
    if not text or text.strip() == "":
        print("❌ Text rỗng")
        return None
    
    if voice_choice not in VOICE_SAMPLES:
        print(f"❌ Giọng không hợp lệ: {voice_choice}")
        return None
    
    raw_text = text.strip()
    
    ref_text_path = VOICE_SAMPLES[voice_choice]["text"]
    try:
        with open(ref_text_path, "r", encoding="utf-8") as f:
            ref_text_raw = f.read()
    except Exception as e:
        print(f"❌ Lỗi đọc file text mẫu: {e}")
        return None

    ref_codes = None
    codes_path = VOICE_SAMPLES[voice_choice]["codes"]
    
    try:
        ref_codes_tensor = torch.load(codes_path, map_location="cpu")
        if isinstance(ref_codes_tensor, torch.Tensor):
            ref_codes = ref_codes_tensor.cpu().numpy()
        else:
            ref_codes = np.array(ref_codes_tensor)
    except Exception as e:
        print(f"❌ Lỗi khi tải codes: {e}")
        return None
    
    if ref_codes is None or len(ref_codes) == 0:
        print("❌ Codes rỗng")
        return None

    text_chunks = split_text_into_chunks(raw_text, max_chars=MAX_CHARS_PER_CHUNK)
    
    all_audio_segments = []
    silence_pad = np.zeros(int(SAMPLE_RATE * 0.15), dtype=np.float32)
    
    start_time = time.time()
    
    try:
        for i, chunk in enumerate(text_chunks):
            print(f"[Internal] Xử lý chunk {i+1}/{len(text_chunks)}")
            
            ref_text_phoneme = phonemize_with_dict(ref_text_raw)
            input_text_phoneme = phonemize_with_dict(chunk)
            
            codes_str = "".join([f"<|speech_{idx}|>" for idx in ref_codes])
            
            prompt = (
                f"user: Convert the text to speech:<|TEXT_PROMPT_START|>{ref_text_phoneme} {input_text_phoneme}"
                f"<|TEXT_PROMPT_END|>\nassistant:<|SPEECH_GENERATION_START|>{codes_str}"
            )
            
            output = backbone(
                prompt,
                max_tokens=2048,
                temperature=1.0,
                top_k=50,
                stop=["<|SPEECH_GENERATION_END|>"],
            )
            output_str = output["choices"][0]["text"]
            
            chunk_wav = decode_audio(output_str, codec)
            
            if chunk_wav is not None and len(chunk_wav) > 0:
                all_audio_segments.append(chunk_wav)
                if i < len(text_chunks) - 1:
                    all_audio_segments.append(silence_pad)
        
        if not all_audio_segments:
            print("❌ Không có audio segment nào")
            add_to_history(raw_text, voice_choice, None, 0, "Lỗi: Không sinh được audio")
            return None
        
        final_wav = np.concatenate(all_audio_segments)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            sf.write(tmp.name, final_wav, SAMPLE_RATE)
            output_path = tmp.name
        
        process_time = time.time() - start_time
        
        permanent_path = add_to_history(raw_text, voice_choice, output_path, process_time, "Thành công")
        
        print(f"✅ Hoàn thành: {permanent_path}")
        return permanent_path
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        add_to_history(raw_text, voice_choice, None, 0, f"Lỗi: {str(e)}")
        return None

# --- SYNTHESIS FUNCTION (UI) ---
def synthesize_speech(text, voice_choice):
    """Main synthesis function với UI feedback"""
    global backbone, codec, model_loaded
    
    if not model_loaded:
        yield None, "⚠️ Model chưa tải. Vui lòng kiểm tra lỗi console!"
        return
    
    if not text or text.strip() == "":
        yield None, "⚠️ Vui lòng nhập văn bản!"
        return
    
    if voice_choice not in VOICE_SAMPLES:
        yield None, "⚠️ Vui lòng chọn giọng mẫu."
        return
    
    raw_text = text.strip()
    
    ref_text_path = VOICE_SAMPLES[voice_choice]["text"]
    try:
        with open(ref_text_path, "r", encoding="utf-8") as f:
            ref_text_raw = f.read()
    except Exception as e:
        yield None, f"❌ Lỗi đọc file text mẫu: {e}"
        return

    yield None, "📄 Đang xử lý Reference..."
    
    ref_codes = None
    codes_path = VOICE_SAMPLES[voice_choice]["codes"]
    
    try:
        ref_codes_tensor = torch.load(codes_path, map_location="cpu")
        if isinstance(ref_codes_tensor, torch.Tensor):
            ref_codes = ref_codes_tensor.cpu().numpy()
        else:
            ref_codes = np.array(ref_codes_tensor)
    except Exception as e:
        yield None, f"❌ Lỗi khi tải codes: {e}"
        return
    
    if ref_codes is None or len(ref_codes) == 0:
        yield None, "❌ Codes tham chiếu không hợp lệ."
        return

    text_chunks = split_text_into_chunks(raw_text, max_chars=MAX_CHARS_PER_CHUNK)
    total_chunks = len(text_chunks)
    
    yield None, f"🚀 Bắt đầu tổng hợp ({total_chunks} đoạn)..."
    
    all_audio_segments = []
    silence_pad = np.zeros(int(SAMPLE_RATE * 0.15), dtype=np.float32)
    
    start_time = time.time()
    
    try:
        for i, chunk in enumerate(text_chunks):
            yield None, f"⏳ Đang xử lý đoạn {i+1}/{total_chunks}... ({int((i/total_chunks)*100)}%)"
            
            ref_text_phoneme = phonemize_with_dict(ref_text_raw)
            input_text_phoneme = phonemize_with_dict(chunk)
            
            codes_str = "".join([f"<|speech_{idx}|>" for idx in ref_codes])
            
            prompt = (
                f"user: Convert the text to speech:<|TEXT_PROMPT_START|>{ref_text_phoneme} {input_text_phoneme}"
                f"<|TEXT_PROMPT_END|>\nassistant:<|SPEECH_GENERATION_START|>{codes_str}"
            )
            
            output = backbone(
                prompt,
                max_tokens=2048,
                temperature=1.0,
                top_k=50,
                stop=["<|SPEECH_GENERATION_END|>"],
            )
            output_str = output["choices"][0]["text"]
            
            chunk_wav = decode_audio(output_str, codec)
            
            if chunk_wav is not None and len(chunk_wav) > 0:
                all_audio_segments.append(chunk_wav)
                if i < total_chunks - 1:
                    all_audio_segments.append(silence_pad)
        
        if not all_audio_segments:
            yield None, "❌ Không sinh được audio nào."
            add_to_history(raw_text, voice_choice, None, 0, "Lỗi: Không sinh được audio")
            return
        
        yield None, "💾 Đang ghép file và lưu..."
        
        final_wav = np.concatenate(all_audio_segments)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            sf.write(tmp.name, final_wav, SAMPLE_RATE)
            output_path = tmp.name
        
        process_time = time.time() - start_time
        
        permanent_path = add_to_history(raw_text, voice_choice, output_path, process_time, "Thành công")
        
        yield permanent_path, f"✅ Hoàn tất! (Tổng thời gian: {process_time:.2f}s | RTF: {process_time/(len(final_wav)/SAMPLE_RATE):.3f})"
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        add_to_history(raw_text, voice_choice, None, 0, f"Lỗi: {str(e)}")
        yield None, f"❌ Lỗi tổng hợp: {str(e)}"

def refresh_history():
    """Làm mới danh sách lịch sử"""
    return get_history_list()

def load_history_item(item_index):
    """Tải một item từ lịch sử theo số thứ tự"""
    if not item_index or item_index.strip() == "":
        return None, "", "", "⚠️ Vui lòng nhập số thứ tự"
    
    try:
        index = int(item_index.strip()) - 1
        history = load_history()
        
        if index < 0 or index >= len(history):
            return None, "", "", f"❌ Số thứ tự không hợp lệ (1-{len(history)})"
        
        record = history[index]
        
        audio_path = None
        if record.get('audio_path') and os.path.exists(record['audio_path']):
            audio_path = record['audio_path']
        
        info = f"""
📅 Thời gian: {record['timestamp']}
⏱️ Thời lượng: {record['duration']:.2f}s
🎭 Giọng: {record['voice']}
📊 Trạng thái: {record['status']}
🆔 ID: {record['id']}
        """.strip()
        
        return audio_path, record['full_text'], record['voice'], info
        
    except ValueError:
        return None, "", "", "❌ Vui lòng nhập số hợp lệ"
    except Exception as e:
        return None, "", "", f"❌ Lỗi: {str(e)}"

# --- UI SETUP ---
theme = gr.themes.Ocean(
    primary_hue="indigo",
    secondary_hue="cyan",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont('Inter'), 'ui-sans-serif', 'system-ui'],
).set(
    button_primary_background_fill="linear-gradient(90deg, #6366f1 0%, #0ea5e9 100%)",
    button_primary_background_fill_hover="linear-gradient(90deg, #4f46e5 0%, #0284c7 100%)",
)

css = """
.container { max-width: 1400px; margin: auto; }
.header-box {
    text-align: center;
    margin-bottom: 25px;
    padding: 25px;
    background: linear-gradient(135deg, #0f172a 0%, #1e293b 100%);
    border-radius: 12px;
    color: white;
}
.header-title {
    font-size: 2.5rem;
    font-weight: 800;
}
.gradient-text {
    background: -webkit-linear-gradient(45deg, #60A5FA, #22D3EE);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
.status-box {
    font-weight: bold;
    text-align: center;
    border: none;
    background: transparent;
}
.model-card-content {
    display: flex;
    flex-wrap: wrap;
    justify-content: center;
    align-items: center;
    gap: 15px;
    font-size: 0.9rem;
    color: #cbd5e1;
}
.model-card-item {
    display: flex;
    align-items: center;
    justify-content: center;
    gap: 6px;
    color: #94a3b8;
}
.model-card-link {
    color: #3b82f6;
    text-decoration: none;
    font-weight: 500;
    transition: color 0.2s;
}
.model-card-link:hover {
    color: #2563eb;
    text-decoration: underline;
}
.history-container {
    max-height: 650px;
    overflow-y: auto;
    padding: 10px;
    background: #f1f5f9;
    border-radius: 8px;
}
.info-box {
    background: #f8fafc;
    padding: 12px;
    border-radius: 6px;
    border-left: 4px solid #3b82f6;
    margin: 10px 0;
}
"""

EXAMPLES_LIST = [
    ["Về miền Tây không chỉ để ngắm nhìn sông nước hữu tình, mà còn để cảm nhận tấm chân tình của người dân nơi đây.", "Vĩnh (nam miền Nam)"],
    ["Hà Nội những ngày vào thu mang một vẻ đẹp trầm mặc và cổ kính đến lạ thường.", "Bình (nam miền Bắc)"],
    ["Thành phố Hồ Chí Minh luôn chuyển mình không ngừng với nhịp sống hối hả, năng động.", "Dung (nữ miền Nam)"],
]

initial_status = f"✅ Model đã tải thành công! (Chạy trên **{DEVICE_INFO}**). Hỗ trợ xử lý background và lưu lịch sử." if model_loaded else "❌ Lỗi khi tải model."

with gr.Blocks(title="VieNeu-TTS", theme=theme, css=css) as demo:
    with gr.Column(elem_classes="container"):
        gr.HTML(f"""
<div class="header-box">
    <h1 class="header-title">
        <span class="header-icon">🦜</span>
        <span class="gradient-text">VieNeu-TTS Studio</span>
    </h1>
    <p style="margin: 10px 0;">Chế độ: {DEVICE_INFO} | Background Processing ✅ | History ✅ | Custom Voices ✅</p>
    <p style="font-size: 0.85rem; color: #94a3b8;">📁 Dữ liệu lưu tại: {BASE_DIR}</p>
    <div class="model-card-content">
        <div class="model-card-item">
            <strong>Repository:</strong>
            <a href="https://github.com/pnnbao97/VieNeu-TTS" target="_blank" class="model-card-link">GitHub</a>
        </div>
        <div class="model-card-item">
            <strong>Tác giả:</strong>
            <span>Phạm Nguyễn Ngọc Bảo</span>
        </div>
    </div>
</div>
        """)
        
        status_banner = gr.Markdown(initial_status)

        # --- TABS ---
        with gr.Tabs():
            # TAB 1: Tổng hợp
            with gr.Tab("🎙️ Tổng hợp"):
                with gr.Row(elem_classes="container"):
                    with gr.Column(scale=3):
                        text_input = gr.Textbox(
                            label=f"📝 Văn bản (Chia chunk: {MAX_CHARS_PER_CHUNK} ký tự)",
                            lines=7,
                            placeholder="Nhập văn bản tiếng Việt cần tổng hợp...",
                            value="Hà Nội, trái tim của Việt Nam, là một thành phố ngàn năm văn hiến với bề dày lịch sử và văn hóa độc đáo. Bước chân trên những con phố cổ kính quanh Hồ Hoàn Kiếm, du khách như được du hành ngược thời gian.",
                        )
                        
                        voice_select = gr.Dropdown(
                            choices=get_voice_list(),
                            value=get_voice_list()[0] if get_voice_list() else None,
                            label="👤 Chọn giọng mẫu",
                        )
                        
                        with gr.Row():
                            btn_generate = gr.Button("🎵 Bắt đầu tổng hợp", variant="primary", size="lg", interactive=model_loaded)
                            btn_clear = gr.Button("🗑️ Xóa", variant="secondary", size="lg")
                    
                    with gr.Column(scale=2):
                        audio_output = gr.Audio(
                            label="🔊 Kết quả",
                            type="filepath",
                            autoplay=True
                        )
                        status_output = gr.Textbox(
                            label="📊 Trạng thái", 
                            elem_classes="status-box", 
                            value="Chờ nhập văn bản...",
                            lines=2
                        )

                gr.Examples(
                    examples=EXAMPLES_LIST,
                    inputs=[text_input, voice_select],
                    outputs=[audio_output, status_output],
                    fn=synthesize_speech,
                    cache_examples=False,
                    label="💡 Các ví dụ nhanh"
                )
            
            # TAB 2: Quản lý giọng
            with gr.Tab("🎤 Quản lý giọng"):
                gr.Markdown("### ➕ Thêm giọng tùy chỉnh")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        new_voice_name = gr.Textbox(
                            label="📝 Tên giọng",
                            placeholder="VD: Mai (nữ miền Trung)",
                            info="Đặt tên mô tả cho giọng mới"
                        )
                        
                        new_voice_audio = gr.Audio(
                            label="🎵 File audio mẫu (WAV)",
                            type="filepath",
                            sources=["upload"]
                        )
                        
                        new_voice_text = gr.Textbox(
                            label="📄 Text tương ứng với audio",
                            lines=5,
                            placeholder="Nhập chính xác nội dung trong file audio...",
                            info="Text này phải khớp chính xác với audio mẫu"
                        )
                        
                        with gr.Row():
                            btn_add_voice = gr.Button("➕ Thêm giọng", variant="primary", size="lg")
                            btn_refresh_voices = gr.Button("🔄 Làm mới", variant="secondary", size="sm")
                        
                        add_voice_status = gr.Textbox(
                            label="📊 Trạng thái",
                            value="Chờ thêm giọng...",
                            interactive=False
                        )
                    
                    with gr.Column(scale=1):
                        gr.Markdown("### 💡 Hướng dẫn")
                        gr.Markdown("""
                        **Yêu cầu file audio:**
                        - Format: WAV (khuyến nghị)
                        - Sample rate: 24kHz (tối ưu)
                        - Độ dài: 5-30 giây
                        - Chất lượng: Rõ ràng, ít nhiễu
                        
                        **Lưu ý:**
                        - Text phải khớp chính xác với audio
                        - Giọng mẫu càng rõ, kết quả càng tốt
                        - Tránh nhiễu và echo
                        - Nói tự nhiên, không quá nhanh/chậm
                        
                        **⚠️ Chú ý:**
                        Quá trình trích xuất codes có thể mất vài giây
                        """)
                
                gr.Markdown("---")
                gr.Markdown("### 📋 Danh sách giọng tùy chỉnh")
                
                with gr.Row():
                    with gr.Column(scale=3):
                        custom_voices_display = gr.HTML(
                            value=get_custom_voices_display(),
                            elem_classes="history-container"
                        )
                    
                    with gr.Column(scale=2):
                        gr.Markdown("### 🗑️ Xóa giọng")
                        
                        delete_voice_select = gr.Dropdown(
                            choices=[v for v in get_voice_list() if VOICE_SAMPLES[v].get("is_custom", False)],
                            label="Chọn giọng cần xóa",
                            info="Chỉ hiển thị giọng tùy chỉnh"
                        )
                        
                        btn_delete_voice = gr.Button("🗑️ Xóa giọng", variant="stop")
                        
                        delete_voice_status = gr.Textbox(
                            label="📊 Trạng thái",
                            interactive=False
                        )
                        
                        gr.Markdown(f"""
                        ### 📊 Thống kê
                        - **Giọng mặc định**: {len([v for v in VOICE_SAMPLES.values() if not v.get('is_custom', False)])}
                        - **Giọng tùy chỉnh**: {len([v for v in VOICE_SAMPLES.values() if v.get('is_custom', False)])}
                        - **Tổng**: {len(VOICE_SAMPLES)}
                        """)
            
            # TAB 3: Lịch sử
            with gr.Tab("📜 Lịch sử"):
                with gr.Row():
                    with gr.Column(scale=3):
                        gr.Markdown("### 📋 Danh sách lịch sử")
                        
                        with gr.Row():
                            btn_refresh = gr.Button("🔄 Làm mới", size="sm", variant="secondary")
                            btn_clear_all = gr.Button("🗑️ Xóa toàn bộ", size="sm", variant="stop")
                            stats_display = gr.Textbox(
                                value=get_processing_stats(),
                                label="",
                                show_label=False,
                                interactive=False,
                                container=False
                            )
                        
                        history_display = gr.HTML(
                            value=get_history_list(),
                            elem_classes="history-container"
                        )
                    
                    with gr.Column(scale=2):
                        gr.Markdown("### 🔍 Chi tiết bản ghi")
                        
                        with gr.Row():
                            history_select = gr.Textbox(
                                label="Nhập số thứ tự (vd: 1, 2, 3...)",
                                placeholder="Nhập số...",
                                scale=3
                            )
                            btn_load_history = gr.Button("📂 Tải", variant="primary", scale=1)
                        
                        history_info = gr.Textbox(
                            label="ℹ️ Thông tin",
                            lines=6,
                            elem_classes="info-box"
                        )
                        
                        history_audio = gr.Audio(
                            label="🔊 Audio", 
                            type="filepath"
                        )
                        
                        history_voice = gr.Textbox(
                            label="🎭 Giọng đã dùng",
                            interactive=False
                        )
                        
                        history_text = gr.Textbox(
                            label="📄 Văn bản đầy đủ",
                            lines=5,
                            interactive=False
                        )
                        
                        with gr.Row():
                            btn_reuse = gr.Button("♻️ Tái sử dụng văn bản", variant="secondary")
            
            # TAB 4: Thông tin
            with gr.Tab("ℹ️ Thông tin"):
                gr.Markdown(f"""
                ## 🎯 Về VieNeu-TTS
                
                **VieNeu-TTS** là hệ thống tổng hợp giọng nói tiếng Việt sử dụng công nghệ AI tiên tiến.
                
                ### ⚙️ Cấu hình hiện tại
                - **Model Backbone**: Q4 GGUF (llama-cpp)
                - **Codec**: ONNX Decoder
                - **Sample Rate**: {SAMPLE_RATE} Hz
                - **Max Chunk Size**: {MAX_CHARS_PER_CHUNK} ký tự
                - **Thư mục dữ liệu**: `{BASE_DIR}`
                
                ### 🎭 Giọng mẫu
                - **Giọng mặc định**: {len([v for v in VOICE_SAMPLES.values() if not v.get('is_custom', False)])}
                - **Giọng tùy chỉnh**: {len([v for v in VOICE_SAMPLES.values() if v.get('is_custom', False)])}
                - **Tổng cộng**: {len(VOICE_SAMPLES)} giọng
                
                ### 📌 Hướng dẫn sử dụng
                
                **1. Tổng hợp giọng nói**:
                   - Nhập văn bản vào ô "Văn bản"
                   - Chọn giọng mẫu phù hợp
                   - Nhấn "Bắt đầu tổng hợp"
                   - Đợi hệ thống xử lý và nghe kết quả
                
                **2. Thêm giọng tùy chỉnh**:
                   - Vào tab "Quản lý giọng"
                   - Nhập tên giọng (VD: "Hùng (nam miền Trung)")
                   - Upload file audio mẫu (WAV, 5-30s, rõ ràng)
                   - Nhập text tương ứng chính xác
                   - Nhấn "Thêm giọng" và đợi xử lý
                   - Giọng mới sẽ xuất hiện trong dropdown
                
                **3. Xem lịch sử**:
                   - Vào tab "Lịch sử"
                   - Nhấn "Làm mới" để cập nhật danh sách
                   - Nhập số thứ tự và nhấn "Tải" để xem chi tiết
                
                **4. Xóa giọng tùy chỉnh**:
                   - Vào tab "Quản lý giọng"
                   - Chọn giọng cần xóa từ dropdown
                   - Nhấn "Xóa giọng" (chỉ xóa được giọng tùy chỉnh)
                
                ### 🔧 Tính năng nâng cao
                - ✅ Xử lý background (không phụ thuộc UI)
                - ✅ Lưu lịch sử tự động
                - ✅ Quản lý giọng tùy chỉnh
                - ✅ Persistent storage (giữ dữ liệu sau khi tắt)
                - ✅ Chia chunk thông minh
                - ✅ Thread-safe operations
                - ✅ Tự động xóa file cũ khi vượt quá 100 bản ghi
                
                ### 📊 Thống kê
                {get_processing_stats()}
                
                ### 🔗 Liên kết
                - [GitHub Repository](https://github.com/pnnbao97/VieNeu-TTS)
                - [Model Hub](https://huggingface.co/pnnbao-ump)
                
                ---
                **Phiên bản**: 3.0 Custom Voice | **Tác giả**: Phạm Nguyễn Ngọc Bảo
                """)
        
        # Event handlers - Tab Tổng hợp
        btn_generate.click(
            fn=synthesize_speech,
            inputs=[text_input, voice_select],
            outputs=[audio_output, status_output]
        )
        
        btn_clear.click(
            fn=lambda: ("", None, "Đã xóa"),
            outputs=[text_input, audio_output, status_output]
        )
        
        # Event handlers - Tab Quản lý giọng
        def handle_add_voice(name, audio, text):
            success, message = add_custom_voice(name, audio, text)
            
            if success:
                # Cập nhật dropdown và display
                new_choices = get_voice_list()
                custom_display = get_custom_voices_display()
                delete_choices = [v for v in new_choices if VOICE_SAMPLES[v].get("is_custom", False)]
                
                return (
                    message,  # add_voice_status
                    gr.update(choices=new_choices),  # voice_select (tab 1)
                    custom_display,  # custom_voices_display
                    gr.update(choices=delete_choices),  # delete_voice_select
                    "",  # new_voice_name
                    None,  # new_voice_audio
                    ""  # new_voice_text
                )
            else:
                return (
                    message,
                    gr.update(),
                    gr.update(),
                    gr.update(),
                    gr.update(),
                    gr.update(),
                    gr.update()
                )
        
        btn_add_voice.click(
            fn=handle_add_voice,
            inputs=[new_voice_name, new_voice_audio, new_voice_text],
            outputs=[
                add_voice_status,
                voice_select,
                custom_voices_display,
                delete_voice_select,
                new_voice_name,
                new_voice_audio,
                new_voice_text
            ]
        )
        
        def handle_delete_voice(voice_name):
            if not voice_name:
                return "⚠️ Vui lòng chọn giọng", gr.update(), gr.update(), gr.update()
            
            success, message = delete_custom_voice(voice_name)
            
            if success:
                new_choices = get_voice_list()
                custom_display = get_custom_voices_display()
                delete_choices = [v for v in new_choices if VOICE_SAMPLES[v].get("is_custom", False)]
                
                return (
                    message,
                    gr.update(choices=new_choices),
                    custom_display,
                    gr.update(choices=delete_choices, value=None)
                )
            else:
                return message, gr.update(), gr.update(), gr.update()
        
        btn_delete_voice.click(
            fn=handle_delete_voice,
            inputs=delete_voice_select,
            outputs=[
                delete_voice_status,
                voice_select,
                custom_voices_display,
                delete_voice_select
            ]
        )
        
        def refresh_all_voices():
            custom_display = get_custom_voices_display()
            delete_choices = [v for v in get_voice_list() if VOICE_SAMPLES[v].get("is_custom", False)]
            
            return (
                custom_display,
                gr.update(choices=delete_choices),
                "✅ Đã làm mới"
            )
        
        btn_refresh_voices.click(
            fn=refresh_all_voices,
            outputs=[custom_voices_display, delete_voice_select, add_voice_status]
        )
        
        # Event handlers - Tab Lịch sử
        btn_refresh.click(
            fn=refresh_history,
            outputs=history_display
        ).then(
            fn=get_processing_stats,
            outputs=stats_display
        )
        
        btn_load_history.click(
            fn=load_history_item,
            inputs=history_select,
            outputs=[history_audio, history_text, history_voice, history_info]
        )
        
        def clear_all():
            msg = clear_all_history()
            return get_history_list(), msg
        
        btn_clear_all.click(
            fn=clear_all,
            outputs=[history_display, stats_display]
        )
        
        def reuse_text(text, voice):
            return text, voice, "✅ Đã tải văn bản vào tab Tổng hợp"
        
        btn_reuse.click(
            fn=reuse_text,
            inputs=[history_text, history_voice],
            outputs=[text_input, voice_select, status_output]
        )

if __name__ == "__main__":
    print(f"\n{'='*60}")
    print(f"🚀 VieNeu-TTS Studio đã sẵn sàng!")
    print(f"📂 Dữ liệu lưu tại: {BASE_DIR}")
    print(f"📜 Lịch sử: {HISTORY_DIR}")
    print(f"🎤 Giọng tùy chỉnh: {CUSTOM_VOICES_DIR}")
    print(f"🎭 Tổng số giọng: {len(VOICE_SAMPLES)} ({len([v for v in VOICE_SAMPLES.values() if not v.get('is_custom', False)])} mặc định + {len([v for v in VOICE_SAMPLES.values() if v.get('is_custom', False)])} tùy chỉnh)")
    print(f"⚙️  Chế độ: {DEVICE_INFO}")
    print(f"{'='*60}\n")
    
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )