wild-pc/castle-api/src/castle_api/services.py

"""Service management — start/stop/restart systemd-managed components."""

from __future__ import annotations

import asyncio
import time

from fastapi import APIRouter, HTTPException, status
from starlette.responses import JSONResponse

from castle_core.generators.systemd import (
    generate_timer,
    generate_unit_from_deployed,
    unit_name,
)

from castle_api.config import get_castle_root, get_registry
from castle_api.health import check_all_health
from castle_api.models import HealthStatus
from castle_api.stream import broadcast

router = APIRouter(prefix="/services", tags=["services"])

UNIT_PREFIX = "castle-"
SELF_NAME = "castle-api"


async def _systemctl(action: str, unit: str) -> tuple[bool, str]:
    """Run a systemctl --user command. Returns (success, output)."""
    proc = await asyncio.create_subprocess_exec(
        "systemctl",
        "--user",
        action,
        unit,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, stderr = await proc.communicate()
    output = (stdout or stderr or b"").decode().strip()
    return proc.returncode == 0, output


async def _get_unit_status(unit: str) -> str:
    """Get the active status of a systemd unit."""
    proc = await asyncio.create_subprocess_exec(
        "systemctl",
        "--user",
        "is-active",
        unit,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    stdout, _ = await proc.communicate()
    return (stdout or b"").decode().strip()


def _managed(name: str):
    """The managed deployment with this name (a name may span kinds; take the
    managed systemd one — service or job), or None."""
    return next((d for d in get_registry().named(name) if d.managed), None)


def _validate_managed(name: str) -> None:
    """Raise 404 if the component isn't managed in the registry."""
    if _managed(name) is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"'{name}' is not a managed service",
        )


async def _broadcast_health_with_override(
    override_name: str, override_status: str
) -> None:
    """Run health checks but override one component's status from systemd."""
    registry = get_registry()
    statuses = await check_all_health(registry)

    result = []
    for s in statuses:
        if s.id == override_name:
            result.append(
                HealthStatus(
                    id=override_name,
                    status="down" if override_status != "active" else "up",
                    latency_ms=None,
                )
            )
        else:
            result.append(s)

    await broadcast(
        "health",
        {
            "statuses": [s.model_dump() for s in result],
            "timestamp": time.time(),
        },
    )


async def _deferred_systemctl(action: str, unit: str, delay: float = 0.5) -> None:
    """Run a systemctl action after a delay, allowing the HTTP response to flush."""
    await asyncio.sleep(delay)
    await _systemctl(action, unit)


async def _do_action(name: str, action: str) -> JSONResponse:
    """Execute a systemctl action and broadcast updated health."""
    deployed = _managed(name)
    unit = (
        unit_name(name, deployed.kind) if deployed else f"{UNIT_PREFIX}{name}.service"
    )

    # Self-restart: defer the systemctl call so the response can be sent first
    if name == SELF_NAME and action in ("restart", "stop"):
        asyncio.create_task(_deferred_systemctl(action, unit))
        return JSONResponse(
            status_code=202,
            content={"program": name, "action": action, "status": "accepted"},
        )

    ok, output = await _systemctl(action, unit)
    unit_status = await _get_unit_status(unit)

    if not ok:
        raise HTTPException(status_code=500, detail=output or f"Failed to {action}")

    await _broadcast_health_with_override(name, unit_status)

    return JSONResponse(
        content={"program": name, "action": action, "status": unit_status},
    )


@router.get("/{name}/unit")
def get_unit(name: str) -> dict[str, str | None]:
    """Return the generated systemd unit file(s) for a managed component."""
    _validate_managed(name)
    deployed = _managed(name)

    # Get systemd spec from config if repo available
    systemd_spec = None
    schedule = None
    description = None
    root = get_castle_root()
    if root:
        from castle_core.config import load_config

        config = load_config(root)
        dep = config.deployment(deployed.kind, name)
        if dep is not None:
            manage = getattr(dep, "manage", None)
            if manage and manage.systemd:
                systemd_spec = manage.systemd
            description = dep.description
            schedule = getattr(dep, "schedule", None)

    unit = generate_unit_from_deployed(name, deployed, systemd_spec)
    timer = generate_timer(name, schedule, description) if schedule else None
    return {"service": unit, "timer": timer}


@router.post("/{name}/restart")
async def restart_service(name: str) -> JSONResponse:
    """Restart a systemd-managed service — the imperative bounce.

    Lifecycle (start/stop/enable/disable) is convergence now: set `enabled` in the
    deployment config and POST /apply. Restart stays as a way to re-actualize the
    current desired state without changing it.
    """
    return await _do_action(name, "restart")