"""sglang-omni model weights: install state on disk, (un)install actions.

Model weights are not part of the pip install: each server fetches its
HuggingFace repo into the standard hub cache on first start. This module
pre-fetches ("Install") and deletes ("Uninstall") those cache directories
per model — via the venv's hf CLI, exactly what a first server start
would do — plus the per-model companion packages the server needs
(``extras`` in the catalog, e.g. the qwen-tts --no-deps stack or the
Descript DAC codec).

The cache helpers mirror huggingface_hub's own directory layout and
environment overrides (the same files ``from_pretrained`` writes), so an
install lands exactly where a server start would look. The cache is
shared with the other backends: a repo both host (Qwen3-TTS 1.7B Base /
VoiceDesign exist in the qwen backend too) is downloaded once and its
deletion affects both — the same convention every backend here accepts.
"""

import shutil
from pathlib import Path
from typing import List, Optional

from backends import common, envs
from backends.sglomni.catalog import ENTRIES, Extra, ModelEntry, \
    entry_by_key, entry_by_repo, extra_import_name
from backends.sglomni.constants import SERVER_NAME, SGLOMNI_PIP_PKG
from backends.sglomni.pythonenv import SGLOMNI_ENV, prepare_env

# The HF-cache primitives live in backends.common (shared with the qwen
# backend, which fetches into the same cache); the module-level wrappers
# below keep the names internal callers (and the tests) reference.


def _hf_cache_dir() -> Path:
    """The cache directory HF keeps repos in (models--<org>--<name>
    folders); common.hf_cache_dir resolves the environment overrides."""
    return common.hf_cache_dir()


def repo_dir(repo_id: str) -> Path:
    """The cache directory HF keeps REPO_ID's weights in."""
    return _hf_cache_dir() / ("models--" + repo_id.replace("/", "--"))


def model_repo_dir(entry: ModelEntry) -> Path:
    """The cached-weights directory for a catalog entry."""
    return repo_dir(entry.repo)


def _tree_has_file(path: Path) -> bool:
    """True when any file or symlink exists under PATH (recursively)."""
    return common.hf_tree_has_file(path)


def model_installed(entry: ModelEntry) -> bool:
    """True when ENTRY's weights look complete in the local HF cache.

    A fetched repo has refs/main plus at least one file under snapshots/;
    anything less counts as not installed. An interrupted download simply
    resumes — via Install, or the next server start for that model.
    """
    directory = model_repo_dir(entry)
    if not (directory / "refs" / "main").is_file():
        return False
    return _tree_has_file(directory / "snapshots")


def installed_entries() -> List[ModelEntry]:
    """The catalog entries whose weights are already on disk."""
    return [entry for entry in ENTRIES if model_installed(entry)]


def installed_keys() -> List[str]:
    """The installed entries' catalog keys, in catalog order."""
    return [entry.key for entry in installed_entries()]


def preset_voices(entry: ModelEntry) -> List[str]:
    """The preset voice names ENTRY can speak with.

    Catalog-declared speakers first (the Qwen3-TTS CustomVoice table);
    otherwise the checkpoint's own ``voice_embedding/*.pt`` presets are
    read from the downloaded snapshot (how Voxtral ships its named
    voices). Empty when the model declares none or is not downloaded.
    """
    if entry.speakers:
        return list(entry.speakers)
    try:
        snapshots = model_repo_dir(entry) / "snapshots"
        for snapshot in sorted(snapshots.iterdir()):
            voice_dir = snapshot / "voice_embedding"
            if voice_dir.is_dir():
                names = sorted(item.stem for item in voice_dir.glob("*.pt")
                               if item.is_file())
                if names:
                    return names
    except OSError:
        pass
    return []


def system_dep_missing(entry: ModelEntry) -> Optional[str]:
    """Remediation text when ENTRY's system binary is absent (None = ok)."""
    if entry.system_dep and not shutil.which(entry.system_dep):
        return (f"{entry.system_dep} (system package) was not found — "
                f"{entry.system_hint}. The weights still download, but the "
                "server will fail to synthesize with this model until it "
                "is installed.")
    return None


def missing_companions(entry: ModelEntry) -> List[Extra]:
    """ENTRY's companion packages absent from the sglang-omni venv.

    One import probe per extra (its top-level module, see
    ``extra_import_name``) against the venv's interpreter, so a model whose
    weights landed in the shared HF cache by another route — or whose
    companions failed to pip-install during setup, which install_model only
    warns about — is detected before the server dies on the import. A venv
    that does not exist at all yields no verdict: the start flow fails on
    the missing ``sgl-omni`` executable anyway."""
    if not entry.extras or not envs.env_exists(SGLOMNI_ENV):
        return []
    return [extra for extra in entry.extras
            if not envs.module_available(extra_import_name(extra[0]),
                                         SGLOMNI_ENV)]


def install_companions(entry: ModelEntry, *, emit=None, cancel=None,
                       force: bool = False) -> int:
    """pip-install ENTRY's companion packages into the venv.

    The same recipe ``install_model`` runs (the catalog's ``--no-deps``
    flags preserved — the Qwen3-TTS companions must not replace the pinned
    Transformers 5 stack), limited to what the import probe found absent,
    so a start-time heal touches as little of the pinned environment as
    possible. With FORCE every extra's pip spec re-runs instead — a
    satisfied pin is a pip no-op, so the update flow uses that to heal
    version drift the import probe cannot see (the protobuf re-pin
    especially). Returns the first failing exit code, 0 when all present."""
    wanted = list(entry.extras) if force else missing_companions(entry)
    for spec, no_deps in wanted:
        args = ["--no-deps"] if no_deps else None
        rc = common.pip_install([spec], emit=emit, cancel=cancel,
                                env_dir=SGLOMNI_ENV, extra_args=args)
        if rc != 0:
            print(f"[ERROR] pip install {spec} failed (exit {rc}); "
                  f"install it into {SGLOMNI_ENV} manually")
            return rc
    return 0


def install_model(key: str, *, emit=None, cancel=None) -> int:
    """Install a catalog model: companion packages, then its weights.

    Companion ``extras`` pip-install into the sglang-omni venv exactly as
    upstream instructs (``--no-deps`` where upstream says so — the
    Qwen3-TTS companions must not replace the pinned Transformers 5
    stack), and a missing system binary is a loud warning, not a stop:
    the download is still useful and the remediation stays on screen.
    The weights pre-download via the venv's hf CLI (resumable, streamed,
    cancelable). Returns the exit code.
    """
    entry = entry_by_key(key)
    if entry is None:
        print(f"[ERROR] Unknown sglang-omni model: {key!r}")
        return 1
    rc = prepare_env(emit=emit, cancel=cancel)
    if rc != 0:
        return rc
    warning = system_dep_missing(entry)
    if warning:
        print(f"[WARNING] {warning}")
    # A GPU the model's default FP8 pipeline cannot run gets the bf16
    # fallback note up front (the install itself is still useful: the
    # weights download either way).
    from backends.sglomni import status as sg_status
    note = sg_status.gpu_fallback_note(entry)
    if note:
        print(f"[WARNING] {note}")
    for spec, no_deps in missing_companions(entry):
        args = ["--no-deps"] if no_deps else None
        rc = common.pip_install([spec], emit=emit, cancel=cancel,
                                env_dir=SGLOMNI_ENV, extra_args=args)
        if rc != 0:
            print(f"[WARNING] pip install {spec} failed (exit {rc}); "
                  f"install it into {SGLOMNI_ENV} manually")
    prefix = _hf_download_prefix()
    if prefix is None:
        print("[ERROR] No hf CLI found in the sglang-omni venv; pip "
              f"install {SGLOMNI_PIP_PKG} first")
        return 1
    print(f"[INFO] Downloading {entry.repo} into {_hf_cache_dir()}...")
    rc = common.run_console_subprocess(
        prefix + ["download", entry.repo], emit=emit, cancel=cancel)
    if rc == 0:
        print(f"[OK] {entry.label} downloaded.")
    return rc


def uninstall_model(key: str, *, emit=None, cancel=None) -> int:
    """Remove a model's cached weights (the inverse of install_model).

    A locally-managed server currently hosting the model is stopped first
    (best-effort) so its weights are not deleted under a live process.
    CANCEL is honored after that stop phase only. Returns the exit code.
    """
    entry = entry_by_key(key)
    if entry is None:
        print(f"[ERROR] Unknown sglang-omni model: {key!r}")
        return 1
    if _managed_running_repo() == entry.repo:
        from backends import servers
        servers.stop(SERVER_NAME)
    if common.cancel_requested(cancel):
        return 130
    delete_model_weights([entry])
    return 0


def delete_model_weights(entries: Optional[List[ModelEntry]] = None) -> int:
    """Delete the cached HF weight dirs of ENTRIES (every model by default).

    Delegates to common.hf_delete_model_weights: best-effort rmtree of
    each ``models--<org>--<name>`` directory; only those directories are
    ever touched — the rest of the HF cache may be shared with unrelated
    tools. Returns how many were present and removed.
    """
    if entries is None:
        entries = list(ENTRIES)
    return common.hf_delete_model_weights([entry.repo for entry in entries])


def _hf_download_prefix() -> Optional[List[str]]:
    """The sglang-omni venv's hf CLI argv prefix (None when absent)."""
    return common.hf_download_prefix(SGLOMNI_ENV)


def _managed_running_repo() -> Optional[str]:
    """The repo id a locally-managed, up-and-running server hosts."""
    from backends import probe, servers
    from converter import config
    if servers.pid_for(SERVER_NAME) is None:
        return None
    if not servers.alive(SERVER_NAME):
        return None
    return probe.sglomni_served_model(config.SGLOMNI_API_URL)


def resolve_model(key: Optional[str]) -> ModelEntry:
    """The catalog entry a run with MODEL_KEY uses.

    An explicit KEY must exist in the catalog and be installed (a hosted
    model without weights cannot boot). Without KEY the single installed
    model is auto-selected; several installed models need an explicit pick
    (the CLI --model flag or the Generate form's Model menu). Raises
    RuntimeError with an actionable message otherwise.
    """
    if key is not None:
        entry = entry_by_key(key)
        if entry is None:
            known = ", ".join(e.key for e in ENTRIES)
            raise RuntimeError(
                f"Unknown sglang-omni model {key!r} (installed models are "
                f"picked by catalog key; known keys: {known})")
        if not model_installed(entry):
            raise RuntimeError(
                f"{entry.label} is not downloaded — install it via "
                "Configure Backends → SGLang-Omni, or pick an installed "
                "model.")
        return entry
    installed = installed_entries()
    if not installed:
        raise RuntimeError(
            "No sglang-omni models are downloaded — install one via "
            "Configure Backends → SGLang-Omni (Configure).")
    if len(installed) > 1:
        names = ", ".join(e.key for e in installed)
        raise RuntimeError(
            "Several sglang-omni models are installed; pick one with "
            f"--model KEY (installed: {names})")
    return installed[0]


def entry_for_served_repo(repo: Optional[str]) -> Optional[ModelEntry]:
    """The catalog entry a served /v1/models repo id belongs to."""
    return entry_by_repo(repo) if repo else None