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)
    if "volumes" in data and not data["volumes"]:
        data.pop("volumes", 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")

    gpu_service.remove_gpu(svc)

    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),
        )
    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 list_services(yaml_str: str) -> list[str]:
    data = _load(yaml_str)
    return list(data["services"].keys())