import inspect
import os
import random
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path
from typing import Optional

import gradio as gr

MODEL_ID = "IndexTeam/IndexTTS-2"
MODEL_DIR = Path("/tmp/index-tts-2")
OUTPUT_DIR = Path(tempfile.gettempdir()) / "index-tts-2-outputs"

os.environ.setdefault("HF_HOME", "/tmp/huggingface")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

_runtime_lock = threading.Lock()
_runtime_ready = False
_tts = None


def _install_runtime_dependencies() -> None:
    """Install the official IndexTTS runtime without changing Space SDK pins."""
    global _runtime_ready
    if _runtime_ready:
        return

    with _runtime_lock:
        if _runtime_ready:
            return

        commands = [
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--no-cache-dir",
                "--no-deps",
                "tokenizers==0.21.0",
            ],
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--no-cache-dir",
                "--no-deps",
                "transformers==4.52.1",
            ],
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--no-cache-dir",
                "--no-deps",
                "git+https://github.com/index-tts/index-tts.git",
            ],
        ]

        for command in commands:
            subprocess.check_call(command)

        _runtime_ready = True


def _download_model() -> Path:
    """Download the model snapshot to local disk (CPU only, costs no ZeroGPU quota)."""
    from huggingface_hub import snapshot_download

    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    marker = MODEL_DIR / ".snapshot-complete"

    if not marker.exists():
        snapshot_download(
            repo_id=MODEL_ID,
            local_dir=str(MODEL_DIR),
        )
        marker.touch()

    return MODEL_DIR


def _download_aux_models() -> None:
    """Pre-download auxiliary models (w2v-bert-2.0, MaskGCT codec, CAMPPlus,
    BigVGAN) into the model cache on the CPU.

    IndexTTS2 downloads these itself when given aux_paths=None, which would
    happen inside the ZeroGPU allocation and waste quota on ~5 GB of downloads.
    """
    from indextts.utils.model_download import ensure_models_available

    ensure_models_available(str(MODEL_DIR))


def _patch_numpy_compat() -> None:
    """Restore NumPy aliases removed in NumPy 2.x for legacy pinned libraries.

    Old tensorboard/keras code (imported via audiotools) references np.bool8,
    which no longer exists in NumPy 2.x, crashing synthesis at import time.
    """
    import numpy as np

    if not hasattr(np, "bool8"):
        np.bool8 = np.bool_


ASSETS_DIR = Path(__file__).resolve().parent / "assets"
MY_VOICES = [
    ("🎙️ My Voice 1", ASSETS_DIR / "voice_1.wav"),
    ("🎙️ My Voice 2", ASSETS_DIR / "voice_2.wav"),
]
MY_VOICES = [(label, str(path)) for label, path in MY_VOICES if path.is_file()]
DEFAULT_VOICE = MY_VOICES[0][1] if MY_VOICES else None


def _warm_up() -> None:
    """Prepare dependencies and model files at Space startup, on the CPU.

    This keeps the paid ZeroGPU allocation focused on actual synthesis instead
    of pip installs and multi-GB downloads, so the daily GPU quota lasts much
    longer.
    """
    try:
        _patch_numpy_compat()
        _install_runtime_dependencies()
        _download_model()
        _download_aux_models()
        print("[warmup] runtime dependencies, model snapshot and aux models are ready", flush=True)
    except Exception as exc:
        print(
            f"[warmup] deferred, will retry on first request: {type(exc).__name__}: {exc}",
            flush=True,
        )


_warmup_thread = threading.Thread(target=_warm_up, daemon=True)
_warmup_thread.start()

def _find_config(model_dir: Path) -> str:
    candidates = [
        model_dir / "config.yaml",
        model_dir / "config.yml",
        model_dir / "configs" / "inference.yaml",
        model_dir / "configs" / "config.yaml",
    ]

    for candidate in candidates:
        if candidate.is_file():
            return str(candidate)

    discovered = sorted(model_dir.rglob("*.yaml")) + sorted(model_dir.rglob("*.yml"))
    if discovered:
        return str(discovered[0])

    raise FileNotFoundError(f"No IndexTTS configuration file was found under {model_dir}.")


def _get_tts():
    global _tts

    if _tts is not None:
        return _tts

    with _runtime_lock:
        if _tts is not None:
            return _tts

        try:
            _patch_numpy_compat()
            _install_runtime_dependencies()
            model_dir = _download_model()
            from indextts.infer_v2 import IndexTTS2

            _tts = IndexTTS2(
                cfg_path=_find_config(model_dir),
                model_dir=str(model_dir),
                use_fp16=False,
                device="cpu",
                use_cuda_kernel=False,
                use_deepspeed=False,
                use_accel=False,
                use_torch_compile=False,
                use_qwen_emo=True,
            )
            return _tts
        except Exception:
            _tts = None
            raise


def _validate_audio(path: Optional[str], label: str) -> str:
    if not path:
        raise gr.Error(f"{label} is required.")

    audio = Path(path)
    if not audio.is_file():
        raise gr.Error(f"{label} could not be found.")
    if audio.stat().st_size == 0:
        raise gr.Error(f"{label} is empty.")

    return str(audio)


def _clean_reference(path: str) -> str:
    """Trim silence, cap length and normalise the reference clip for a closer clone."""
    try:
        import librosa
        import numpy as np
        import soundfile as sf

        audio, sample_rate = librosa.load(path, sr=None, mono=True)
        trimmed, _ = librosa.effects.trim(audio, top_db=30)
        if len(trimmed) < sample_rate // 2:
            trimmed = audio
        trimmed = trimmed[: int(sample_rate * 20.0)]
        peak = float(np.max(np.abs(trimmed))) if trimmed.size else 0.0
        if peak > 0:
            trimmed = trimmed * (0.9 / peak)
        cleaned = f"{path}.cleaned.wav"
        sf.write(cleaned, trimmed, sample_rate)
        return cleaned
    except Exception as exc:
        print(f"[reference-clean] skipped: {type(exc).__name__}: {exc}", flush=True)
        return path


def _set_seed(seed: int) -> None:
    random.seed(seed)

    try:
        import numpy as np

        np.random.seed(seed)
    except ImportError:
        pass

    try:
        import torch

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:
        pass


def _supported_kwargs(function, values: dict) -> dict:
    """Keep compatibility with small upstream inference API changes."""
    parameters = inspect.signature(function).parameters
    if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()):
        return values
    return {key: value for key, value in values.items() if key in parameters}


def synthesize(
    reference_audio: Optional[str],
    text: str,
    permission: bool,
    emotion_audio: Optional[str],
    emotion_strength: float,
    use_text_emotion: bool,
    seed,
):
    if not permission:
        raise gr.Error("Please confirm that you have permission to use the reference voice.")

    reference_audio = _clean_reference(_validate_audio(reference_audio, "Reference audio"))
    text = (text or "").strip()

    if not text:
        raise gr.Error("English text is required.")
    if len(text) > 2000:
        raise gr.Error("Please keep the English text to 2,000 characters or fewer.")

    try:
        numeric_seed = float(seed)
        if not numeric_seed.is_integer() or numeric_seed < 0:
            raise ValueError
        numeric_seed = int(numeric_seed)
    except (TypeError, ValueError, OverflowError) as exc:
        raise gr.Error("Seed must be a non-negative integer.") from exc

    try:
        emotion_strength = float(emotion_strength)
    except (TypeError, ValueError) as exc:
        raise gr.Error("Emotion strength must be a number between 0 and 1.") from exc

    if not 0 <= emotion_strength <= 1:
        raise gr.Error("Emotion strength must be between 0 and 1.")

    if not isinstance(use_text_emotion, bool):
        raise gr.Error("Use text emotion must be a boolean.")

    emotion_audio = (
        _clean_reference(_validate_audio(emotion_audio, "Emotion audio")) if emotion_audio else None
    )
    output_dir = OUTPUT_DIR
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "index-tts-2-output.wav"

    if output_path.exists():
        output_path.unlink()

    started = time.perf_counter()
    try:
        _set_seed(numeric_seed)
        tts = _get_tts()

        values = {
            "spk_audio_prompt": reference_audio,
            "text": text,
            "output_path": str(output_path),
            "use_random": False,
        }

        if emotion_audio:
            values.update(
                {
                    "emo_audio_prompt": emotion_audio,
                    "emo_alpha": emotion_strength,
                }
            )

        if use_text_emotion:
            values.update(
                {
                    "use_emo_text": True,
                    "emo_alpha": min(emotion_strength, 0.6),
                }
            )

        tts.infer(**_supported_kwargs(tts.infer, values))

        if not output_path.is_file() or output_path.stat().st_size == 0:
            raise RuntimeError("IndexTTS-2 returned without creating an audio file.")

        elapsed = round(time.perf_counter() - started, 1)
        return str(output_path), elapsed
    except gr.Error:
        raise
    except Exception as exc:
        raise gr.Error(f"Synthesis failed: {type(exc).__name__}: {exc}") from exc


# --------------------------------------------------------------------------
# Quota panel helpers (CPU only, never touches the GPU)
# --------------------------------------------------------------------------

DAILY_HINT = "CPU mode has no daily quota — generate as much as you like. It is just slower than GPU."


def _read_ip_token(request) -> dict:
    """Decode the ZeroGPU identity token HF attaches to every request."""
    import base64
    import json

    headers = getattr(request, "headers", None) or {}
    token = headers.get("x-ip-token") or headers.get("X-IP-Token")
    if not token:
        return {}
    try:
        payload = json.loads(base64.urlsafe_b64decode(f"{token.split('.')[1]}=="))
        return payload if isinstance(payload, dict) else {}
    except Exception:
        return {}


def _quota_html(request, used_seconds: float) -> str:
    claims = _read_ip_token(request)
    identity = claims.get("user") or ("signed-in visitor" if claims.get("sub") else "guest (not signed in)")

    used = max(0.0, float(used_seconds or 0))

    return f"""
    <div class="quota-card">
      <div class="quota-title">🖥️ CPU mode — unlimited generations</div>
      <div class="quota-row"><span>Generation time this session</span><b>~{used:.0f}s</b></div>
      <div class="quota-note">{DAILY_HINT}</div>
      <div class="quota-note">💡 Keep texts short — each generation takes a few minutes on CPU. For fast GPU results use the <a href="https://huggingface.co/spaces/Offexb/Ai" target="_blank" style="color:#818cf8;">GPU version</a>.</div>
    </div>
    """


def _update_session_usage(used_seconds, last_elapsed, request: gr.Request):
    total = float(used_seconds or 0) + float(last_elapsed or 0)
    return total, _quota_html(request, total)


def _refresh_quota(request: gr.Request):
    return _quota_html(request, 0)


# --------------------------------------------------------------------------
# Premium UI
# --------------------------------------------------------------------------

CUSTOM_CSS = """
.gradio-container {max-width: 1180px !important; margin: 0 auto !important;}
html, body {overflow-x: hidden !important;}
.gradio-container .block, .gradio-container .form, .gradio-container .wrap,
.gradio-container .gap, .gradio-container .form-row, .gradio-container .component-wrapper {min-width: 0 !important; max-width: 100% !important;}
#hero {
  background: linear-gradient(120deg, #4f46e5 0%, #7c3aed 45%, #db2777 100%);
  border-radius: 18px; padding: 26px 30px; margin-bottom: 8px;
  box-shadow: 0 12px 32px rgba(79, 70, 229, 0.35);
}
#hero h1 {color: #ffffff; font-size: 2rem; font-weight: 800; margin: 0 0 6px 0; letter-spacing: -0.02em;}
#hero p {color: rgba(255, 255, 255, 0.88); margin: 0; font-size: 0.98rem;}
#hero .hero-badges {margin-top: 12px;}
#hero .badge {
  display: inline-block; background: rgba(255, 255, 255, 0.16); color: #fff;
  border: 1px solid rgba(255, 255, 255, 0.25); border-radius: 999px;
  padding: 3px 12px; font-size: 0.78rem; margin-right: 6px;
}
#generate-btn {
  background: linear-gradient(120deg, #4f46e5, #db2777) !important;
  border: none !important; color: #fff !important; font-weight: 700 !important;
  font-size: 1.02rem !important; height: 52px !important;
  box-shadow: 0 8px 22px rgba(219, 39, 119, 0.45) !important;
}
#generate-btn:hover {filter: brightness(1.08);}
.quota-card {
  background: #0f172a; border: 1px solid #1e293b; border-radius: 14px;
  padding: 16px 20px; margin-top: 10px;
}
.quota-title {font-weight: 700; font-size: 1rem; margin-bottom: 10px; color: #e2e8f0;}
.quota-row {display: flex; justify-content: space-between; font-size: 0.9rem; margin: 4px 0; color: #94a3b8;}
.quota-row b {color: #e2e8f0;}
.quota-bar {height: 10px; background: #1e293b; border-radius: 999px; overflow: hidden; margin: 10px 0;}
.quota-fill {height: 100%; border-radius: 999px; transition: width .4s ease;}
.quota-note {font-size: 0.78rem; color: #64748b; margin-top: 6px; line-height: 1.5;}
.tips-card {font-size: 0.92rem; line-height: 1.7;}
footer {display: none !important;}
#foot {text-align: center; color: #64748b; margin-top: 20px; font-size: 0.85rem;}
#foot span {opacity: .75;}
#hero {word-break: break-word;}
body {overflow-x: hidden;}
@media (max-width: 768px) {
  #hero {padding: 16px 14px; border-radius: 14px;}
  #hero h1 {font-size: 1.3rem;}
  #hero p {font-size: 0.82rem;}
  #adv-row {flex-direction: column !important;}
  #voice-btns {flex-wrap: wrap !important; gap: 8px !important;}
  #voice-btns button {flex: 1 1 45% !important; min-width: 0 !important;}
  .gradio-container {padding: 8px !important;}
  .quota-row {flex-direction: column; gap: 1px;}
  #generate-btn {height: 46px !important; font-size: 0.95rem !important;}
}
"""

HEADER_HTML = """
<div id="hero">
  <h1>🎙️ Offex Voice Studio</h1>
  <p>Clone any voice from a short clip — high fidelity, powered by IndexTTS-2.</p>
  <div class="hero-badges">
    <span class="badge">🎧 Voice cloning</span>
    <span class="badge">🎭 Emotion control</span>
    <span class="badge">⚡ Quota-smart</span>
  </div>
</div>
"""

TIPS_MD = """
**For the closest possible clone:**
1. **Reference clip:** use 5–15 seconds of **clean speech** from one speaker only (no music, no background noise). The app auto-trims silence and normalises volume, but it cannot remove background noise.
2. **Emotion reference:** add a second clip of the same voice speaking with the emotion you want — this dramatically improves expressiveness.
3. **Keep texts short** — shorter texts synthesize faster and use less of your daily GPU quota.
4. **Seed:** keep the same seed to get consistent output across retries; change it to explore variations.
"""

with gr.Blocks(title="Offex Voice Studio — AI Voice Cloning", theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"), css=CUSTOM_CSS) as demo:
    gr.HTML(HEADER_HTML)

    reference_audio = gr.Audio(
        label="Reference voice",
        sources=["upload", "microphone"],
        type="filepath",
        value=DEFAULT_VOICE,
    )
    if MY_VOICES:
        gr.Markdown("**⚡ Quick-pick — no upload needed:**")
        with gr.Row(elem_id="voice-btns"):
            for label, path in MY_VOICES:
                gr.Button(label, size="sm").click(
                    fn=lambda p=path: p, outputs=reference_audio
                )
    with gr.Accordion("🎭 Emotion reference (optional — big boost)", open=False):
        emotion_audio = gr.Audio(
            label="Emotion reference",
            sources=["upload", "microphone"],
            type="filepath",
        )
    text = gr.Textbox(
        label="Text to speak",
        placeholder="Type the text you want in the cloned voice...",
        lines=5,
        max_lines=12,
    )
    permission = gr.Checkbox(
        label="✅ I have permission to use this voice",
        value=False,
    )
    generate = gr.Button("✨ Generate Voice", variant="primary", size="lg", elem_id="generate-btn")
    output = gr.Audio(label="🔊 Generated voice", type="filepath", autoplay=True, show_download_button=True)

    with gr.Accordion("⚙️ Advanced settings", open=False):
        with gr.Row(elem_id="adv-row"):
            emotion_strength = gr.Slider(
                minimum=0,
                maximum=1,
                value=0.6,
                step=0.05,
                label="Emotion strength",
            )
            use_text_emotion = gr.Checkbox(
                label="Derive emotion from the text itself",
                value=False,
            )
            seed = gr.Number(
                label="Seed (same seed = consistent output)",
                value=1234,
                precision=0,
                minimum=0,
            )

    usage_state = gr.State(0.0)
    elapsed_box = gr.Number(visible=False, value=0)
    quota_html = gr.HTML(elem_id="quota-panel")

    generate.click(
        fn=synthesize,
        inputs=[
            reference_audio,
            text,
            permission,
            emotion_audio,
            emotion_strength,
            use_text_emotion,
            seed,
        ],
        outputs=[output, elapsed_box],
    ).then(
        fn=_update_session_usage,
        inputs=[usage_state, elapsed_box],
        outputs=[usage_state, quota_html],
    )

    demo.load(fn=_refresh_quota, inputs=None, outputs=quota_html)

    with gr.Accordion("💡 Tips for the most accurate clone", open=False):
        gr.Markdown(TIPS_MD)

    gr.HTML("<div id='foot'>Built with <span>❤️</span> in India · Powered by IndexTTS-2 on ZeroGPU</div>")

    if MY_VOICES:
        gr.Examples(
            examples=[
                [MY_VOICES[0][1], "Hello there! This is my cloned voice, speaking naturally and clearly."],
                [MY_VOICES[-1][1], "Welcome to Offex Voice Studio. Today is a wonderful day to create something amazing."],
            ],
            inputs=[reference_audio, text],
            label="One-tap examples with your saved voices",
        )


if __name__ == "__main__":
    demo.queue(max_size=8).launch(ssr_mode=False)