"""Gradio demo for Diamond speech restoration.

Designed to run both locally and on HuggingFace Spaces. The checkpoint
(`step_00100000.safetensors` + `.json` sidecar) is downloaded from
`kyrgyz-ai/diamond-s2s` on first launch and cached in ./ckpt/.

Local:
    uv run python space/app.py

HuggingFace Space: drop this folder at the root of a Gradio Space repo.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

import gradio as gr
import soundfile as sf
import torch
from huggingface_hub import hf_hub_download

# Best-effort import of the HF Spaces `spaces` package. On ZeroGPU hardware,
# @spaces.GPU schedules the call onto the shared GPU pool; on CPU hardware
# (or local dev) it's a no-op decorator.
try:
    import spaces  # type: ignore
except ImportError:
    class _NoSpaces:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn
            return decorator if (args and callable(args[0])) is False else args[0]
    spaces = _NoSpaces()  # type: ignore

# Allow `import speech_restoration` whether this file runs from repo root
# (uv run python space/app.py) or from the HF Space root (where src/ is at the
# same level as app.py — we add it explicitly to be safe).
HERE = Path(__file__).resolve().parent
REPO_ROOT = HERE.parent
for cand in (REPO_ROOT / "src", HERE / "src"):
    if cand.exists():
        sys.path.insert(0, str(cand))

from speech_restoration.infer.infer import (  # noqa: E402
    build_pipeline,
    load_audio,
    restore,
)

MODEL_REPO = os.environ.get("DIAMOND_REPO", "kyrgyz-ai/diamond-s2s")
MODEL_FILE = os.environ.get("DIAMOND_CKPT", "step_00100000.safetensors")
CKPT_DIR = Path(os.environ.get("DIAMOND_CKPT_DIR", "ckpt"))


def _fetch_checkpoint() -> str:
    """Download checkpoint + config sidecar from HF Hub (cached)."""
    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    config_file = Path(MODEL_FILE).with_suffix(".json").name
    ckpt = hf_hub_download(MODEL_REPO, MODEL_FILE, local_dir=str(CKPT_DIR))
    hf_hub_download(MODEL_REPO, config_file, local_dir=str(CKPT_DIR))
    return ckpt


# ── One-time pipeline init (always on CPU at module load) ──────────────────
# Even on ZeroGPU hardware, module-level code runs in the CPU container before
# any @spaces.GPU function is invoked. The model migrates to cuda on first
# inference call.
print("Resolving device...")
INIT_DEVICE = torch.device("cpu")
print(f"Device at init: {INIT_DEVICE}")

print(f"Fetching checkpoint {MODEL_REPO}/{MODEL_FILE}...")
CKPT_PATH = _fetch_checkpoint()
print(f"Checkpoint: {CKPT_PATH}")

print("Building inference pipeline...")
MODEL, MEL, DAC, CFG, STEP = build_pipeline(CKPT_PATH, INIT_DEVICE)
print(f"Loaded checkpoint @ step {STEP}")

_PIPELINE_DEVICE = "cpu"


def _migrate_pipeline(target: str) -> None:
    """Move MODEL / MEL / DAC to `target` ('cuda' or 'cpu'). Idempotent."""
    global _PIPELINE_DEVICE
    if _PIPELINE_DEVICE == target:
        return
    MODEL.to(target)
    MEL.to(target)
    DAC.model.to(target)
    DAC.device = target
    _PIPELINE_DEVICE = target
    print(f"Pipeline migrated to {target}")


# ── Inference callback ─────────────────────────────────────────────────────

@spaces.GPU(duration=120)
@torch.no_grad()
def run(audio_path: str | None,
        chunk_sec: float,
        overlap_sec: float,
        trim_leadin: bool) -> tuple[int, "any"] | None:
    """Gradio callback: file path in → (sample_rate, waveform) tuple out.

    On ZeroGPU, this function is dispatched onto the shared GPU pool by the
    @spaces.GPU decorator; the pipeline is migrated to cuda on the first call.
    On CPU hardware (or local dev), the decorator is a no-op and inference
    runs on CPU.
    """
    if not audio_path:
        raise gr.Error("Upload a file or record from the microphone first.")

    target = "cuda" if torch.cuda.is_available() else "cpu"
    _migrate_pipeline(target)

    wav, sr = load_audio(audio_path)
    if wav.size == 0:
        raise gr.Error("Empty audio.")
    restored = restore(
        wav, sr, MODEL, MEL, DAC, CFG,
        chunk_sec=chunk_sec, overlap_sec=overlap_sec,
        warmup_sec=1.0, tail_pad_sec=1.0,
        trim_leadin=trim_leadin, normalize=True, verbose=True,
    )
    return DAC.sample_rate, restored


# ── UI ──────────────────────────────────────────────────────────────────────

EXAMPLES_DIR = HERE / "examples"

# Map degraded example path → pre-rendered restored path. When the user clicks
# an example, we return the pre-rendered file directly instead of running the
# model — this is what the user expected to compare "before/after" against, and
# it costs zero compute (no GPU, no inference). Files follow the convention
# <id>_degraded.<ext> + <id>_restored.<ext> living side by side in examples/.
# Keyed by the degraded file's basename — Gradio copies clicked examples to a
# session temp dir, so we can't match on full path. Basename is stable.
PRECOMPUTED_OUTPUTS: dict[str, str] = {}
_example_input_paths: list[Path] = []
for deg in sorted(EXAMPLES_DIR.glob("*_degraded.*")):
    stem = deg.name.removesuffix(deg.suffix).removesuffix("_degraded")
    for ext in (".wav", ".mp3", ".flac", ".ogg"):
        rest = EXAMPLES_DIR / f"{stem}_restored{ext}"
        if rest.exists():
            PRECOMPUTED_OUTPUTS[deg.name] = str(rest.resolve())
            _example_input_paths.append(deg)
            break

examples = [[str(p), 2.5, 0.4, True] for p in _example_input_paths] or None


def _load_audio_as_numpy(path: str):
    """Read a wav/mp3/flac as (sample_rate, mono ndarray) for Gradio Audio."""
    import numpy as np
    wav, sr = sf.read(path, dtype="float32", always_2d=False)
    if wav.ndim > 1:
        wav = wav.mean(axis=-1).astype(np.float32)
    return sr, wav


def run_example(audio_path: str | None, *_) -> tuple[int, "any"]:
    """Examples-only callback: hand back the pre-rendered restored audio for
    this degraded input. No model call, no GPU."""
    if not audio_path:
        raise gr.Error("Click an example below.")
    name = Path(audio_path).name
    if name not in PRECOMPUTED_OUTPUTS:
        # Defensive: unknown example → fall back to real inference.
        return run(audio_path, 2.5, 0.4, True)
    return _load_audio_as_numpy(PRECOMPUTED_OUTPUTS[name])

DESCRIPTION = """
# Diamond — Speech Restoration

Upload a degraded recording (phone call, low-bitrate, noisy mic) or record
from your microphone, and Diamond will resynthesize it as near-studio
44.1 kHz audio.

[Model card](https://huggingface.co/kyrgyz-ai/diamond-s2s) ·
[Code](https://github.com/KaniTTS-research-team/S2S-inference-diamond)
"""

with gr.Blocks(title="Diamond — Speech Restoration") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(
                label="Degraded input (upload or record)",
                sources=["upload", "microphone"],
                type="filepath",
            )
            with gr.Accordion("Advanced", open=False):
                chunk_sec = gr.Slider(
                    0.5, 6.0, value=2.5, step=0.1,
                    label="Chunk length (s)",
                    info="Decoder chunk size. Smaller = safer for long files, "
                         "more chunks. 0 not exposed here — use the CLI.")
                overlap_sec = gr.Slider(
                    0.0, 1.5, value=0.4, step=0.05,
                    label="Crossfade overlap (s)",
                    info="Overlap between adjacent chunks, crossfaded.")
                trim_leadin = gr.Checkbox(
                    value=True, label="Trim non-speech intro",
                    info="Cut room tone / handling noise before the first word. "
                         "Useful for phone recordings.")
            run_btn = gr.Button("Restore", variant="primary")

        with gr.Column():
            audio_out = gr.Audio(
                label="Restored output (44.1 kHz)",
                type="numpy",
                interactive=False,
            )

    run_btn.click(
        run,
        inputs=[audio_in, chunk_sec, overlap_sec, trim_leadin],
        outputs=audio_out,
    )

    if examples:
        gr.Examples(
            examples=examples,
            inputs=[audio_in, chunk_sec, overlap_sec, trim_leadin],
            outputs=audio_out,
            # Examples bypass real inference: run_example resolves the click
            # to a pre-rendered restored file (PRECOMPUTED_OUTPUTS map). No
            # model call, no GPU — output appears instantly.
            fn=run_example,
            # cache_examples=True: gradio pre-runs run_example for each row at
            # app launch (~1 s total, just reads the prerendered wavs from
            # disk). On click, BOTH the degraded input AND the cached restored
            # output are populated immediately — no waiting, no Restore press.
            cache_examples=True,
            label="Examples — click to hear Diamond's pre-rendered output",
        )


if __name__ == "__main__":
    demo.queue(max_size=8).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        show_error=True,
    )