import os
import random
import numpy as np
import torch
from pathlib import Path

# Đảm bảo torch.load luôn map về CPU khi cần
_orig_torch_load = torch.load
def _torch_load_cpu(f, *args, **kwargs):
    if "map_location" not in kwargs:
        kwargs["map_location"] = torch.device("cpu")
    return _orig_torch_load(f, *args, **kwargs)
torch.load = _torch_load_cpu

from huggingface_hub import snapshot_download
from chatterbox.src.chatterbox.tts import ChatterboxTTS
import gradio as gr

# --- CẤU HÌNH MODEL TỪ HUGGINGFACE ---
MODEL_REPO = "LTTEAM/TTS_Pro"
LOCAL_MODEL_DIR = Path(os.getcwd()) / "models" / "tts_pro"
MODELS = {}

# Download model một lần (cache)
if not LOCAL_MODEL_DIR.exists():
    print(f"📥 Đang tải model từ HuggingFace repo {MODEL_REPO} …")
    snapshot_download(
        repo_id=MODEL_REPO,
        repo_type="model",
        local_dir=str(LOCAL_MODEL_DIR),
        local_dir_use_symlinks=False
    )
    print(f"✅ Đã tải xong vào {LOCAL_MODEL_DIR}")

def get_or_load_model(device_str: str):
    """
    Lấy hoặc load model ChatterboxTTS trên device 'cpu' hoặc 'cuda'.
    device_str = "cpu" hoặc "gpu".
    """
    device = "cuda" if (device_str == "gpu" and torch.cuda.is_available()) else "cpu"
    if device not in MODELS:
        print(f"📂 Loading model lên {device} …")
        model = ChatterboxTTS.from_local(str(LOCAL_MODEL_DIR), device)
        MODELS[device] = model
        print(f"✅ Model đã được load lên {device}")
    return MODELS[device]

def set_seed(seed: int):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)

def chunk_text(text: str, chunk_size: int = 300):
    """Chia text dài thành các đoạn tối đa chunk_size ký tự, giữ nguyên từ."""
    words = text.split()
    chunks, current = [], ""
    for w in words:
        if len(current) + len(w) + 1 > chunk_size:
            chunks.append(current.strip())
            current = w
        else:
            current = f"{current} {w}".strip()
    if current:
        chunks.append(current.strip())
    return chunks

def generate_tts_audio(
    device_choice: str,
    text_input: str,
    audio_prompt_path: str,
    exaggeration: float,
    temperature: float,
    seed_num: int,
    cfg_weight: float
):
    """
    Sinh audio từ văn bản không giới hạn: chia thành chunk, generate từng chunk, ghép nối.
    Trả về (sample_rate, numpy.ndarray).
    """
    model = get_or_load_model(device_choice)
    if seed_num != 0:
        set_seed(int(seed_num))

    chunks = chunk_text(text_input, chunk_size=300)
    waves, sr = [], model.sr

    for idx, chunk in enumerate(chunks, start=1):
        print(f"🔊 Sinh đoạn {idx}/{len(chunks)} trên {model.device}")
        wav = model.generate(
            chunk,
            audio_prompt_path=audio_prompt_path,
            exaggeration=exaggeration,
            temperature=temperature,
            cfg_weight=cfg_weight,
        )
        waves.append(wav.squeeze(0).cpu().numpy())

    full_wave = np.concatenate(waves, axis=0)
    print("✅ Hoàn thành sinh toàn bộ audio.")
    return sr, full_wave

# --- GIAO DIỆN GRADIO TIẾNG VIỆT ---
with gr.Blocks(title="LTTEAM TTS") as demo:
    gr.Markdown(
        """
        # LTTEAM TTS  
        **Phát triển bởi: Lý Trần**  
        Ứng dụng chuyển văn bản thành giọng nói chất lượng cao, hỗ trợ đầu vào không giới hạn.
        """
    )
    with gr.Row():
        with gr.Column():
            device_choice = gr.Radio(
                choices=["cpu", "gpu"],
                value="gpu" if torch.cuda.is_available() else "cpu",
                label="Chọn thiết bị"
            )
            text = gr.Textbox(
                label="Văn bản (không giới hạn độ dài)",
                lines=8,
                placeholder="Dán hoặc nhập văn bản vào đây..."
            )
            ref_wav = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="Âm thanh mẫu (tùy chọn)",
                value="brian.wav"
            )
            exaggeration = gr.Slider(
                minimum=0.25, maximum=2, step=0.05,
                value=0.5,
                label="Mức nhấn nhá (Exaggeration)"
            )
            cfg_weight = gr.Slider(
                minimum=0.2, maximum=1, step=0.05,
                value=0.5,
                label="Trọng số CFG / Tốc độ"
            )
            with gr.Accordion("Tùy chọn thêm", open=False):
                seed_num = gr.Number(0, label="Seed (0 = random)")
                temperature = gr.Slider(
                    minimum=0.05, maximum=5, step=0.05,
                    value=0.8,
                    label="Nhiệt độ (Temperature)"
                )
            run = gr.Button("Chuyển giọng", variant="primary")

        with gr.Column():
            out_audio = gr.Audio(label="Kết quả âm thanh")

    run.click(
        fn=generate_tts_audio,
        inputs=[device_choice, text, ref_wav, exaggeration, temperature, seed_num, cfg_weight],
        outputs=[out_audio],
    )

if __name__ == "__main__":
    # Phát hiện Colab qua biến môi trường
    is_colab = "COLAB_GPU" in os.environ
    if is_colab:
        demo.launch(share=True)
    else:
        # Dùng host/port để hỗ trợ HuggingFace Spaces
        demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))