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
import spaces

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
_warmup_state = {"done": False, "started_at": time.time()}


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 Voice 3", ASSETS_DIR / "voice_3.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()
        _warmup_state["done"] = True
        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:
            import torch

            if not torch.cuda.is_available():
                raise RuntimeError("CUDA GPU is not available in this Space.")

            _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=True,
                device="cuda",
                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 _estimate_duration(*args, **kwargs) -> int:
    """Request only as much ZeroGPU time as this request realistically needs.

    The platform pre-checks the declared duration against the visitor's
    remaining quota, so smaller realistic durations let people fit far more
    generations into their daily allowance. Accepts *args/**kwargs because
    Gradio may inject extra parameters (e.g. gr.Progress).
    """
    if _tts is None:
        return 120  # cold start: one-time model load plus synthesis

    text = kwargs.get("text")
    if not text and len(args) > 1:
        text = args[1]
    characters = len((text or "").strip())
    return int(min(120, 30 + characters * 0.05))


@spaces.GPU(duration=_estimate_duration)
def synthesize(
    reference_audio: Optional[str],
    text: str,
    permission: bool,
    emotion_audio: Optional[str],
    emotion_strength: float,
    use_text_emotion: bool,
    seed,
    progress=gr.Progress(),
):
    progress(0.05, desc="Validating input")

    if not permission:
        raise gr.Error("Please confirm that you have permission to use the reference voice.")

    progress(0.1, desc="Preparing reference audio")

    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)
        if _tts is None:
            progress(0.18, desc="Loading AI model - first run takes 1-3 min")
        else:
            progress(0.45, desc="Model warm - starting synthesis")
        tts = _get_tts()
        progress(0.6, desc="Synthesizing voice")

        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),
                }
            )

        _ticker_stop = threading.Event()

        def _tick_progress() -> None:
            value = 0.62
            while not _ticker_stop.wait(2.0) and value < 0.92:
                value += 0.01
                try:
                    progress(value, desc="Synthesizing voice")
                except Exception:
                    break

        _ticker = threading.Thread(target=_tick_progress, daemon=True)
        _ticker.start()
        try:
            tts.infer(**_supported_kwargs(tts.infer, values))
        finally:
            _ticker_stop.set()

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

        progress(0.97, desc="Saving 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 = "Free accounts get ~5 min/day, PRO ~40 min/day. The window resets 24h after your first use."


def backend_status():
    """Lightweight, non-GPU status probe for external frontends."""
    return {
        "status": "ready" if _warmup_state["done"] else "preparing",
        "model_loaded": _tts is not None,
        "uptime_minutes": round((time.time() - _warmup_state["started_at"]) / 60, 1),
    }


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))
    bar = min(100, round(used / 300 * 100))  # free-tier daily allowance estimate
    color = "#22c55e" if bar < 50 else "#eab308" if bar < 80 else "#ef4444"

    return f"""
    <div class="quota-card">
      <div class="quota-title">⚡ ZeroGPU quota</div>
      <div class="quota-row"><span>Identity</span><b>{identity}</b></div>
      <div class="quota-row"><span>GPU time used on this Space (this session)</span><b>~{used:.0f}s</b></div>
      <div class="quota-bar"><div class="quota-fill" style="width:{bar}%;background:{color};"></div></div>
      <div class="quota-note">Bar = your session usage vs a ~5 min free daily allowance (estimate). {DAILY_HINT}</div>
      <div class="quota-note">💡 This Space auto-trims requests to the minimum GPU time needed, so your quota lasts longer. Upgrade to <a href="https://huggingface.co/pricing" target="_blank" style="color:#818cf8;">HF PRO</a> for ~8x quota.</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)

    status_probe = gr.Button(visible=False)
    status_out = gr.JSON(visible=False)
    status_probe.click(fn=backend_status, inputs=None, outputs=status_out, api_name="backend_status", queue=False)

    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)