path: root/app/backends/audiocpp/models.py

"""Model install state: what is on disk, what is missing, how to fetch it."""

import json
import os
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

from backends import common
from . import catalog as _catalog

# No-output watchdog for model downloads: huggingface_hub streams steady
# byte progress, so 5 minutes of silence means the transfer wedged. The
# runner kills it and reports exit 124 (see run_console_subprocess).
DOWNLOAD_STALL_TIMEOUT = 300

def _installed_display_names(audiocpp_dir: Path,
                             model_entries: Optional[List[dict]],
                             install_guidance: List[Tuple[str, str]]
                             ) -> Set[str]:
    """Display names from INSTALL_GUIDANCE whose model files are on disk.

    MODEL_ENTRIES and INSTALL_GUIDANCE are built in lockstep by
    ``_build_entries`` (one guidance pair per entry), so the pairs resolve
    positionally: each entry's ``path`` is checked against the checkout
    exactly like ``_all_models_present`` resolves it. Returns an empty set
    when ENTRIES is None or does not line up with the guidance (no
    filtering — every model counts as not installed).
    """
    if model_entries is None or len(model_entries) != len(install_guidance):
        return set()
    installed: Set[str] = set()
    for entry, (name, _install_id) in zip(model_entries, install_guidance):
        rel = entry.get("path")
        if not isinstance(rel, str) or not rel:
            continue
        path = Path(rel) if Path(rel).is_absolute() else audiocpp_dir / rel
        if _model_path_present(path):
            installed.add(name)
    return installed


def _split_pending_and_installed(
        install_guidance: List[Tuple[str, str]],
        installed_names: Set[str]) -> Tuple[List[Tuple[str, str]], List[str]]:
    """Partition guidance into (pending installs, installed display names).

    PENDING keeps only models whose display name is not INSTALLED_NAMES,
    de-duped by install id (the same package may host several entries) in
    first-occurrence order. INSTALLED lists each installed display name
    once, also in first-occurrence order.
    """
    seen: Set[str] = set()
    pending: List[Tuple[str, str]] = []
    noted: List[str] = []
    for name, install_id in install_guidance:
        if name in installed_names:
            if name not in noted:
                noted.append(name)
            continue
        if install_id in seen:
            continue
        seen.add(install_id)
        pending.append((name, install_id))
    return pending, noted


def _install_models(audiocpp_dir: Path,
                    install_guidance: List[Tuple[str, str]],
                    download: bool, emit=None, cancel=None,
                    model_entries: Optional[List[dict]] = None) -> int:
    """Report and optionally run the model install commands.

    When MODEL_ENTRIES (built in lockstep with INSTALL_GUIDANCE by
    ``_build_entries``) is given, models already on disk are reported as
    installed and never re-downloaded or printed as commands; when every
    selected model is present nothing runs at all. The remaining models
    get one ``python <manager> install <id>`` command each (de-duped by
    install id). When DOWNLOAD is True each command is run in the audio.cpp
    checkout via ``subprocess`` so the models are downloaded automatically;
    a failing install is reported as a warning and does not abort the
    remaining downloads. When DOWNLOAD is False (or the model manager is
    missing) the commands are only printed after a note that setup downloads
    them automatically — copy-pasteable for a manual install.

    With EMIT given (the in-TUI task view) each download streams its output
    to EMIT and — when the checkout's ``model_manager_v2.py`` supports it —
    runs with ``--progress --cancel-file`` so the view can show a real byte
    progress bar and cancel gracefully. CANCEL aborts a running download.

    Returns 0 when every command succeeded (or nothing needed running),
    130 when cancelled, 1 when any download failed.
    """
    if not install_guidance:
        return 0
    manager = audiocpp_dir / "tools" / "model_manager_v2.py"
    installed_names = _installed_display_names(
        audiocpp_dir, model_entries, install_guidance)
    pending, installed_noted = _split_pending_and_installed(
        install_guidance, installed_names)

    supports_progress = emit is not None and _manager_supports_progress(manager)

    if download and not manager.is_file():
        print(f"[WARNING] {manager} not found; printing the install commands "
              "instead of running them")
        download = False

    for name in installed_noted:
        print(f"[OK] {name} is already installed.")
    if not pending:
        print("[OK] All selected models are already installed.")
        return 0

    if not download:
        print("[INFO] Models are downloaded automatically by this tool's "
              "setup — to download them manually instead, run:")
        for _, install_id in pending:
            print(f"python {manager} install {install_id}")
        return 0

    failed = False
    for _, install_id in pending:
        print(f"[INFO] Downloading {install_id}...")
        argv = [sys.executable, str(manager), "install", install_id]
        cancel_file: Optional[Path] = None
        on_cancel = None
        if supports_progress:
            fd, cancel_path = tempfile.mkstemp(
                prefix="audiocpp_cancel_", suffix=".cancel")
            os.close(fd)
            cancel_file = Path(cancel_path)
            cancel_file.unlink()  # absent = not cancelled
            argv += ["--progress", "--cancel-file", str(cancel_file)]
            on_cancel = cancel_file.touch
        try:
            rc = common.run_console_subprocess(
                argv, cwd=str(audiocpp_dir), emit=emit, cancel=cancel,
                on_cancel=on_cancel,
                stall_timeout=(DOWNLOAD_STALL_TIMEOUT
                               if supports_progress else None))
        except OSError as exc:
            print(f"[WARNING] Could not run python {manager} install "
                  f"{install_id}: {exc}")
            rc = 1
        finally:
            if cancel_file is not None:
                try:
                    cancel_file.unlink()
                except OSError:
                    pass
        if rc == 130 or (cancel is not None and cancel.is_set()):
            return 130
        if rc != 0:
            failed = True
            print(f"[WARNING] install {install_id} exited with code "
                  f"{rc}; the model may need to be downloaded "
                  "by hand")
    return 1 if failed else 0


def _manager_supports_progress(manager: Path) -> bool:
    """True when MANAGER (model_manager_v2.py) supports --progress output.

    The ``--progress``/``--cancel-file`` flags are relatively recent; an
    older audio.cpp checkout may not have them, so probe the script source
    once instead of failing the download with an unknown flag.
    """
    try:
        text = manager.read_text(encoding="utf-8", errors="ignore")
    except OSError:
        return False
    return "AUDIOCPP_PROGRESS" in text and "--cancel-file" in text


def download_applicable(audiocpp_dir: Path, model_entries: List[dict]) -> bool:
    """True when the wizard's "download models automatically?" row applies.

    The audio.cpp model manager must be present (otherwise the install
    commands can only be printed), and at least one selected model must be
    missing from disk (see ``_all_models_present``), so an already-configured
    checkout is not asked to re-download models it already has.
    """
    manager = audiocpp_dir / "tools" / "model_manager_v2.py"
    if not manager.is_file():
        return False
    return not _all_models_present(audiocpp_dir, model_entries)


def _build_tree_families(catalog: List[dict]) -> List[dict]:
    """Shape the catalog into the checkbox_tree widget's family list."""
    families: List[dict] = []
    for entry in catalog:
        capabilities = ["tts"]
        if "clone" in entry["tasks"]:
            capabilities.append("cloning")
        if "design" in entry["tasks"]:
            capabilities.append("design")
        name = entry["display_name"]
        options = []
        for opt in _catalog.package_dir_options(entry):
            options.append({
                "key": opt["target_directory"],
                "label": opt["install_id"],
                "recommended": opt["recommended"],
            })
        families.append({
            "label": name,
            "detail": ", ".join(capabilities),
            "options": options,
        })
    return families


def _model_path_present(path: Path) -> bool:
    """True when a server.json model path holds actual model files.

    A present path is either a file (a single-model package) or a non-empty
    directory (the usual GGUF package target directory; an empty one means a
    download that never ran or was cleaned up halfway).
    """
    try:
        if path.is_file():
            return True
        if path.is_dir():
            return any(path.iterdir())
    except OSError:
        return False
    return False


def _all_models_present(audiocpp_dir: Path, model_entries: List[dict]) -> bool:
    """True when every selected model entry's path already holds files on disk.

    Paths resolve against the checkout (where model_manager_v2.py installs
    them), honoring absolute paths. Used by the wizard to skip the
    "Automatically download the selected models" prompt when nothing is
    actually missing. An empty selection is treated as not-present.
    """
    if not model_entries:
        return False
    for entry in model_entries:
        rel = entry.get("path")
        if not isinstance(rel, str) or not rel:
            return False
        path = Path(rel) if Path(rel).is_absolute() else audiocpp_dir / rel
        if not _model_path_present(path):
            return False
    return True


def missing_model_entries(server_json: Path) -> List[dict]:
    """Return the server.json model entries whose files are not on disk.

    Paths resolve exactly like audiocpp_server resolves them (relative paths
    against the server.json's directory). Each returned entry carries the
    entry ``id`` and ``rel`` (the configured path string); used by ``detect``
    to warn that a conversion would fail until the models are installed.
    """
    try:
        data = json.loads(server_json.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return []
    if not isinstance(data, dict):
        return []
    base = server_json.parent
    missing: List[dict] = []
    for entry in data.get("models") or []:
        if not isinstance(entry, dict):
            continue
        rel = entry.get("path")
        if not isinstance(rel, str) or not rel:
            continue
        path = Path(rel) if Path(rel).is_absolute() else base / rel
        if _model_path_present(path):
            continue
        missing.append({"id": str(entry.get("id") or rel), "rel": rel})
    return missing


def _install_id_by_path(audiocpp_dir: Path) -> Dict[str, str]:
    """Map ``models/<target_directory>`` -> catalog install id.

    The catalog package that installs a model is derived from the
    ``default_path`` of each TTS family; an entry whose path matches no
    catalog package has no install id.
    """
    by_path: Dict[str, str] = {}
    try:
        for entry in _catalog.load_model_catalog(audiocpp_dir):
            by_path[entry["default_path"]] = entry["install_id"]
    except (NotADirectoryError, OSError):
        pass
    return by_path


def installed_model_entries(server_json: Path) -> List[dict]:
    """Return the server.json model entries whose files ARE on disk.

    The complement of ``missing_model_entries``: each returned entry carries
    the entry ``id`` and ``rel`` (the configured path string), resolved
    exactly like ``missing_model_entries`` (relative against the server.json's
    directory). Used by the wizard's "Delete unused models?" step to find
    already-downloaded models that were unselected.
    """
    try:
        data = json.loads(server_json.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return []
    if not isinstance(data, dict):
        return []
    base = server_json.parent
    installed: List[dict] = []
    for entry in data.get("models") or []:
        if not isinstance(entry, dict):
            continue
        rel = entry.get("path")
        if not isinstance(rel, str) or not rel:
            continue
        path = Path(rel) if Path(rel).is_absolute() else base / rel
        if _model_path_present(path):
            installed.append({"id": str(entry.get("id") or rel), "rel": rel})
    return installed


def missing_model_install_guidance(audiocpp_dir: Path,
                                   missing: List[dict]) -> List[Tuple[str, str]]:
    """Map MISSING model entries to (display name, install id) pairs.

    The install id is derived from each entry's configured path via the
    catalog (see ``_install_id_by_path``); entries whose path matches no
    catalog package are skipped (there is no ``model_manager_v2.py install``
    command for them). Feeds ``_install_models`` for the "Download Missing
    Models" action.
    """
    by_path = _install_id_by_path(audiocpp_dir)
    guidance: List[Tuple[str, str]] = []
    for item in missing:
        install_id = by_path.get(item["rel"])
        if install_id:
            guidance.append((item["id"], install_id))
    return guidance


def model_install_hints(audiocpp_dir: Path,
                        missing: List[dict]) -> List[str]:
    """Remediation lines for MISSING model entries (see missing_model_entries).

    Maps each entry's configured path back to the catalog package that
    installs it (``models/<target_directory>`` -> install id) so the line
    carries the exact ``model_manager_v2.py install`` command; entries whose
    directory matches no catalog package just name the path.
    """
    by_path = _install_id_by_path(audiocpp_dir)
    hints: List[str] = []
    for item in missing:
        install_id = by_path.get(item["rel"])
        hint = f"model not downloaded: {item['id']} ({item['rel']})"
        if install_id:
            hint += (f" — install with: python tools/model_manager_v2.py "
                     f"install {install_id}")
        hints.append(hint)
    return hints


def install_models(audiocpp_dir: Path,
                   guidance: List[Tuple[str, str]],
                   emit=None, cancel=None) -> int:
    """Download the (display name, install id) models via the helper script.

    Runs ``model_manager_v2.py install`` for each de-duped install id in the
    checkout, streaming to the console (or to EMIT, the in-TUI task view); a
    failing install is reported as a warning and does not abort the rest.
    Returns 0 when every download succeeded, 130 when cancelled, 1 when any
    failed. Used by the hub's "Download Missing Models" action (see
    ``missing_model_install_guidance`` for the mapping).
    """
    return _install_models(audiocpp_dir, guidance, download=True,
                           emit=emit, cancel=cancel)


def hand_install_guidance(audiocpp_dir: Path,
                          missing: List[dict]) -> str:
    """Explain how to install MISSING model entries by hand.

    Returns a multi-line message listing each missing model's id and the
    path its files must be placed in (``rel``, resolved against the
    checkout). Used when the missing models cannot be mapped to
    a ``model_manager_v2.py install`` command, so the user still knows what
    to download and where to put it.
    """
    lines = [
        "None of the missing models map to a model_manager_v2.py install "
        "command.",
        "Download them by hand and place the files at these paths:",
    ]
    for item in missing:
        lines.append(f"  {item['id']}  ->  {item['rel']}")
    lines.append(f"(paths are relative to {audiocpp_dir})")
    return "\n".join(lines)


def unused_installed_entries(server_json: Path,
                             new_paths: Set[str]) -> List[dict]:
    """Return installed server.json entries whose path is not in NEW_PATHS.

    The already-downloaded models (see ``installed_model_entries``) that the
    new selection does not host any more — the candidates for the wizard's
    "Delete unused models?" prompt. Entries whose files are not on disk are
    never listed (there is nothing to delete).
    """
    return [entry for entry in installed_model_entries(server_json)
            if entry["rel"] not in new_paths]


def delete_model_files(server_json: Path, entries: List[dict]) -> int:
    """Remove the on-disk model files for ENTRIES ({id, rel}) from disk.

    Each entry's ``rel`` is resolved exactly like the server resolves it
    (relative against ``server_json``'s directory; absolute paths honored),
    then removed as a directory tree or a single file. Missing entries are
    ignored. Returns the number of paths removed. Used by the wizard's
    "Delete unused models?" step — the regenerated server.json already only
    lists the kept models, so no entry cleanup is needed here.
    """
    base = server_json.parent
    removed = 0
    for item in entries:
        rel = item.get("rel")
        if not isinstance(rel, str) or not rel:
            continue
        path = Path(rel) if Path(rel).is_absolute() else base / rel
        try:
            if not path.exists():
                continue
            if path.is_dir():
                shutil.rmtree(path, ignore_errors=True)
            else:
                path.unlink()
        except OSError as exc:
            print(f"[WARNING] Could not remove {path}: {exc}")
            continue
        print(f"[OK] Removed unused model {path}")
        removed += 1
    return removed