stackpilot/backend/services/compose_edit_service.py

"""Server-side merging of wizard fragments into a compose YAML string.

PyYAML round-trips lose comments; that is acceptable here because these helpers
are invoked from the editor where the returned YAML replaces the buffer and the
user reviews before saving.
"""
from __future__ import annotations

import yaml

from services import gpu_service, volume_service


class EditError(Exception):
    pass


def _load(yaml_str: str) -> dict:
    try:
        data = yaml.safe_load(yaml_str) or {}
    except yaml.YAMLError as exc:
        raise EditError(f"Invalid YAML: {exc}") from exc
    if not isinstance(data, dict):
        raise EditError("Top-level compose document must be a mapping")
    data.setdefault("services", {})
    if not isinstance(data["services"], dict):
        raise EditError("`services` must be a mapping")
    return data


def _dump(data: dict) -> str:
    # Drop empty top-level keys we may have created.
    if not data.get("services"):
        data.pop("services", None)
    for key in ("volumes", "secrets", "configs"):
        if key in data and not data[key]:
            data.pop(key, None)
    return yaml.safe_dump(data, sort_keys=False, default_flow_style=False)


def _get_service(data: dict, service: str) -> dict:
    svc = data["services"].get(service)
    if svc is None:
        svc = {}
        data["services"][service] = svc
    if not isinstance(svc, dict):
        raise EditError(f"Service '{service}' is not a mapping")
    return svc


# --------------------------------------------------------------------------- #
# Volumes
# --------------------------------------------------------------------------- #


def add_volume(yaml_str: str, service: str, spec: dict) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)

    # Top-level volume declaration (named/nfs/smb).
    definition = volume_service.build_volume_definition(spec)
    if definition:
        top = data.setdefault("volumes", {})
        if not isinstance(top, dict):
            raise EditError("`volumes` must be a mapping")
        top.update(definition)

    mount = volume_service.build_service_mount(spec)
    vols = svc.setdefault("volumes", [])
    if not isinstance(vols, list):
        raise EditError(f"Service '{service}' volumes must be a list")
    if mount not in vols:
        vols.append(mount)
    return _dump(data)


# --------------------------------------------------------------------------- #
# GPU
# --------------------------------------------------------------------------- #


def set_gpu(yaml_str: str, service: str, config: dict) -> str:
    """config = { mode: 'none'|'nvidia'|'dri', ... }"""
    data = _load(yaml_str)
    svc = _get_service(data, service)
    mode = config.get("mode", "none")

    # Detect host iGPU group GIDs so we can both inject and clean them up.
    try:
        dri_gids = gpu_service.dri_group_gids()
    except Exception:  # noqa: BLE001 - detection is best-effort
        dri_gids = set()
    gpu_service.remove_gpu(svc, dri_gids)

    if mode == "nvidia":
        gpu_service.inject_nvidia(
            svc,
            device_ids=config.get("device_ids") or None,
            count=config.get("count"),
            capabilities=config.get("capabilities") or ["gpu"],
        )
    elif mode == "dri":
        gpu_service.inject_dri(
            svc,
            vendor=config.get("vendor", "intel"),
            add_render_group=config.get("add_render_group", True),
            add_video_group=config.get("add_video_group", False),
            set_libva=config.get("set_libva", False),
            render_gid=config.get("render_gid"),
            video_gid=config.get("video_gid"),
        )
    return _dump(data)


# --------------------------------------------------------------------------- #
# Devices / privileged
# --------------------------------------------------------------------------- #


def add_device(yaml_str: str, service: str, host_path: str, target: str | None = None) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)
    entry = f"{host_path}:{target or host_path}"
    devices = svc.setdefault("devices", [])
    if not isinstance(devices, list):
        raise EditError(f"Service '{service}' devices must be a list")
    if entry not in devices:
        devices.append(entry)
    return _dump(data)


def remove_device(yaml_str: str, service: str, host_path: str) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)
    devices = svc.get("devices", [])
    if isinstance(devices, list):
        svc["devices"] = [
            d for d in devices if not str(d).startswith(host_path + ":") and d != host_path
        ]
        if not svc["devices"]:
            svc.pop("devices", None)
    return _dump(data)


def set_privileged(yaml_str: str, service: str, value: bool) -> str:
    data = _load(yaml_str)
    svc = _get_service(data, service)
    if value:
        svc["privileged"] = True
    else:
        svc.pop("privileged", None)
    return _dump(data)


def set_resources(
    yaml_str: str,
    service: str,
    *,
    cpus: float | None = None,
    memory: str | None = None,
    cpus_reserve: float | None = None,
    memory_reserve: str | None = None,
) -> str:
    """Set deploy.resources limits/reservations for a service.

    Pass None / empty to clear an individual value.
    """
    data = _load(yaml_str)
    svc = _get_service(data, service)
    deploy = svc.setdefault("deploy", {})
    resources = deploy.setdefault("resources", {})

    limits = resources.get("limits", {})
    if cpus:
        limits["cpus"] = str(cpus)
    else:
        limits.pop("cpus", None)
    if memory:
        limits["memory"] = memory
    else:
        limits.pop("memory", None)
    if limits:
        resources["limits"] = limits
    else:
        resources.pop("limits", None)

    reservations = resources.get("reservations", {})
    if cpus_reserve:
        reservations["cpus"] = str(cpus_reserve)
    else:
        reservations.pop("cpus", None)
    if memory_reserve:
        reservations["memory"] = memory_reserve
    else:
        reservations.pop("memory", None)
    # Don't drop reservations entirely — it may hold GPU devices.
    if reservations:
        resources["reservations"] = reservations
    elif "reservations" in resources and not resources["reservations"]:
        resources.pop("reservations", None)

    if not resources:
        deploy.pop("resources", None)
    if not deploy:
        svc.pop("deploy", None)
    return _dump(data)


def list_services(yaml_str: str) -> list[str]:
    data = _load(yaml_str)
    return list(data["services"].keys())


# --------------------------------------------------------------------------- #
# Secrets & configs (compose file-based)
# --------------------------------------------------------------------------- #


def add_secret(yaml_str: str, service: str, name: str, file_path: str) -> str:
    """Define a top-level file-based secret and attach it to ``service``."""
    data = _load(yaml_str)
    svc = _get_service(data, service)
    secrets = data.setdefault("secrets", {})
    if not isinstance(secrets, dict):
        raise EditError("`secrets` must be a mapping")
    secrets[name] = {"file": file_path}
    refs = svc.setdefault("secrets", [])
    if not isinstance(refs, list):
        raise EditError(f"Service '{service}' secrets must be a list")
    if name not in refs:
        refs.append(name)
    return _dump(data)


def remove_secret(yaml_str: str, service: str, name: str) -> str:
    """Detach a secret from ``service``; drop the top-level def if now unused."""
    data = _load(yaml_str)
    svc = data["services"].get(service)
    if isinstance(svc, dict) and isinstance(svc.get("secrets"), list):
        svc["secrets"] = [s for s in svc["secrets"] if s != name]
        if not svc["secrets"]:
            svc.pop("secrets", None)
    _prune_top_level(data, "secrets", name, _secret_still_used)
    return _dump(data)


def add_config(yaml_str: str, service: str, name: str, file_path: str, target: str) -> str:
    """Define a top-level file-based config and mount it into ``service`` at target."""
    data = _load(yaml_str)
    svc = _get_service(data, service)
    configs = data.setdefault("configs", {})
    if not isinstance(configs, dict):
        raise EditError("`configs` must be a mapping")
    configs[name] = {"file": file_path}
    refs = svc.setdefault("configs", [])
    if not isinstance(refs, list):
        raise EditError(f"Service '{service}' configs must be a list")
    if not any(isinstance(e, dict) and e.get("source") == name for e in refs):
        refs.append({"source": name, "target": target})
    return _dump(data)


def remove_config(yaml_str: str, service: str, name: str) -> str:
    data = _load(yaml_str)
    svc = data["services"].get(service)
    if isinstance(svc, dict) and isinstance(svc.get("configs"), list):
        svc["configs"] = [
            e for e in svc["configs"]
            if not (e == name or (isinstance(e, dict) and e.get("source") == name))
        ]
        if not svc["configs"]:
            svc.pop("configs", None)
    _prune_top_level(data, "configs", name, _config_still_used)
    return _dump(data)


def _secret_still_used(data: dict, name: str) -> bool:
    for svc in data.get("services", {}).values():
        if isinstance(svc, dict) and name in (svc.get("secrets") or []):
            return True
    return False


def _config_still_used(data: dict, name: str) -> bool:
    for svc in data.get("services", {}).values():
        if not isinstance(svc, dict):
            continue
        for entry in svc.get("configs") or []:
            if entry == name or (isinstance(entry, dict) and entry.get("source") == name):
                return True
    return False


def _prune_top_level(data: dict, key: str, name: str, still_used) -> None:
    top = data.get(key)
    if isinstance(top, dict) and name in top and not still_used(data, name):
        top.pop(name, None)