import spaces
import os
from download_models import download_all_models
from huggingface_hub import login

import gradio as gr

# ======================= HF LOGIN & DOWNLOAD MODEL =======================
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)

# Tải model khi Space ACTIVE
download_all_models()

from infer import run_zipvoice

# NEW: ASR + DENOISE
from chunkformer import ChunkFormerModel
from clearvoice import ClearVoice
from proccess_wav import enhance_ref_audio, transcribe_ref_audio

# (Nếu 2 dòng test này không cần thì bạn có thể xoá bớt cho nhẹ)
enhanced = enhance_ref_audio("Toại.wav")
text = transcribe_ref_audio(enhanced)


def infer_ref_text_ui(ref_audio_path: str) -> str:
    """
    Dùng cho nút 'Infer Text':
    - Enhance WAV (ClearVoice + xử lý khoảng lặng + cắt 5–10s)
    - ASR theo khoảng lặng
    - Đổ kết quả vào ô Reference Text
    """
    if not ref_audio_path:
        raise gr.Error("Vui lòng upload file giọng mẫu trước khi infer text.")

    try:
        enhanced = enhance_ref_audio(ref_audio_path)
        text = transcribe_ref_audio(enhanced)
    except Exception as e:
        raise gr.Error(f"Lỗi khi nhận dạng từ audio tham chiếu: {e}")

    if not text:
        raise gr.Error("Không nhận dạng được nội dung từ audio tham chiếu.")
    return text


# ======================= CẤU HÌNH DEMO SẴN =======================
SAMPLE_CONFIGS = [
    {
        "name": "Sample 1 – Kể chuyện",
        "ref_audio": "Toại.wav",
        "ref_text": "Trong bóng tối, Toại nói cái gì đó mà Thoan không nghe thấy.",
        "gen_text": "Đêm nay trời nhiều mây, ánh trăng bị che khuất, chỉ còn lại một dải sáng yếu ớt rơi xuống con đường đất trải dài giữa cánh đồng. Cậu bé tên Tín đang dắt chiếc xe đạp cũ đi về nhà, bánh xe bị cán đinh nên lăn nặng và chậm như con trâu mệt nhọc sau vụ mùa. Gió thổi lạnh buốt, mùi bùn đất ngai ngái quấn lấy chân cậu. Tới đoạn rẽ dẫn vào xóm, Tín nghe tiếng nước chảy khe khẽ từ con mương bên đường. Tiếng ấy vẫn quen thuộc, nhưng tối nay lại vang khác lạ, như có giọng người đang hòa vào nhịp nước, lúc trầm lúc cao, nghe mơ hồ mà lạnh sống lưng. Cậu dừng lại, nghiêng tai lắng nghe, tim đập nhanh như muốn vượt khỏi lồng ngực.",
        "out_audio": "Toại_output.wav",
    },
    {
        "name": "Sample 2 – Nữ",
        "ref_audio": "Trung.wav",
        "ref_text": "Mùa hè không chỉ là khoảng thời gian nghỉ ngơi, mà còn là khoảng thời gian tuyệt vời.",
        "gen_text": "Từ các kết quả này, chúng tôi đề xuất rằng sự kết hợp nhuần nhuyễn giữa adaptive optimization, robust training pipelines và interpretable model design sẽ là chìa khóa để phát triển các hệ thống ây ai vừa mạnh mẽ vừa đáng tin cậy trong môi trường thực tế.",
        "out_audio": "Trung_output.wav",
    },
    {
        "name": "Sample 3 – English",
        "ref_audio": "T_English.wav",
        "ref_text": "And turning to the pole which he had dragged, He drew it close beneath the widowed bough, And what was of it unto it left bound.",
        "gen_text": "Recent experiments indicate that the current model architecture still exhibits significant overfitting, especially when evaluated on out of distribution samples. Although the training accuracy remains consistently high, the performance drops sharply when the model is exposed to noise perturbed inputs, suggesting limited robustness.",
        "out_audio": "T_English_output.wav",
    },
]

# Hàm dùng khi bấm "Dùng sample này"
def make_sample_loader(sample):
    def _load_sample():
        return (
            sample["ref_audio"],  # ref_audio -> input Audio
            sample["ref_text"],   # ref_text -> Textbox
            sample["gen_text"],   # gen_text -> Textbox
            sample["out_audio"],  # output_audio -> Audio
        )
    return _load_sample


# ======================= STYLE TÙY CHỈNH (LÀM SÁNG HƠN) =======================
custom_css = """
#app-container {
    max-width: 1000px;
    margin: 0 auto;
}
.gradio-container {
    background: radial-gradient(circle at top, #ffffff 0, #f9fafb 55%);
    color: #111827;
}

/* Tiêu đề lớn */
#title-block h1 {
    font-size: 2.4rem !important;
    font-weight: 800 !important;
    background: linear-gradient(120deg, #f97316, #eab308, #22c55e);
    -webkit-background-clip: text;
    color: transparent;
    text-align: center;
}
#title-block p {
    text-align:center;
    font-size: 0.95rem;
    color: #6b7280;
}

/* Card sáng hơn */
.sample-card {
    border-radius: 16px;
    padding: 16px;
    background: rgba(255, 255, 255, 0.96);
    border: 1px solid rgba(148, 163, 184, 0.6);
    box-shadow: 0 18px 28px rgba(148, 163, 184, 0.35);
}

/* Nút bấm */
button.primary {
    border-radius: 999px !important;
    font-weight: 600 !important;
}

/* Tabs */
.svelte-1ipelgc, .tabitem {
    font-weight: 600;
}
"""

# ======================= XỬ LÝ TEXT (NẾU CẦN) =======================
def post_process(text: str) -> str:
    text = " " + text + " "
    text = text.replace(" . . ", " . ")
    text = " " + text + " "
    text = text.replace(" .. ", " . ")
    text = " " + text + " "
    text = text.replace(" , , ", " , ")
    text = " " + text + " "
    text = text.replace(" ,, ", " , ")
    text = " " + text + " "
    text = text.replace('"', "")
    return " ".join(text.split())


@spaces.GPU
def infer_tts(ref_audio_path, ref_text, gen_text, steps, request: gr.Request = None):
    if not ref_audio_path:
        raise gr.Error("Please upload a sample audio file.")

    if not gen_text.strip():
        raise gr.Error("Please enter the text content to generate voice.")

    # Giới hạn độ dài nội dung (4000 từ)
    if len(gen_text.split()) > 4000:
        raise gr.Error("Please enter text content with less than 4000 words.")

    # 1) Enhance ref audio: clearvoice + xử lý khoảng lặng + cắt 5–10s
    try:
        enhanced_ref_audio = enhance_ref_audio(ref_audio_path)
    except Exception as e:
        raise gr.Error(f"Lỗi khi xử lý audio tham chiếu: {e}")

    # 2) Nếu không có ref_text thì chạy ASR theo khoảng lặng
    if not ref_text or not ref_text.strip():
        try:
            inferred = transcribe_ref_audio(enhanced_ref_audio)
            if not inferred:
                raise gr.Error(
                    "Không nhận dạng được nội dung từ audio tham chiếu. "
                    "Vui lòng nhập Reference Text thủ công."
                )
            ref_text = inferred
            print(f"[ASR] Inferred ref_text: {ref_text}")
        except gr.Error:
            raise
        except Exception as e:
            raise gr.Error(f"Lỗi khi tự động nhận dạng Reference Text: {e}")

    try:
        out_path = "result.wav"

        run_zipvoice(
            model_name="zipvoice",
            prompt_wav=enhanced_ref_audio,  # dùng file đã xử lý
            prompt_text=ref_text.strip() if ref_text else "xin chào các bạn",
            text=gen_text,
            res_wav_path=out_path,
            lang="vi",
            tokenizer_name="espeak",
            num_step=steps,
            seed=123456,
            speed=1.0,
        )

        return out_path

    except Exception as e:
        raise gr.Error(f"Error generating voice: {e}")


# ======================= UI =======================
with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    with gr.Column(elem_id="app-container"):
        # --------- TIÊU ĐỀ ----------
        gr.Markdown(
            """
<div id="title-block">
  <h1>🎤 ZipVoice – Zero-shot Vietnamese TTS</h1>
  <p>Upload một mẫu giọng + nhập nội dung &rarr; hệ thống sẽ bắt chước giọng nói và đọc đoạn text của bạn.</p>
</div>
            """,
            elem_id="title-block",
        )

        with gr.Tabs():
            # Chỉ còn 1 tab chính, demo cũng nằm trong tab này
            with gr.TabItem("🎯 Tự tạo giọng nói"):
                # --------- KHỐI INPUT / OUTPUT CHÍNH ----------
                with gr.Row():
                    with gr.Column(elem_classes=["sample-card"]):
                        gr.Markdown("#### 1️⃣ Tải giọng mẫu & nhập text")

                        ref_audio = gr.Audio(
                            label="🔊 Sample Voice (upload hoặc kéo thả)",
                            type="filepath",
                        )

                        ref_text = gr.Textbox(
                            label="📝 Reference Text (optional)",
                            placeholder="Nội dung đang được nói trong file giọng mẫu (nên tự viết cho chính xác)",
                            lines=3,
                        )

                        # Nút infer text từ audio tham chiếu (ASR + khử nhiễu)
                        btn_infer_text = gr.Button(
                            "✨ Infer Text từ audio tham chiếu"
                        )

                        gen_text = gr.Textbox(
                            label="📝 Text to Generate",
                            placeholder="Nhập nội dung tiếng Việt bạn muốn tổng hợp...",
                            lines=6,
                        )

                        steps = gr.Slider(
                            8,
                            64,
                            value=25,
                            step=1,
                            label="⚡ Step (càng lớn, càng tốt, càng lâu)",
                        )

                        btn_synthesize = gr.Button(
                            "🔥 Generate Voice",
                            variant="primary",
                        )

                    with gr.Column(elem_classes=["sample-card"]):
                        gr.Markdown("#### 2️⃣ Kết quả tổng hợp")
                        output_audio = gr.Audio(
                            label="🎧 Generated Audio",
                            type="filepath",
                        )
                        gr.Markdown(
                            """
- Bạn có thể tải file `.wav` về sau khi tạo.
- Nếu nghe chưa ổn, hãy thử:
  - Dùng **ref audio ngắn 3-8s, phát âm chuẩn hơn.
                            """
                        )

                # mapping nút Generate -> infer_tts
                btn_synthesize.click(
                    infer_tts,
                    inputs=[ref_audio, ref_text, gen_text, steps],
                    outputs=[output_audio],
                )

                # mapping nút Infer Text -> điền ref_text (có khử nhiễu trước)
                btn_infer_text.click(
                    infer_ref_text_ui,
                    inputs=[ref_audio],
                    outputs=[ref_text],
                )

                # --------- KHỐI DEMO NẰM NGAY TRONG TAB CHÍNH ----------
                gr.Markdown(
                    """
### 🎧 Demo có sẵn
Click vào một sample bên dưới để tự động nạp:
- 🔊 Giọng mẫu (ref voice)  
- 📝 Reference text  
- 📝 Text to generate  
- 🎧 Output audio mẫu  
                    """
                )

                for sample in SAMPLE_CONFIGS:
                    with gr.Column(elem_classes=["sample-card"]):
                        gr.Markdown(f"### {sample['name']}")
                        with gr.Row():
                            gr.Audio(
                                value=sample["ref_audio"],
                                label="🔊 Reference Voice",
                                interactive=False,
                            )
                            gr.Textbox(
                                value=sample["ref_text"],
                                label="📝 Reference Text",
                                interactive=False,
                                lines=3,
                            )

                        gr.Audio(
                            value=sample["out_audio"],
                            label="🎧 Generated Sample (TTS)",
                            interactive=False,
                        )

                        if sample.get("gen_text"):
                            gr.Markdown(
                                f"**Text dùng để synth:** {sample['gen_text']}"
                            )

                        # Nút này sẽ fill luôn ref_audio, ref_text, gen_text, output_audio
                        use_btn = gr.Button(f"➡️ Dùng {sample['name']}")

                        use_btn.click(
                            make_sample_loader(sample),
                            inputs=[],
                            outputs=[ref_audio, ref_text, gen_text, output_audio],
                        )

        gr.Markdown(
            """
### ⚠️ Model Limitations
1. Có thể xử lý chưa tốt với số, ngày tháng, ký tự đặc biệt.  
2. Nhịp điệu đôi khi chưa tự nhiên.  
3. Chất lượng phụ thuộc khá nhiều vào chất lượng ref audio.  
"""
        )

demo.queue().launch()