import os
import re
import sys
import subprocess
import tempfile
from typing import Optional, List, Tuple

# ==========================================
# AUTO-INSTALL DEPENDENCIES (SELF-HEALING)
# ==========================================
# Chuyên gia Fix: Tự động cài đặt các thư viện thiếu nếu môi trường chưa có
def install_dependencies():
    required_packages = {
        "neucodec": "neucodec",
        "vinorm": "vinorm",
        "phonemizer": "phonemizer",
        "scipy": "scipy",
        "soundfile": "soundfile"
    }
    for import_name, package_name in required_packages.items():
        try:
            __import__(import_name)
        except ImportError:
            print(f"⏳ Đang cài đặt thiếu sót: {package_name}...")
            try:
                subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
                print(f"✅ Đã cài xong {package_name}")
            except Exception as e:
                print(f"⚠️ Không thể cài {package_name}: {e}")

install_dependencies()

# ==========================================
# IMPORTS SAU KHI ĐÃ CÀI ĐẶT
# ==========================================
import numpy as np
import gradio as gr
import torch
import librosa
import soundfile as sf

from transformers import AutoTokenizer, AutoModelForCausalLM

# Import an toàn cho các thư viện đặc thù
try:
    from neucodec import NeuCodec
    from vinorm import TTSnorm
    from phonemizer.backend import EspeakBackend
except ImportError as e:
    print(f"❌ Lỗi Import nghiêm trọng: {e}. Vui lòng kiểm tra lại môi trường.")
    # Sẽ không crash ngay ở đây để Gradio vẫn có thể hiện UI báo lỗi nếu cần
    NeuCodec = None
    TTSnorm = None
    EspeakBackend = None

# ==========================================
# CẤU HÌNH HỆ THỐNG (SYSTEM CONFIG)
# ==========================================
MODEL_ID = "dinhthuan/neutts-air-vi"
CODEC_ID = "neuphonic/neucodec"

# Cấu hình Audio
REF_SR = 16000     # Sample rate yêu cầu cho Reference
OUT_SR = 24000     # Sample rate đầu ra
MIN_REF_DUR = 3.0  # Giây tối thiểu
MAX_REF_DUR = 30.0 # Giây tối đa (để tránh lỗi bộ nhớ encoding)

# Các Token đặc biệt
SPEECH_START = "<|SPEECH_GENERATION_START|>"
SPEECH_END   = "<|SPEECH_GENERATION_END|>"

# Thiết bị xử lý
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Gradio 6/Latest Transformers warning fix: use dtype instead of torch_dtype if warned
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32

# ==========================================
# KHỞI TẠO MODEL (RUNTIME SETUP)
# ==========================================
print(f"🚀 Đang khởi tạo trên thiết bị: {DEVICE}...")

model = None
codec = None
tokenizer = None
phonemizer = None

def load_models():
    global model, codec, tokenizer, phonemizer
    try:
        if tokenizer is None:
            tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        
        if model is None:
            # FIX: Thay torch_dtype bằng dtype theo warning log
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_ID,
                dtype=DTYPE,
                trust_remote_code=True,
            ).to(DEVICE)
            model.eval()

        if codec is None and NeuCodec:
            codec = NeuCodec.from_pretrained(CODEC_ID).to(DEVICE)
            codec.eval()

        # Khởi tạo Phonemizer
        if phonemizer is None and EspeakBackend:
            try:
                phonemizer = EspeakBackend(language="vi", preserve_punctuation=True, with_stress=True)
                print("✅ Khởi tạo Phonemizer thành công!")
            except RuntimeError as e:
                print("❌ Lỗi Espeak: Không tìm thấy thư viện espeak-ng.")
                print("👉 HÃY THÊM 'espeak-ng' VÀO FILE packages.txt TRÊN HUGGINGFACE SPACES.")
                phonemizer = None
                
        print("✅ Khởi tạo Model hoàn tất!")
    except Exception as e:
        print(f"❌ Lỗi khởi tạo model: {str(e)}")

# Gọi load models ngay khi chạy
load_models()

# ==========================================
# HÀM HỖ TRỢ (HELPERS)
# ==========================================

def _normalize_text(text: str) -> str:
    """Chuẩn hóa và chuyển đổi văn bản sang phoneme."""
    if phonemizer is None:
        raise gr.Error("Lỗi hệ thống: Chưa cài đặt 'espeak-ng'. Vui lòng kiểm tra System Logs.")
        
    text = (text or "").strip()
    if not text:
        return ""
    
    if TTSnorm:
        # Chuẩn hóa tiếng Việt
        t = TTSnorm(text, punc=False, unknown=True, lower=False, rule=False)
    else:
        t = text

    # Chuyển sang phoneme
    return phonemizer.phonemize([t])[0]

def _check_audio_duration(filepath: str) -> bool:
    """Kiểm tra độ dài file audio có phù hợp không."""
    try:
        dur = librosa.get_duration(path=filepath)
        if dur < MIN_REF_DUR:
            raise gr.Error(f"Audio quá ngắn ({dur:.1f}s). Vui lòng dùng audio > {MIN_REF_DUR}s.")
        if dur > MAX_REF_DUR:
            gr.Warning(f"Audio hơi dài ({dur:.1f}s). Hệ thống sẽ tự động cắt hoặc chỉ lấy phần đầu.")
        return True
    except Exception as e:
        print(f"Lỗi kiểm tra audio: {e}")
        return False

def _encode_ref_audio(ref_wav_path: str) -> torch.Tensor:
    """Đọc file audio, resample về 16k và encode thành speech codes."""
    if codec is None:
        raise gr.Error("Lỗi hệ thống: NeuCodec chưa được tải.")

    # Kiểm tra file
    if not os.path.exists(ref_wav_path):
        raise gr.Error("Không tìm thấy file audio reference.")

    _check_audio_duration(ref_wav_path)

    # Load và convert sang mono, 16kHz
    wav, _ = librosa.load(ref_wav_path, sr=REF_SR, mono=True)
    
    # Cắt ngắn nếu quá dài để tránh tràn bộ nhớ khi encode
    max_len = int(MAX_REF_DUR * REF_SR)
    if len(wav) > max_len:
        wav = wav[:max_len]

    wav_t = torch.from_numpy(wav).float().unsqueeze(0).unsqueeze(0).to(DEVICE)
    
    with torch.inference_mode():
        codes = codec.encode_code(audio_or_path=wav_t).squeeze(0).squeeze(0).detach().cpu()
    
    return codes

def _extract_generated_codes(text: str) -> List[int]:
    """Trích xuất mã âm thanh từ chuỗi token trả về."""
    if SPEECH_START in text and SPEECH_END in text:
        # Lấy phần nằm giữa thẻ start và end
        text = text.split(SPEECH_START, 1)[1].split(SPEECH_END, 1)[0]
    # Tìm tất cả các pattern speech_123
    return [int(x) for x in re.findall(r"<\|speech_(\d+)\|>", text)]

# ==========================================
# HÀM XỬ LÝ CHÍNH (CORE LOGIC)
# ==========================================
@torch.inference_mode()
def run_tts(
    target_text: str, 
    ref_audio_path: Optional[str], 
    ref_text: str, 
    max_new_tokens: int
) -> str:
    """
    Hàm xử lý chính cho Gradio.
    """
    # Đảm bảo model đã load
    if model is None or tokenizer is None:
        load_models()
        if model is None:
            raise gr.Error("Model chưa sẵn sàng. Vui lòng thử lại sau giây lát hoặc kiểm tra Logs.")

    # 1. Validation đầu vào
    target_text = (target_text or "").strip()
    ref_text = (ref_text or "").strip()

    if not target_text:
        raise gr.Error("Vui lòng nhập 'Văn bản cần đọc'.")
    
    if not ref_audio_path:
        raise gr.Error("Vui lòng Upload hoặc Ghi âm giọng mẫu (Reference Audio).")
        
    if not ref_text:
        raise gr.Error("Vui lòng nhập nội dung văn bản của file ghi âm (Reference Text).")

    # Giới hạn token để tránh tạo audio quá dài (> 5 phút)
    if max_new_tokens > 4000:
        gr.Warning("Số lượng tokens quá lớn có thể khiến việc tạo giọng mất nhiều thời gian hoặc lỗi.")

    # 2. Xử lý Phoneme
    print(f"--- Processing: {target_text[:20]}... ---")
    
    target_phones = _normalize_text(target_text)
    ref_phones = _normalize_text(ref_text)
    
    if not target_phones:
        raise gr.Error("Lỗi: Không thể chuyển văn bản đích sang phoneme (Kiểm tra espeak).")
    if not ref_phones:
        raise gr.Error("Lỗi: Không thể chuyển văn bản mẫu sang phoneme (Kiểm tra espeak).")

    # 3. Encode Reference Audio
    ref_codes = _encode_ref_audio(ref_audio_path)
    
    # Chuyển codes thành chuỗi token
    ref_codes_str = "".join([f"<|speech_{i}|>" for i in ref_codes.tolist()])

    # 4. Tạo Prompt cho mô hình
    # Cấu trúc: Reference Phones + Target Phones
    combined_phones = ref_phones + " " + target_phones
    
    chat_prompt = (
        "user: Convert the text to speech:"
        f"<|TEXT_PROMPT_START|>{combined_phones}<|TEXT_PROMPT_END|>\n"
        f"assistant:{SPEECH_START}{ref_codes_str}"
    )

    # 5. Generate (Inference)
    input_ids = tokenizer.encode(chat_prompt, return_tensors="pt").to(DEVICE)
    speech_end_id = tokenizer.convert_tokens_to_ids(SPEECH_END)

    print("-> Dang generate...")
    out = model.generate(
        input_ids,
        max_new_tokens=int(max_new_tokens),
        temperature=1.0,
        top_k=50,
        eos_token_id=speech_end_id,
        pad_token_id=tokenizer.eos_token_id,
    )

    # 6. Decode output tokens
    out_text = tokenizer.decode(out[0], skip_special_tokens=False)
    all_codes = _extract_generated_codes(out_text)

    # Loại bỏ phần ref codes ban đầu, chỉ lấy phần mới sinh ra
    # Logic: Nếu model lặp lại ref code, ta cắt bỏ nó đi
    if len(all_codes) > len(ref_codes):
        gen_codes = all_codes[len(ref_codes):]
    else:
        # Fallback nếu model sinh ngắn quá hoặc lỗi
        gen_codes = all_codes

    if len(gen_codes) < 10:
        raise gr.Error("Mô hình không sinh được âm thanh. Hãy thử Reference Audio khác rõ ràng hơn.")

    # 7. Decode về Audio Waveform
    codes_tensor = torch.tensor(gen_codes, dtype=torch.long).view(1, 1, -1).to(DEVICE)
    audio_numpy = codec.decode_code(codes_tensor).detach().cpu().numpy()[0, 0, :]
    
    # Clip để tránh rè
    audio_numpy = np.clip(audio_numpy, -1.0, 1.0)

    # 8. Lưu file tạm thời (Tempfile) để Gradio hiển thị nút Download
    # Sử dụng tempfile an toàn, file sẽ tồn tại đến khi app đóng hoặc hệ thống dọn dẹp
    temp_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(temp_file.name, audio_numpy, OUT_SR)
    
    print(f"✅ Đã tạo xong: {temp_file.name}")
    return temp_file.name

# ==========================================
# GIAO DIỆN NGƯỜI DÙNG (GRADIO UI)
# ==========================================
css = """
.container { max-width: 900px; margin: auto; }
.note-box { background-color: #f0f9ff; border-left: 5px solid #00b4d8; padding: 15px; border-radius: 5px; margin-bottom: 20px; }
"""

# FIX: Xóa css và theme khỏi constructor Blocks (Gradio 6.0 compatibility)
with gr.Blocks(title="Expert Voice Cloning - NeuTTS (Auto-Fix Build)") as demo:
    
    gr.Markdown("# 🎙️ Clone Giọng Nói Tiếng Việt (NeuTTS-Air Refactored)")
    
    # Hộp ghi chú chuyên gia
    gr.HTML("""
    <div class="note-box">
        <h3>💡 Ghi chú từ chuyên gia:</h3>
        <ul>
            <li><b>Lưu ý Build:</b> Hệ thống đã tự động vá lỗi thư viện. Nếu vẫn lỗi "phonemize", hãy đảm bảo bạn đã cài <code>espeak-ng</code>.</li>
            <li><b>Môi trường:</b> Phòng yên tĩnh, mic để gần miệng, hạn chế tiếng vang (echo).</li>
            <li><b>Reference Audio (Mẫu):</b> Tốt nhất khoảng <b>15 giây</b>. Đọc tốc độ vừa phải, tự nhiên.</li>
            <li><b>Reference Text:</b> Phải khớp tuyệt đối với nội dung bạn đọc trong file mẫu (khoảng 100-120 từ).</li>
        </ul>
    </div>
    """)

    with gr.Row():
        # CỘT TRÁI: CẤU HÌNH & INPUT
        with gr.Column(scale=1):
            gr.Markdown("### 1. Nội dung cần đọc")
            text_input = gr.Textbox(
                label="Nhập văn bản tiếng Việt",
                lines=5,
                placeholder="Nhập nội dung bạn muốn AI đọc bằng giọng của bạn...",
                value="Xin chào, đây là thử nghiệm tái tạo giọng nói chính chủ."
            )
            
            with gr.Accordion("Cài đặt nâng cao", open=False):
                max_tokens = gr.Slider(
                    minimum=256, 
                    maximum=3072, 
                    value=1536, 
                    step=128, 
                    label="Độ dài tối đa (Max Tokens)",
                    info="Tăng lên nếu văn bản dài. Đừng quá 5 phút audio."
                )

        # CỘT PHẢI: REFERENCE & OUTPUT
        with gr.Column(scale=1):
            gr.Markdown("### 2. Dữ liệu giọng mẫu (Reference)")
            
            # Tích hợp cả Micro và Upload vào 1 component
            ref_audio_input = gr.Audio(
                label="Chọn nguồn Audio mẫu (15s)",
                sources=["microphone", "upload"], # Cho phép cả 2
                type="filepath", # Luôn trả về đường dẫn file
                editable=True
            )
            
            ref_text_input = gr.Textbox(
                label="Nội dung Reference Text (Bắt buộc)",
                lines=3,
                placeholder="Gõ lại chính xác những gì bạn vừa nói ở trên...",
                info="Khớp tuyệt đối với audio mẫu. Khoảng 100-120 từ."
            )

    # NÚT BẤM VÀ KẾT QUẢ
    btn_generate = gr.Button("🚀 Tạo giọng nói (Start Cloning)", variant="primary", size="lg")
    
    gr.Markdown("### 3. Kết quả")
    
    # FIX: Xóa tham số show_download_button=True (Gradio 6.0 error fix)
    audio_output = gr.Audio(
        label="Audio đã tạo", 
        type="filepath",
        interactive=False # Output mode
    )

    # SỰ KIỆN
    btn_generate.click(
        fn=run_tts,
        inputs=[text_input, ref_audio_input, ref_text_input, max_tokens],
        outputs=audio_output
    )

if __name__ == "__main__":
    # FIX: Chuyển theme và css vào launch()
    demo.queue().launch(
        share=False, 
        theme=gr.themes.Soft(), 
        css=css
    )