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

"""The model_specs catalog, server.json building and selection views."""

import json
import re
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

from .constants import TASK_TTS

DESIGN_PACKAGE_RE = re.compile(r"voice[\s_\-]?design", re.IGNORECASE)


def request_options_families(audiocpp_dir: Path) -> Dict[str, dict]:
    """Map the families whose spec defines per-request options.

    Reads every ``model_specs/<family>.json`` in AUDIOCPP_DIR once and
    returns ``{family_key: {"display_name": ...}}`` for the specs that
    list request options (a non-empty ``options.request`` array) — the
    families a "Request options" field makes sense for. The key is the
    spec's ``family`` field (falling back to the file stem), matching
    what ``GET /v1/models`` reports, so callers can look an entry up by
    its family id. A missing or unreadable specs directory yields {}
    (every family then counts as unknown rather than unsupported).
    """
    specs_dir = audiocpp_dir / "model_specs"
    if not specs_dir.is_dir():
        return {}
    families: Dict[str, dict] = {}
    for spec_path in sorted(specs_dir.glob("*.json")):
        try:
            spec = json.loads(spec_path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            continue
        options = spec.get("options")
        request = options.get("request") if isinstance(options, dict) else None
        if not isinstance(request, list) or not request:
            continue
        family = str(spec.get("family") or spec_path.stem)
        families[family] = {
            "display_name": str(spec.get("display_name") or family),
        }
    return families


def supports_request_options(families: Dict[str, dict],
                             family: str) -> Optional[bool]:
    """Whether FAMILY accepts per-request options — None when unknown.

    True only when FAMILIES (from request_options_families) lists the
    family; False when it was read but does not define request options;
    None when support cannot be determined from the local specs (no
    checkout, or a family the specs do not describe).
    """
    if not families:
        return None
    return family in families


_BACKEND_DESCRIPTIONS = (
    ("cuda", "NVIDIA GPUs"),
    ("vulkan", "cross-vendor GPU"),
    ("hip", "AMD GPUs"),
    ("cpu", "no GPU required"),
)


def _backend_options(detected: Optional[str] = None
                     ) -> Tuple[List[Tuple[str, str]], int]:
    """Build the aligned backend menu options and the default index.

    The backend names are padded to a common width so the ``-`` dashes
    before the descriptions line up. When DETECTED matches one of the
    options, that option gets ``[auto-detected]`` appended and is the
    default (cursor/start) selection; otherwise the first option is the
    default as before. Returns (options, default_index).
    """
    width = max(len(name) for name, _ in _BACKEND_DESCRIPTIONS)
    options: List[Tuple[str, str]] = []
    default_index = 0
    for index, (name, desc) in enumerate(_BACKEND_DESCRIPTIONS):
        label = f"{name.ljust(width)} - {desc}"
        if detected == name:
            label += " [auto-detected]"
            default_index = index
        options.append((label, name))
    return options, default_index


_BACKEND_TOKEN_RE = re.compile(r"-(cuda|vulkan|hip|cpu|metal)(?:-|$)")


def detect_backend(audiocpp_dir: Path) -> Optional[str]:
    """Best-effort detection of the backend audiocpp_server was built for.

    Scans ``audiocpp_dir/build/*`` for build directories that contain a
    built ``bin/audiocpp_server`` (``.exe`` allowed on Windows) and reads
    the backend token out of the directory name (``-cuda-``, ``-vulkan-``,
    ``-hip-`` or ``-cpu-``; ``-metal-`` is mapped to ``cpu``). Returns the
    backend only when exactly one distinct backend was built, so a checkout
    with builds for several backends does not silently pick one. Returns
    None when there is no ``build/`` directory, no built server, or more
    than one distinct backend.
    """
    build_root = audiocpp_dir / "build"
    if not build_root.is_dir():
        return None
    backends: Set[str] = set()
    try:
        build_dirs = sorted(build_root.iterdir(),
                            key=lambda p: p.name.lower())
    except OSError:
        return None
    for build_dir in build_dirs:
        if not build_dir.is_dir():
            continue
        server = build_dir / "bin" / "audiocpp_server"
        if not server.exists():
            server_exe = build_dir / "bin" / "audiocpp_server.exe"
            if not server_exe.exists():
                continue
        match = _BACKEND_TOKEN_RE.search(build_dir.name.lower())
        if not match:
            continue
        token = match.group(1)
        backends.add("cpu" if token == "metal" else token)
    if len(backends) == 1:
        return next(iter(backends))
    return None


def _default_package(packages: List[dict]) -> Optional[dict]:
    """Pick the default package from a list of packages.

    Prefers the package flagged ``default: true``, then the first GGUF
    package, then the first package overall. Returns None for an empty list.
    """
    if not packages:
        return None
    for package in packages:
        if package.get("default"):
            return package
    for package in packages:
        if package.get("format") == "gguf":
            return package
    return packages[0]


def load_model_catalog(audiocpp_dir: Path) -> List[dict]:
    """Read model_specs/*.json and return the TTS-capable families.

    Each returned entry has: family, display_name, description, languages,
    clone_capable, packages (the full list from the spec), install_id
    (recommended package id), and default_path (``models/<target_directory>``).
    All families are treated equally and listed in alphabetical order by
    display name.
    """
    specs_dir = audiocpp_dir / "model_specs"
    if not specs_dir.is_dir():
        raise NotADirectoryError(
            f"{audiocpp_dir} has no model_specs/ directory; re-run setup "
            "to refresh the audio.cpp checkout")
    entries: List[dict] = []
    for spec_path in sorted(specs_dir.glob("*.json")):
        try:
            spec = json.loads(spec_path.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            continue
        tasks = spec.get("tasks") or []
        if "tts" not in tasks and spec.get("category") != "tts":
            continue
        family = spec.get("family") or spec_path.stem
        packages = spec.get("packages") or []
        package = _default_package(packages)
        if package is None:
            # No installable package: skip (cannot be hosted from a path).
            continue
        target_directory = package.get("target_directory") or family
        languages = spec.get("languages") or []
        display_name = spec.get("display_name") or family
        description = spec.get("description") or ""
        entries.append({
            "family": family,
            "display_name": display_name,
            "description": description,
            "languages": languages,
            "tasks": list(tasks),
            "clone_capable": "clone" in tasks,
            "packages": packages,
            "install_id": package.get("id") or family,
            "default_path": f"models/{target_directory}",
        })

    # All families are treated equally: alphabetical by display name.
    entries.sort(key=lambda entry: entry["display_name"].lower())
    return entries


def is_design_package(package: dict) -> bool:
    """Return True when a package's name marks it a voice-design model.

    audio.cpp voice-design packages (whose id, display name, or target
    directory mentions "voice design") are the only packages that must be
    hosted with task "vdes"; their role is not in the schema, only in those
    strings, so it is detected from them.
    """
    text = " ".join(str(package.get(key, ""))
                    for key in ("id", "display_name", "target_directory"))
    return bool(DESIGN_PACKAGE_RE.search(text))


def package_dir_options(entry: dict) -> List[dict]:
    """Return one option per distinct target_directory of a family's packages.

    Each option is a dict with: target_directory, install_id (the recommended
    package id inside that directory), design (voice-design package flag), and
    recommended (whether it holds the family's default package). Precisions
    that share a directory (q8_0/bf16/...) collapse to a single option.
    """
    packages = entry.get("packages") or []
    default_pkg = _default_package(packages)
    default_dir = (default_pkg or {}).get("target_directory") or entry["family"]
    by_dir: Dict[str, List[dict]] = {}
    order: List[str] = []
    for package in packages:
        directory = package.get("target_directory") or entry["family"]
        if directory not in by_dir:
            by_dir[directory] = []
            order.append(directory)
        by_dir[directory].append(package)
    options: List[dict] = []
    for directory in order:
        package = _default_package(by_dir[directory])
        options.append({
            "target_directory": directory,
            "install_id": (package or {}).get("id") or directory,
            "design": is_design_package(package or {}),
            "recommended": directory == default_dir,
        })
    # Put the recommended package first for a friendlier checklist.
    options.sort(key=lambda opt: not opt["recommended"])
    return options


def build_model_entry(family: str, model_id: str, model_path: str,
                      task: str = TASK_TTS) -> dict:
    """Assemble one server.json model entry.

    ``task`` defaults to "tts"; voice design packages are hosted with
    "vdes" so the server runs its design session for speech requests
    (audiobook.py then requires --instructions with that entry).
    """
    return {
        "id": model_id,
        "family": family,
        "path": model_path,
        "task": task,
        "mode": "offline",
    }


def build_server_config(host: str, port: int, backend: str, lazy_load: bool,
                        model_entries: List[dict],
                        voice_dir: Optional[str] = None) -> dict:
    """Assemble the server.json document.

    ``voice_dir`` is a server-level cloning voice library; when set, every
    hosted clone-capable family can use its voices with ``--voice``.
    """
    config_doc = {
        "host": host,
        "port": port,
        "backend": backend,
        "lazy_load": lazy_load,
        "models": model_entries,
    }
    if voice_dir:
        config_doc["voice_dir"] = voice_dir
    return config_doc


def load_server_config(server_json: Path) -> Optional[dict]:
    """Read server.json into a dict, or None when it cannot be used.

    Returns None for a missing file, unreadable content, or a non-dict
    document. Used by the wizard's modify flow to pre-fill its screens
    from an existing config instead of prompting to overwrite it.
    """
    if not server_json.exists():
        return None
    try:
        data = json.loads(server_json.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    if not isinstance(data, dict):
        return None
    return data


def server_config_selections(server_config: dict,
                             catalog: List[dict]
                             ) -> Tuple[Dict[str, List[str]],
                                        Dict[Tuple[str, str], str]]:
    """Map an existing server.json's models back to catalog selections.

    Returns ``(selected_dirs, tasks)``: ``selected_dirs`` maps a catalog
    family to the target directories it hosts (``models/<target>`` paths
    with the ``models/`` prefix stripped, in server.json order), and
    ``tasks`` maps ``(family, target_directory)`` to the entry's task
    (``"tts"`` or ``"vdes"``) so the wizard can preserve how design
    packages were hosted. Entries whose family is not in the CATALOG are
    ignored — the wizard cannot offer them again.
    """
    families = {entry["family"] for entry in catalog}
    selected_dirs: Dict[str, List[str]] = {}
    tasks: Dict[Tuple[str, str], str] = {}
    for entry in server_config.get("models") or []:
        if not isinstance(entry, dict):
            continue
        family = entry.get("family")
        if not isinstance(family, str) or family not in families:
            continue
        path = entry.get("path")
        if not isinstance(path, str):
            continue
        target = path[len("models/"):] if path.startswith("models/") else path
        if family not in selected_dirs:
            selected_dirs[family] = []
        if target not in selected_dirs[family]:
            selected_dirs[family].append(target)
        tasks[(family, target)] = str(entry.get("task") or TASK_TTS)
    return selected_dirs, tasks