import io
import os
import random
from typing import Optional

import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download

from irodori_tts.inference_runtime import (
    InferenceRuntime,
    RuntimeKey,
    SamplingRequest,
)

# ---------------------------------------------------------------------------
# Konfigurasi CPU untuk kecepatan maksimal
# ---------------------------------------------------------------------------
torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", "4")))
torch.set_default_dtype(torch.float32)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MODEL_REPO = os.environ.get("MODEL_REPO", "Aratako/Irodori-TTS-500M")
CODEC_REPO = "facebook/dacvae-watermarked"
FIXED_SECONDS = 30.0
MAX_GRADIO_CANDIDATES = int(os.environ.get("MAX_GRADIO_CANDIDATES", "32"))
GRADIO_AUDIO_COLS_PER_ROW = 8

# Global state
_runtime: Optional[InferenceRuntime] = None

# Contoh teks acak (bahasa Jepang)
EXAMPLE_TEXTS = [
    "こんにちは、今日はいい天気ですね。",
    "私は毎朝コーヒーを飲みます。",
    "この曲はとても美しいです。",
    "明日、友達と映画を見に行きます。",
    "富士山は日本で一番高い山です。",
    "夏休みは海に行きたいです。",
    "日本語の勉強は楽しいです。",
    "ありがとうございます。",
]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _parse_optional_float(raw: str | None, label: str) -> float | None:
    if raw is None:
        return None
    text = str(raw).strip()
    if text == "" or text.lower() == "none":
        return None
    try:
        return float(text)
    except ValueError as exc:
        raise ValueError(f"{label} must be a float or blank.") from exc


def _parse_optional_int(raw: str | None, label: str) -> int | None:
    if raw is None:
        return None
    text = str(raw).strip()
    if text == "" or text.lower() == "none":
        return None
    try:
        return int(text)
    except ValueError as exc:
        raise ValueError(f"{label} must be an int or blank.") from exc


# ---------------------------------------------------------------------------
# Model Loading (dioptimasi untuk CPU)
# ---------------------------------------------------------------------------

def load_models():
    global _runtime

    if _runtime is not None:
        return

    print(f"[Info] Downloading checkpoint from {MODEL_REPO}...")
    checkpoint_path = hf_hub_download(
        repo_id=MODEL_REPO, filename="model.safetensors"
    )

    device = "cpu"
    precision = "fp32"

    key = RuntimeKey(
        checkpoint=checkpoint_path,
        model_device=device,
        codec_repo=CODEC_REPO,
        model_precision=precision,
        codec_device=device,
        codec_precision=precision,
        enable_watermark=False,
    )

    print("[Info] Building runtime (CPU mode)...")
    _runtime = InferenceRuntime.from_key(key)
    print("[Info] All models loaded successfully on CPU.")


# Load models at startup
load_models()


# ---------------------------------------------------------------------------
# Inference (tanpa dekorator GPU agar tidak crash di CPU)
# ---------------------------------------------------------------------------

def run_inference(
    text: str,
    uploaded_audio: Optional[str],
    num_steps: int,
    num_candidates: int,
    seed_raw: str,
    cfg_guidance_mode: str,
    cfg_scale_text: float,
    cfg_scale_speaker: float,
    cfg_scale_raw: str,
    cfg_min_t: float,
    cfg_max_t: float,
    context_kv_cache: bool,
    truncation_factor_raw: str,
    rescale_k_raw: str,
    rescale_sigma_raw: str,
    speaker_kv_scale_raw: str,
    speaker_kv_min_t_raw: str,
    speaker_kv_max_layers_raw: str,
) -> tuple[list[tuple[int, np.ndarray]], str]:
    load_models()

    log_buffer = io.StringIO()

    def stdout_log(msg: str) -> None:
        print(msg, flush=True)
        log_buffer.write(msg + "\n")

    if not str(text).strip():
        raise gr.Error("Please enter text to synthesize.")

    cfg_scale = _parse_optional_float(cfg_scale_raw, "cfg_scale")
    truncation_factor = _parse_optional_float(
        truncation_factor_raw, "truncation_factor"
    )
    rescale_k = _parse_optional_float(rescale_k_raw, "rescale_k")
    rescale_sigma = _parse_optional_float(rescale_sigma_raw, "rescale_sigma")
    speaker_kv_scale = _parse_optional_float(
        speaker_kv_scale_raw, "speaker_kv_scale"
    )
    speaker_kv_min_t = _parse_optional_float(
        speaker_kv_min_t_raw, "speaker_kv_min_t"
    )
    speaker_kv_max_layers = _parse_optional_int(
        speaker_kv_max_layers_raw, "speaker_kv_max_layers"
    )
    seed = _parse_optional_int(seed_raw, "seed")
    requested_candidates = int(num_candidates)
    if requested_candidates <= 0:
        raise gr.Error("num_candidates must be >= 1.")
    if requested_candidates > MAX_GRADIO_CANDIDATES:
        raise gr.Error(
            f"num_candidates must be <= {MAX_GRADIO_CANDIDATES}."
        )

    ref_wav: str | None = None
    no_ref = True
    if uploaded_audio is not None and str(uploaded_audio).strip() != "":
        ref_wav = str(uploaded_audio)
        no_ref = False

    result = _runtime.synthesize(
        SamplingRequest(
            text=str(text),
            ref_wav=ref_wav,
            ref_latent=None,
            no_ref=bool(no_ref),
            num_candidates=requested_candidates,
            decode_mode="sequential",
            seconds=FIXED_SECONDS,
            max_ref_seconds=30.0,
            max_text_len=None,
            num_steps=int(num_steps),
            seed=None if seed is None else int(seed),
            cfg_guidance_mode=str(cfg_guidance_mode),
            cfg_scale_text=float(cfg_scale_text),
            cfg_scale_speaker=float(cfg_scale_speaker),
            cfg_scale=cfg_scale,
            cfg_min_t=float(cfg_min_t),
            cfg_max_t=float(cfg_max_t),
            truncation_factor=truncation_factor,
            rescale_k=rescale_k,
            rescale_sigma=rescale_sigma,
            context_kv_cache=bool(context_kv_cache),
            speaker_kv_scale=speaker_kv_scale,
            speaker_kv_min_t=speaker_kv_min_t,
            speaker_kv_max_layers=speaker_kv_max_layers,
            trim_tail=True,
        ),
        log_fn=stdout_log,
    )

    sample_rate = result.sample_rate
    audio_results: list[tuple[int, np.ndarray]] = []
    for audio in result.audios:
        waveform = audio.squeeze(0).float().numpy()
        audio_results.append((sample_rate, waveform))
    return audio_results, log_buffer.getvalue()


# ---------------------------------------------------------------------------
# Gradio UI (Enhanced with orange random button, credit, warning text & loading)
# ---------------------------------------------------------------------------

def build_demo():
    # CSS hijau-putih + tombol random oranye
    custom_css = """
    body, .gradio-container {
        background-color: #ffffff;
        font-family: 'Inter', sans-serif;
    }
    .gr-button-primary {
        background-color: #2e7d32 !important;
        border: none !important;
        color: white !important;
        transition: 0.2s;
    }
    .gr-button-primary:hover {
        background-color: #1b5e20 !important;
    }
    .orange-random {
        background-color: #f57c00 !important;
        color: white !important;
        border: none !important;
        font-weight: bold !important;
    }
    .orange-random:hover {
        background-color: #e65100 !important;
    }
    .gr-box, .gr-form, .input-panel, .output-panel {
        border-color: #c8e6c9 !important;
    }
    h1, h2, h3, label {
        color: #1b5e20 !important;
    }
    .gr-textbox, .gr-dropdown, .gr-slider {
        border-color: #a5d6a7 !important;
    }
    footer {
        display: none !important;
    }
    .credit-footer {
        text-align: center;
        margin-top: 30px;
        padding: 10px;
        font-size: 14px;
        color: #2e7d32;
        border-top: 1px solid #c8e6c9;
        width: 100%;
    }
    .warning-text {
        text-align: center;
        background-color: #fff3e0;
        color: #e65100;
        padding: 8px;
        border-radius: 10px;
        margin-top: 10px;
        margin-bottom: 10px;
        font-weight: bold;
        border-left: 4px solid #f57c00;
    }
    .header-image {
        margin-bottom: 1rem;
        border-radius: 12px;
    }
    """

    HEADER_IMAGE_URL = "https://huggingface.co/spaces/Plana-Archive/Irodori-TTS/resolve/main/latest.png"

    with gr.Blocks(css=custom_css, title="Irodori-TTS Demo") as demo:
        # Header gambar
        gr.Image(value=HEADER_IMAGE_URL, show_label=False, elem_classes="header-image", height=200)

        gr.Markdown("🍊 Irodori-TTS-500M 🍊")
        gr.Markdown("Generate Voice Anime")

        # Baris input teks + tombol random (oranye)
        with gr.Row():
            text = gr.Textbox(label="📝 Input Teks", lines=3, scale=6, placeholder="Masukkan teks bahasa Jepang di sini...")
            random_btn = gr.Button("🎲 Random Text 🎲", scale=1, variant="secondary", elem_classes="orange-random")

        # Audio referensi
        uploaded_audio = gr.Audio(
            label="🎵 Audio Referensi (opsional, kosongkan untuk tanpa referensi)",
            type="filepath",
        )

        # Parameter sampling
        with gr.Accordion("⚙️ Pengaturan Sampling", open=True):
            with gr.Row():
                num_steps = gr.Slider(label="Langkah", minimum=1, maximum=120, value=40, step=1)
                num_candidates = gr.Slider(label="Jumlah Kandidat", minimum=1, maximum=MAX_GRADIO_CANDIDATES, value=1, step=1)
                seed_raw = gr.Textbox(label="Seed (kosong = acak)", value="")

            with gr.Row():
                cfg_guidance_mode = gr.Dropdown(
                    label="Mode CFG",
                    choices=["independent", "joint", "alternating"],
                    value="independent",
                )
                cfg_scale_text = gr.Slider(label="Skala CFG Teks", minimum=0.0, maximum=10.0, value=3.0, step=0.1)
                cfg_scale_speaker = gr.Slider(label="Skala CFG Speaker", minimum=0.0, maximum=10.0, value=5.0, step=0.1)

        with gr.Accordion("🔧 Lanjutan (Opsional)", open=False):
            cfg_scale_raw = gr.Textbox(label="Timpa Skala CFG", value="")
            with gr.Row():
                cfg_min_t = gr.Number(label="CFG Min t", value=0.5)
                cfg_max_t = gr.Number(label="CFG Max t", value=1.0)
                context_kv_cache = gr.Checkbox(label="Cache KV Context", value=True)
            with gr.Row():
                truncation_factor_raw = gr.Textbox(label="Faktor Trunkasi", value="")
                rescale_k_raw = gr.Textbox(label="Rescale k", value="")
                rescale_sigma_raw = gr.Textbox(label="Rescale sigma", value="")
            with gr.Row():
                speaker_kv_scale_raw = gr.Textbox(label="Skala KV Speaker", value="")
                speaker_kv_min_t_raw = gr.Textbox(label="KV Speaker Min t", value="0.9")
                speaker_kv_max_layers_raw = gr.Textbox(label="KV Speaker Max Layers", value="")

        # Tombol generate
        generate_btn = gr.Button("💫 Generate Voice 💫", variant="primary")

        # Tambahan teks peringatan di bawah tombol generate
        warning_html = gr.HTML(
            '<div class="warning-text">🍁 WARNING MINNA 🍁<br>After generating the voice, it will appear in a few minutes, just wait.</div>'
        )

        # Layout audio hasil
        out_audios: list[gr.Audio] = []
        num_rows = (MAX_GRADIO_CANDIDATES + GRADIO_AUDIO_COLS_PER_ROW - 1) // GRADIO_AUDIO_COLS_PER_ROW
        with gr.Column():
            for row_idx in range(num_rows):
                with gr.Row():
                    for col_idx in range(GRADIO_AUDIO_COLS_PER_ROW):
                        i = row_idx * GRADIO_AUDIO_COLS_PER_ROW + col_idx
                        if i >= MAX_GRADIO_CANDIDATES:
                            break
                        out_audios.append(
                            gr.Audio(
                                label=f"Hasil {i + 1}",
                                type="numpy",
                                visible=(i == 0),
                            )
                        )
        out_log = gr.Textbox(label="📜 Log Proses", lines=6)

        # Credit footer
        gr.HTML('<div class="credit-footer">🌠 CREATED BY MUTSUMI 🌠</div>')

        # Fungsi random text
        def set_random_text():
            return random.choice(EXAMPLE_TEXTS)

        random_btn.click(fn=set_random_text, outputs=text)

        # Fungsi utama generasi dengan loading progress bar orange
        def gradio_inference(
            text_val,
            uploaded_audio_val,
            num_steps_val,
            num_candidates_val,
            seed_raw_val,
            cfg_guidance_mode_val,
            cfg_scale_text_val,
            cfg_scale_speaker_val,
            cfg_scale_raw_val,
            cfg_min_t_val,
            cfg_max_t_val,
            context_kv_cache_val,
            truncation_factor_raw_val,
            rescale_k_raw_val,
            rescale_sigma_raw_val,
            speaker_kv_scale_raw_val,
            speaker_kv_min_t_raw_val,
            speaker_kv_max_layers_raw_val,
            progress=gr.Progress(track_tqdm=True),  # Loading indicator orange
        ):
            try:
                audio_results, log_text = run_inference(
                    text=text_val,
                    uploaded_audio=uploaded_audio_val,
                    num_steps=num_steps_val,
                    num_candidates=num_candidates_val,
                    seed_raw=seed_raw_val,
                    cfg_guidance_mode=cfg_guidance_mode_val,
                    cfg_scale_text=cfg_scale_text_val,
                    cfg_scale_speaker=cfg_scale_speaker_val,
                    cfg_scale_raw=cfg_scale_raw_val,
                    cfg_min_t=cfg_min_t_val,
                    cfg_max_t=cfg_max_t_val,
                    context_kv_cache=context_kv_cache_val,
                    truncation_factor_raw=truncation_factor_raw_val,
                    rescale_k_raw=rescale_k_raw_val,
                    rescale_sigma_raw=rescale_sigma_raw_val,
                    speaker_kv_scale_raw=speaker_kv_scale_raw_val,
                    speaker_kv_min_t_raw=speaker_kv_min_t_raw_val,
                    speaker_kv_max_layers_raw=speaker_kv_max_layers_raw_val,
                )
                audio_updates: list[object] = []
                for i in range(MAX_GRADIO_CANDIDATES):
                    if i < len(audio_results):
                        audio_updates.append(gr.update(value=audio_results[i], visible=True))
                    else:
                        audio_updates.append(gr.update(value=None, visible=False))
                return tuple([*audio_updates, log_text])
            except Exception as e:
                raise gr.Error(str(e))

        generate_btn.click(
            fn=gradio_inference,
            inputs=[
                text,
                uploaded_audio,
                num_steps,
                num_candidates,
                seed_raw,
                cfg_guidance_mode,
                cfg_scale_text,
                cfg_scale_speaker,
                cfg_scale_raw,
                cfg_min_t,
                cfg_max_t,
                context_kv_cache,
                truncation_factor_raw,
                rescale_k_raw,
                rescale_sigma_raw,
                speaker_kv_scale_raw,
                speaker_kv_min_t_raw,
                speaker_kv_max_layers_raw,
            ],
            outputs=[*out_audios, out_log],
            show_progress="full",  # Menampilkan progress bar orange (Gradio built-in)
        )

    return demo


if __name__ == "__main__":
    demo = build_demo()
    demo.queue(default_concurrency_limit=1)
    demo.launch()