"""audio.cpp runtime helpers for the Hugging Face Space.

Everything audio.cpp-specific lives here so it can be imported both by the
Gradio app and by user code in the Python tab:

    import audiocpp_runtime as ac
    ac.list_packages(task="tts")
    result = ac.run(family="pocket_tts", task="tts", request={"text": "Hello"})
"""
from __future__ import annotations

import base64
import ctypes
import glob
import json
import os
import shutil
import socket
import subprocess
import sys
import tarfile
import threading
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path
from typing import Any

APP_DIR = Path(__file__).resolve().parent
AUDIOCPP_VERSION = os.environ.get("AUDIOCPP_VERSION", "v0.9.0")
ASSET_NAME = "audio-{tag}-bin-ubuntu-x64-cuda12.8-colab.tar.gz"
RELEASE_URL = "https://github.com/0xShug0/audio.cpp/releases/download/{tag}/{asset}"

HOME = Path(os.environ.get("AUDIOCPP_HOME", Path.home() / ".audiocpp"))
BIN_ROOT = HOME / "bin"
WORK_ROOT = HOME / "work"


def _pick_models_root() -> Path:
    override = os.environ.get("AUDIOCPP_MODELS_DIR")
    if override:
        return Path(override)
    data = Path("/data")
    if data.is_dir() and os.access(data, os.W_OK):
        return data / "audiocpp-models"  # persistent storage, if the Space has it
    return HOME / "models"


MODELS_ROOT = _pick_models_root()
STUB_SRC = APP_DIR / "stub" / "cuda_stub.c"
STUB_DIR = HOME / "stub"

# Spec task names -> audio.cpp runtime task tokens.
TASK_TOKEN = {
    "tts": "tts", "clone": "clon", "design": "vdes", "asr": "asr", "align": "align",
    "vad": "vad", "diar": "diar", "sep": "sep", "vc": "vc", "svc": "svc", "s2s": "s2s",
    "music": "gen", "sfx": "gen", "edit": "gen", "midi": "midi", "turn": "turn",
    "speaker": "spk",
}
TASK_LABELS = {
    "tts": "Text to speech",
    "clone": "Voice cloning",
    "design": "Voice design",
    "asr": "Speech recognition (ASR)",
    "align": "Forced alignment",
    "vad": "Voice activity detection",
    "diar": "Speaker diarization",
    "sep": "Source separation",
    "vc": "Voice conversion",
    "svc": "Singing voice conversion",
    "s2s": "Speech to speech / restoration",
    "music": "Music generation",
    "sfx": "Sound effects",
    "edit": "Audio editing",
    "midi": "Audio to MIDI / score",
    "turn": "Turn detection",
}
# Families whose loader wants the package directory rather than one GGUF file.
DIR_FAMILIES = {"auk", "lfm2_audio", "liveavatar", "minimax_music3", "yue2"}
# Preset voices documented in docs/tts.md but missing from some specs' "ui" block.
VOICE_FALLBACKS = {
    "supertonic": ([f"M{i}" for i in range(1, 6)] + [f"F{i}" for i in range(1, 6)], "M1"),
}

_setup_lock = threading.Lock()
_install_lock = threading.Lock()
_state: dict[str, Any] = {"ready": False, "error": None, "server": None}


# --------------------------------------------------------------------------- setup
def log(msg: str) -> None:
    print(f"[audiocpp] {msg}", flush=True)


def ensure_runtime() -> Path:
    """Download and unpack the audio.cpp CUDA prebuilt once. Returns the bundle dir."""
    with _setup_lock:
        if _state["ready"]:
            return _state["bundle"]
        bundle = BIN_ROOT / AUDIOCPP_VERSION
        server = bundle / "audiocpp_server"
        if not server.exists():
            bundle.mkdir(parents=True, exist_ok=True)
            asset = ASSET_NAME.format(tag=AUDIOCPP_VERSION)
            url = RELEASE_URL.format(tag=AUDIOCPP_VERSION, asset=asset)
            archive = BIN_ROOT / asset
            log(f"downloading {url}")
            req = urllib.request.Request(url, headers={"User-Agent": "audiocpp-hf-space"})
            with urllib.request.urlopen(req, timeout=600) as resp, open(archive, "wb") as out:
                shutil.copyfileobj(resp, out, length=1 << 20)
            with tarfile.open(archive, "r:gz") as tar:
                tar.extractall(bundle, filter="data")
            archive.unlink(missing_ok=True)
            found = list(bundle.rglob("audiocpp_server"))
            if not found:
                raise RuntimeError("release archive has no audiocpp_server")
            server = found[0]
        server.chmod(0o755)
        MODELS_ROOT.mkdir(parents=True, exist_ok=True)
        WORK_ROOT.mkdir(parents=True, exist_ok=True)
        build_cuda_stub()
        _state.update(ready=True, bundle=server.parent, server=server)
        log(f"runtime ready: {server}")
        return server.parent


def build_cuda_stub() -> None:
    """Compile the tiny libcuda.so.1 stand-in used for CPU-only runs."""
    target = STUB_DIR / "libcuda.so.1"
    if target.exists():
        return
    STUB_DIR.mkdir(parents=True, exist_ok=True)
    prebuilt = APP_DIR / "stub" / "libcuda.so.1"
    if prebuilt.exists():
        shutil.copy2(prebuilt, target)
        return
    cc = shutil.which("gcc") or shutil.which("cc")
    if not cc:
        log("no C compiler found; CPU fallback unavailable")
        return
    res = subprocess.run([cc, "-shared", "-fPIC", "-O2", "-Wl,-soname,libcuda.so.1",
                          "-o", str(target), str(STUB_SRC)], capture_output=True, text=True)
    if res.returncode != 0:
        log("stub build failed: " + res.stderr[-500:])


def cuda_lib_dirs() -> list[str]:
    dirs: set[str] = set()
    for base in sys.path:
        for lib in glob.glob(os.path.join(base, "nvidia", "*", "lib")):
            dirs.add(lib)
    return sorted(dirs)


def gpu_available() -> bool:
    """True when a real NVIDIA driver and device are visible to this process."""
    try:
        ctypes.CDLL("libcuda.so.1")
    except OSError:
        return False
    smi = shutil.which("nvidia-smi")
    if smi:
        try:
            out = subprocess.run([smi, "-L"], capture_output=True, text=True, timeout=20)
            return out.returncode == 0 and "GPU" in out.stdout
        except Exception:
            return False
    return True


def server_env(backend: str) -> dict[str, str]:
    env = {k: v for k, v in os.environ.items()}
    dirs = cuda_lib_dirs()
    if backend == "cpu":
        dirs.append(str(STUB_DIR))  # lets the CUDA build start without a driver
    if env.get("LD_LIBRARY_PATH"):
        dirs.append(env["LD_LIBRARY_PATH"])
    env["LD_LIBRARY_PATH"] = ":".join(dirs)
    env.setdefault("OMP_NUM_THREADS", str(os.cpu_count() or 4))
    return env


# --------------------------------------------------------------------------- catalog
_catalog_cache: dict[str, dict] | None = None


def specs_dir() -> Path:
    return ensure_runtime() / "model_specs"


def catalog() -> dict[str, dict]:
    """family -> spec summary (tasks, packages, voices, options)."""
    global _catalog_cache
    if _catalog_cache is not None:
        return _catalog_cache
    out: dict[str, dict] = {}
    for path in sorted(specs_dir().glob("*.json")):
        try:
            spec = json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            continue
        family = spec.get("family") or path.stem
        defaults = spec.get("package_defaults") or {}
        pkgs = []
        for p in spec.get("packages", []):
            p = dict(p)
            p.setdefault("download", defaults.get("download"))
            pkgs.append(p)
        ui = spec.get("ui") or {}
        voices = list(ui.get("builtin_voices") or [])
        default_voice = ui.get("default_voice")
        if family in VOICE_FALLBACKS and not voices:
            voices, default_voice = VOICE_FALLBACKS[family]
        out[family] = {
            "family": family,
            "display_name": spec.get("display_name", family),
            "description": spec.get("description", ""),
            "category": spec.get("category", ""),
            "status": spec.get("status", ""),
            "tasks": spec.get("tasks", []),
            "languages": spec.get("languages", []),
            "packages": pkgs,
            "recommended": ui.get("recommended_package"),
            "voices": voices,
            "default_voice": default_voice,
            "request_options": (spec.get("options") or {}).get("request", []),
            "docs": ui.get("docs", []),
        }
    _catalog_cache = out
    return out


def families_for_task(task: str) -> list[str]:
    return sorted(f for f, s in catalog().items() if task in s["tasks"] and s["packages"])


def package(pkg_id: str) -> tuple[dict, dict]:
    for fam in catalog().values():
        for p in fam["packages"]:
            if p["id"] == pkg_id:
                return fam, p
    raise KeyError(f"unknown package {pkg_id!r}")


def default_package(family: str) -> str | None:
    fam = catalog()[family]
    ids = [p["id"] for p in fam["packages"]]
    if fam["recommended"] in ids:
        return fam["recommended"]
    for p in fam["packages"]:
        if p.get("default"):
            return p["id"]
    gguf = [p["id"] for p in fam["packages"] if p.get("format") == "gguf"]
    return (gguf or ids or [None])[0]


def list_packages(task: str | None = None) -> list[dict]:
    rows = []
    for fam in catalog().values():
        if task and task not in fam["tasks"]:
            continue
        for p in fam["packages"]:
            rows.append({
                "family": fam["family"], "package": p["id"], "format": p.get("format"),
                "precision": p.get("precision"), "tasks": fam["tasks"],
                "installed": is_installed(p["id"]),
            })
    return rows


def package_dir(pkg: dict) -> Path:
    return MODELS_ROOT / pkg["target_directory"]


def model_path(pkg_id: str) -> Path:
    """Path to hand to the loader, mirroring the audio.cpp WebUI's choice."""
    fam, pkg = package(pkg_id)
    root = package_dir(pkg)
    if pkg.get("format") != "gguf" or fam["family"] in DIR_FAMILIES:
        return root
    files = pkg.get("files") or []
    if fam["family"] == "minimax_h3":
        want = "dit_int8.gguf" if "int8_dit" in pkg_id else "dit.gguf"
        files = [f for f in files if f.lower().endswith("/" + want)]
    ggufs = [f for f in files if f.lower().endswith(".gguf")]
    if not ggufs:
        return root
    rel = ggufs[0].replace("\\", "/")
    prefix = (pkg.get("strip_prefix") or "").rstrip("/")
    if prefix and rel.startswith(prefix + "/"):
        rel = rel[len(prefix) + 1:]
    return root / rel


def session_options_for(pkg_id: str) -> dict[str, str]:
    fam, pkg = package(pkg_id)
    if fam["family"] == "lfm2_audio":
        backbone = next((f for f in pkg.get("files", []) if not Path(f).name.startswith("mmproj-")), None)
        return {"lfm2_audio.model_gguf": Path(backbone).name} if backbone else {}
    if fam["family"] == "minimax_music3":
        table = {
            "minimax_music3_q8_0": ("language_model_q8_0.gguf", "rvq_depth_decoder_q8_0.gguf", "transformer_q8_0.gguf"),
            "minimax_music3_bf16": ("language_model_bf16.gguf", "rvq_depth_decoder_bf16.gguf", "transformer_bf16.gguf"),
            "minimax_music3_q4_0": ("language_model_q4_0.gguf", "rvq_depth_decoder_q8_0.gguf", "transformer_q4_0.gguf"),
        }
        if pkg_id in table:
            lm, rvq, flow = table[pkg_id]
            return {"minimax_music3.language_model_gguf": lm,
                    "minimax_music3.rvq_depth_decoder_gguf": rvq,
                    "minimax_music3.flow_transformer_gguf": flow}
    return {}


def is_installed(pkg_id: str) -> bool:
    try:
        _, pkg = package(pkg_id)
    except KeyError:
        return False
    root = package_dir(pkg)
    if not root.exists():
        return False
    files = pkg.get("files") or []
    prefix = (pkg.get("strip_prefix") or "").rstrip("/")
    for f in files:
        rel = f.replace("\\", "/")
        if prefix and rel.startswith(prefix + "/"):
            rel = rel[len(prefix) + 1:]
        if not (root / rel).exists():
            return False
    return bool(files) or any(root.iterdir())


def package_size(pkg_id: str) -> int | None:
    """Download size in bytes from the Hugging Face API, if available."""
    try:
        from huggingface_hub import HfApi
        _, pkg = package(pkg_id)
        dl = pkg.get("download") or {}
        if dl.get("kind") != "huggingface_snapshot" or not pkg.get("files"):
            return None
        infos = HfApi().get_paths_info(dl["repo"], pkg["files"], revision=dl.get("revision") or "main")
        return sum(getattr(i, "size", 0) or 0 for i in infos)
    except Exception:
        return None


def install(pkg_id: str, progress=None) -> Path:
    """Download one package with the repo's model manager. Returns the model path."""
    ensure_runtime()
    if is_installed(pkg_id):
        return model_path(pkg_id)
    with _install_lock:
        if is_installed(pkg_id):
            return model_path(pkg_id)
        bundle = ensure_runtime()
        cmd = [sys.executable, str(bundle / "tools" / "model_manager_v2.py"),
               "--specs-dir", str(bundle / "model_specs"), "install", pkg_id,
               "--models-root", str(MODELS_ROOT), "--progress"]
        log("installing " + pkg_id)
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        tail: list[str] = []
        for line in proc.stdout:  # type: ignore[union-attr]
            line = line.rstrip()
            tail = (tail + [line])[-30:]
            if progress and line.startswith("AUDIOCPP_PROGRESS"):
                try:
                    kv = dict(part.split("=", 1) for part in line.split()[1:] if "=" in part)
                    done, total = int(kv.get("downloaded", 0)), int(kv.get("total", 0))
                    if total > 0:
                        progress(min(done / total, 1.0),
                                 desc=f"Downloading {pkg_id} ({done / 1e9:.2f} / {total / 1e9:.2f} GB)")
                except Exception:
                    pass
        if proc.wait() != 0:
            raise RuntimeError(f"install of {pkg_id} failed:\n" + "\n".join(tail))
    return model_path(pkg_id)


def uninstall(pkg_id: str) -> None:
    bundle = ensure_runtime()
    subprocess.run([sys.executable, str(bundle / "tools" / "model_manager_v2.py"),
                    "--specs-dir", str(bundle / "model_specs"), "uninstall", pkg_id,
                    "--models-root", str(MODELS_ROOT)], check=False)


def installed_packages() -> list[str]:
    return [r["package"] for r in list_packages() if r["installed"]]


# --------------------------------------------------------------------------- audio
def to_wav(src: str | os.PathLike, dst_dir: str | os.PathLike, rate: int | None = None,
           mono: bool = False, max_seconds: float | None = None) -> str:
    """Convert any audio or video file to 16-bit PCM WAV with ffmpeg."""
    dst = Path(dst_dir) / f"{Path(src).stem[:40]}-{uuid.uuid4().hex[:6]}.wav"
    cmd = ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", str(src), "-vn"]
    if max_seconds:
        cmd += ["-t", str(max_seconds)]
    if rate:
        cmd += ["-ar", str(rate)]
    if mono:
        cmd += ["-ac", "1"]
    cmd += ["-c:a", "pcm_s16le", str(dst)]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode != 0:
        raise RuntimeError(f"ffmpeg could not read {Path(src).name}: {res.stderr.strip()[-400:]}")
    return str(dst)


def input_format(task: str, family: str) -> tuple[int | None, bool]:
    """Sample rate / mono policy matching the audio.cpp WebUI."""
    if task == "sep" or family == "apollo":
        return 44100, False
    if task in {"asr", "vad", "diar", "align", "midi", "turn"}:
        return 16000, True
    return None, False


# --------------------------------------------------------------------------- server
def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def _post(url: str, body: dict, timeout: float) -> dict:
    data = json.dumps(body).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8", "replace"))
    except urllib.error.HTTPError as err:
        payload = err.read().decode("utf-8", "replace")
        try:
            msg = json.loads(payload).get("error", {}).get("message", payload)
        except Exception:
            msg = payload
        raise RuntimeError(f"HTTP {err.code}: {msg}") from None


def _tail(path: Path, n: int = 60) -> str:
    try:
        return "\n".join(path.read_text(errors="replace").splitlines()[-n:])
    except Exception:
        return ""


def _save_result(res: dict, out_dir: Path, stem: str) -> dict:
    """Replace base64 payloads with files on disk."""
    out: dict[str, Any] = {k: v for k, v in res.items()
                           if k not in {"audio", "named_audio_outputs", "artifacts"}}
    files: list[str] = []
    if res.get("audio"):
        p = out_dir / f"{stem}.wav"
        p.write_bytes(base64.b64decode(res["audio"]))
        files.append(str(p))
        out["audio_file"] = str(p)
    for item in res.get("named_audio_outputs") or []:
        name = "".join(c if c.isalnum() or c in "-_" else "_" for c in str(item.get("id", "out")))
        p = out_dir / f"{stem}-{name}.wav"
        p.write_bytes(base64.b64decode(item["audio"]))
        files.append(str(p))
        out.setdefault("named_audio_files", {})[item.get("id", name)] = str(p)
    for art in res.get("artifacts") or []:
        meta = art.get("meta") or {}
        ext = meta.get("extension") or meta.get("ext") or {"midi": "mid", "video": "mp4", "json": "json",
                                                             "text": "txt"}.get(art.get("kind", ""), "bin")
        name = "".join(c if c.isalnum() or c in "-_" else "_" for c in str(art.get("id", "artifact")))
        p = out_dir / f"{stem}-{name}.{ext.lstrip('.')}"
        p.write_bytes(base64.b64decode(art.get("payload", "")))
        files.append(str(p))
        out.setdefault("artifact_files", []).append({"id": art.get("id"), "kind": art.get("kind"),
                                                     "meta": meta, "file": str(p)})
    out["files"] = files
    return out


def execute(job: dict, prefer_gpu: bool = True) -> dict:
    """Start audiocpp_server for one model, run every request in job, stop it.

    job keys: family, task (spec name), model_path, requests (list of request
    dicts), load_options, session_options, timeout (seconds), out_dir.
    """
    ensure_runtime()
    started = time.time()
    # Preset-voice models (e.g. Supertonic) fail without a voice id, so fill the default.
    fam = catalog().get(job["family"], {})
    if job["task"] in {"tts", "design"} and fam.get("default_voice"):
        for r in job["requests"]:
            if not r.get("voice_id") and not r.get("voice_ref"):
                r["voice_id"] = fam["default_voice"]
    timeout = float(job.get("timeout") or 600)
    backends = ["cuda", "cpu"] if (prefer_gpu and gpu_available()) else ["cpu"]
    out_dir = Path(job.get("out_dir") or WORK_ROOT / uuid.uuid4().hex)
    out_dir.mkdir(parents=True, exist_ok=True)
    errors: list[str] = []
    for backend in backends:
        port = _free_port()
        model = {
            "id": "m",
            "family": job["family"],
            "path": str(job["model_path"]),
            "task": TASK_TOKEN.get(job["task"], job["task"]),
            "mode": "offline",
        }
        if job.get("load_options"):
            model["load_options"] = {k: str(v) for k, v in job["load_options"].items()}
        if job.get("session_options"):
            model["session_options"] = {k: str(v) for k, v in job["session_options"].items()}
        cfg = {"host": "127.0.0.1", "port": port, "backend": backend, "device": 0,
               "threads": os.cpu_count() or 4, "lazy_load": True,
               "busy_timeout_ms": int(timeout * 1000), "models": [model]}
        cfg_path = out_dir / f"server-{backend}.json"
        cfg_path.write_text(json.dumps(cfg, indent=1))
        log_path = out_dir / f"server-{backend}.log"
        server = _state["server"]
        with open(log_path, "w") as logf:
            proc = subprocess.Popen([str(server), "--config", str(cfg_path), "--no-ui"],
                                    cwd=str(server.parent), env=server_env(backend),
                                    stdout=logf, stderr=subprocess.STDOUT)
        try:
            base = f"http://127.0.0.1:{port}"
            deadline = time.time() + 60
            while True:
                if proc.poll() is not None:
                    raise RuntimeError("server exited during startup")
                try:
                    with urllib.request.urlopen(base + "/health", timeout=2) as r:
                        if r.status == 200:
                            break
                except Exception:
                    pass
                if time.time() > deadline:
                    raise RuntimeError("server did not become ready")
                time.sleep(0.3)
            results = []
            for i, request in enumerate(job["requests"]):
                remaining = max(10.0, timeout - (time.time() - started))
                t0 = time.time()
                res = _post(base + "/v1/tasks/run", {"model": "m", "request": request}, remaining)
                saved = _save_result(res, out_dir, request.get("id") or f"out{i + 1:02d}")
                saved["wall_s"] = round(time.time() - t0, 2)
                results.append(saved)
            return {"backend": backend, "results": results, "log": _tail(log_path),
                    "elapsed_s": round(time.time() - started, 2), "out_dir": str(out_dir)}
        except Exception as err:
            errors.append(f"[{backend}] {err}\n{_tail(log_path, 25)}")
            # Request-level errors (HTTP 4xx from the model) would repeat on CPU,
            # so only start-up / transport failures fall through to the CPU backend.
            if str(err).startswith("HTTP 4"):
                break
        finally:
            proc.terminate()
            try:
                proc.wait(timeout=10)
            except subprocess.TimeoutExpired:
                proc.kill()
    raise RuntimeError("\n\n".join(errors))


def run(family: str, task: str, request: dict | list[dict] | None = None, package_id: str | None = None,
        load_options: dict | None = None, session_options: dict | None = None,
        timeout: float = 600, prefer_gpu: bool = True, out_dir: str | None = None) -> dict:
    """Convenience wrapper: install the package if needed and run request(s)."""
    package_id = package_id or default_package(family)
    path = install(package_id)
    sess = {**session_options_for(package_id), **(session_options or {})}
    reqs = request if isinstance(request, list) else [request or {}]
    return execute({"family": family, "task": task, "model_path": str(path), "requests": reqs,
                    "load_options": load_options or {}, "session_options": sess,
                    "timeout": timeout, "out_dir": out_dir}, prefer_gpu=prefer_gpu)