stackpilot/backend/services/compose_service.py

"""File-based stack storage and Docker Compose lifecycle.

The compose YAML on disk is always the source of truth. The database only
stores metadata (name, description, timestamps).
"""
from __future__ import annotations

import asyncio
import os
import re
import shutil
import tempfile
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

from config import settings
from docker_client import DockerError, get_client, safe_call

COMPOSE_FILENAMES = ("compose.yaml", "compose.yml", "docker-compose.yml", "docker-compose.yaml")
DEFAULT_COMPOSE_NAME = "compose.yaml"
COMPOSE_LABEL = "com.docker.compose.project"
SERVICE_LABEL = "com.docker.compose.service"


class StackFileError(Exception):
    pass


# --------------------------------------------------------------------------- #
# Slug / paths
# --------------------------------------------------------------------------- #


def slugify(name: str) -> str:
    slug = re.sub(r"[^a-z0-9_-]+", "-", name.strip().lower())
    slug = re.sub(r"-{2,}", "-", slug).strip("-_")
    return slug or "stack"


def stacks_root(override: Optional[str] = None) -> str:
    return override or settings.STACKS_DIR


def stack_dir(stack_id: str, override: Optional[str] = None) -> str:
    return os.path.join(stacks_root(override), stack_id)


def find_compose_file(directory: str) -> Optional[str]:
    for name in COMPOSE_FILENAMES:
        candidate = os.path.join(directory, name)
        if os.path.isfile(candidate):
            return candidate
    return None


def compose_path(stack_id: str, override: Optional[str] = None) -> str:
    directory = stack_dir(stack_id, override)
    return find_compose_file(directory) or os.path.join(directory, DEFAULT_COMPOSE_NAME)


def env_path(stack_id: str, override: Optional[str] = None) -> str:
    return os.path.join(stack_dir(stack_id, override), ".env")


# --------------------------------------------------------------------------- #
# Read / write files
# --------------------------------------------------------------------------- #


def read_compose(stack_id: str, override: Optional[str] = None) -> str:
    path = compose_path(stack_id, override)
    if not os.path.isfile(path):
        return ""
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read()


def read_env(stack_id: str, override: Optional[str] = None) -> str:
    path = env_path(stack_id, override)
    if not os.path.isfile(path):
        return ""
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read()


def write_compose(stack_id: str, content: str, override: Optional[str] = None) -> None:
    directory = stack_dir(stack_id, override)
    os.makedirs(directory, exist_ok=True)
    path = compose_path(stack_id, override)
    # Non-destructive: back up existing file first.
    if os.path.isfile(path):
        shutil.copy2(path, path + ".bak")
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        fh.write(content)
    os.replace(tmp, path)


def write_env(stack_id: str, content: str, override: Optional[str] = None) -> None:
    directory = stack_dir(stack_id, override)
    os.makedirs(directory, exist_ok=True)
    path = env_path(stack_id, override)
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(content)


def delete_stack_files(stack_id: str, override: Optional[str] = None) -> None:
    directory = stack_dir(stack_id, override)
    if os.path.isdir(directory):
        shutil.rmtree(directory)


def clone_stack_files(src_id: str, dst_id: str, override: Optional[str] = None) -> None:
    src = stack_dir(src_id, override)
    dst = stack_dir(dst_id, override)
    if os.path.isdir(dst):
        raise StackFileError(f"Target stack '{dst_id}' already exists")
    shutil.copytree(src, dst)


def discover_stacks(override: Optional[str] = None) -> list[str]:
    """Return ids of all directories under STACKS_DIR that contain a compose file."""
    root = stacks_root(override)
    if not os.path.isdir(root):
        return []
    found = []
    for entry in sorted(os.listdir(root)):
        directory = os.path.join(root, entry)
        if os.path.isdir(directory) and find_compose_file(directory):
            found.append(entry)
    return found


# --------------------------------------------------------------------------- #
# Live status from Docker
# --------------------------------------------------------------------------- #


@dataclass
class ContainerInfo:
    id: str
    name: str
    service: str
    image: str
    state: str  # running, exited, ...
    status: str  # human string
    health: Optional[str] = None
    ports: list[dict] = field(default_factory=list)
    created: Optional[str] = None


def _parse_ports(attrs: dict) -> list[dict]:
    ports = []
    bindings = (attrs.get("NetworkSettings") or {}).get("Ports") or {}
    for container_port, host in (bindings or {}).items():
        if host:
            for binding in host:
                ports.append(
                    {
                        "container": container_port,
                        "host_ip": binding.get("HostIp"),
                        "host_port": binding.get("HostPort"),
                    }
                )
        else:
            ports.append({"container": container_port, "host_port": None})
    return ports


def containers_for_stack(stack_id: str) -> list[ContainerInfo]:
    client = get_client()
    raw = safe_call(
        client.containers.list,
        all=True,
        filters={"label": f"{COMPOSE_LABEL}={stack_id}"},
    )
    result = []
    for c in raw:
        attrs = c.attrs
        state = attrs.get("State", {}) or {}
        health = (state.get("Health") or {}).get("Status")
        result.append(
            ContainerInfo(
                id=c.id,
                name=c.name,
                service=c.labels.get(SERVICE_LABEL, c.name),
                # Use the configured image name from the inspect we already have;
                # reading c.image.tags would trigger a separate image-inspect call
                # per container.
                image=attrs.get("Config", {}).get("Image", "") or attrs.get("Image", ""),
                state=state.get("Status", c.status),
                status=attrs.get("State", {}).get("Status", c.status),
                health=health,
                ports=_parse_ports(attrs),
                created=attrs.get("Created"),
            )
        )
    return result


# in-memory set of stacks currently performing a pull/up
_BUSY: set[str] = set()


def mark_busy(stack_id: str) -> None:
    _BUSY.add(stack_id)


def clear_busy(stack_id: str) -> None:
    _BUSY.discard(stack_id)


def is_busy(stack_id: str) -> bool:
    return stack_id in _BUSY


def _status_from_states(states: list[str]) -> str:
    if not states:
        return "stopped"
    if any(s == "dead" for s in states):
        return "error"
    running = [s for s in states if s == "running"]
    if len(running) == len(states):
        return "running"
    if running:
        return "partial"
    return "stopped"


def compute_status(stack_id: str, containers: Optional[list[ContainerInfo]] = None) -> str:
    """Status for one stack. Pass already-fetched ``containers`` to avoid a
    redundant Docker round-trip (the detail view already has them)."""
    if stack_id in _BUSY:
        return "updating"
    try:
        if containers is None:
            containers = containers_for_stack(stack_id)
    except DockerError:
        return "unknown"
    return _status_from_states([c.state for c in containers])


def stack_status_summaries() -> dict[str, dict]:
    """One cheap Docker call → ``{project: {status, total, running}}`` for every
    stack at once.

    Uses the low-level container *summary* list (``GET /containers/json``, no
    per-container inspect) grouped by the compose project label. This is far
    cheaper than calling :func:`containers_for_stack` (which full-inspects every
    container) once per stack for the overview — turning the stacks list from
    O(stacks × containers) Docker round-trips into a single one.
    """
    client = get_client()
    raw = safe_call(client.api.containers, all=True)
    by_project: dict[str, list[str]] = {}
    for c in raw:
        project = (c.get("Labels") or {}).get(COMPOSE_LABEL)
        if project:
            by_project.setdefault(project, []).append(c.get("State", ""))
    return {
        project: {
            "status": _status_from_states(states),
            "total": len(states),
            "running": sum(1 for s in states if s == "running"),
        }
        for project, states in by_project.items()
    }


# --------------------------------------------------------------------------- #
# Compose CLI lifecycle (async subprocess)
# --------------------------------------------------------------------------- #


def _compose_base_cmd(stack_id: str, override: Optional[str] = None) -> list[str]:
    directory = stack_dir(stack_id, override)
    compose_file = find_compose_file(directory) or os.path.join(directory, DEFAULT_COMPOSE_NAME)
    return [
        "docker",
        "compose",
        "-p",
        stack_id,
        "--project-directory",
        directory,
        "-f",
        compose_file,
    ]


async def run_compose(
    stack_id: str,
    args: list[str],
    override: Optional[str] = None,
    timeout: float = 600.0,
) -> dict:
    """Run a `docker compose` subcommand. Returns {returncode, stdout, stderr}."""
    cmd = _compose_base_cmd(stack_id, override) + args
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError as exc:
        proc.kill()
        raise StackFileError(f"compose command timed out: {' '.join(args)}") from exc
    return {
        "returncode": proc.returncode,
        "stdout": stdout_b.decode("utf-8", "replace"),
        "stderr": stderr_b.decode("utf-8", "replace"),
        "command": " ".join(args),
    }


async def validate_yaml(content: str, env_content: str = "") -> dict:
    """Validate compose YAML by running ``docker compose config -q`` against it
    in a throwaway directory. Returns ``{"ok": bool, "errors": str}`` — the
    errors string carries compose's own diagnostics (parse errors, unknown
    keys, bad interpolation) so the editor can surface them before saving."""
    with tempfile.TemporaryDirectory(prefix="sp-validate-") as tmp:
        compose_file = os.path.join(tmp, DEFAULT_COMPOSE_NAME)
        with open(compose_file, "w", encoding="utf-8") as fh:
            fh.write(content)
        if env_content:
            with open(os.path.join(tmp, ".env"), "w", encoding="utf-8") as fh:
                fh.write(env_content)
        cmd = [
            "docker", "compose",
            "--project-directory", tmp,
            "-f", compose_file,
            "config", "-q",
        ]
        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            _out, err = await asyncio.wait_for(proc.communicate(), timeout=30.0)
        except asyncio.TimeoutError as exc:
            raise StackFileError("compose config validation timed out") from exc
        ok = proc.returncode == 0
        return {"ok": ok, "errors": "" if ok else err.decode("utf-8", "replace").strip()}


async def stream_compose(
    stack_id: str, args: list[str], override: Optional[str] = None
):
    """Yield lines from a `docker compose` subcommand as they are produced."""
    cmd = _compose_base_cmd(stack_id, override) + args
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )
    assert proc.stdout is not None
    async for raw in proc.stdout:
        yield raw.decode("utf-8", "replace").rstrip("\n")
    await proc.wait()


_json_progress: Optional[bool] = None


async def supports_json_progress() -> bool:
    """Whether this Docker Compose understands ``--progress json``.

    The JSON progress stream carries per-layer ``current``/``total`` bytes, which
    the deploy console turns into a real progress bar. Older compose releases
    reject the value, so probe once (cheap, no side effects) and cache it; on a
    negative result callers fall back to the plain text stream.
    """
    global _json_progress
    if _json_progress is None:
        try:
            proc = await asyncio.create_subprocess_exec(
                "docker", "compose", "--progress", "json", "version",
                stdout=asyncio.subprocess.DEVNULL,
                stderr=asyncio.subprocess.DEVNULL,
            )
            rc = await asyncio.wait_for(proc.wait(), timeout=15.0)
            _json_progress = rc == 0
        except Exception:  # noqa: BLE001 - probe failure just disables the feature
            _json_progress = False
    return _json_progress


async def stream_up(stack_id: str, override: Optional[str] = None):
    """Run `compose up -d` streaming combined output, so the deploy console can
    show image-pull and container-create progress live.

    Yields ``("log", line)`` for each output line, then ``("done", returncode)``.
    """
    cmd = _compose_base_cmd(stack_id, override)
    if await supports_json_progress():
        # Global flag, must precede the subcommand.
        cmd += ["--progress", "json"]
    cmd += ["up", "-d", "--remove-orphans"]
    proc = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )
    assert proc.stdout is not None
    async for raw in proc.stdout:
        yield ("log", raw.decode("utf-8", "replace").rstrip("\n"))
    await proc.wait()
    yield ("done", proc.returncode)


# Convenience lifecycle wrappers ------------------------------------------------


async def up(stack_id: str, override: Optional[str] = None) -> dict:
    mark_busy(stack_id)
    try:
        return await run_compose(stack_id, ["up", "-d", "--remove-orphans"], override)
    finally:
        clear_busy(stack_id)


async def down(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["down"], override)


async def start(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["start"], override)


async def stop(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["stop"], override)


async def restart(stack_id: str, override: Optional[str] = None) -> dict:
    return await run_compose(stack_id, ["restart"], override)


async def pull(stack_id: str, override: Optional[str] = None) -> dict:
    mark_busy(stack_id)
    try:
        return await run_compose(stack_id, ["pull"], override)
    finally:
        clear_busy(stack_id)


async def update(stack_id: str, override: Optional[str] = None) -> dict:
    """Pull then up -d."""
    mark_busy(stack_id)
    try:
        pull_res = await run_compose(stack_id, ["pull"], override)
        up_res = await run_compose(stack_id, ["up", "-d", "--remove-orphans"], override)
        return {
            "returncode": up_res["returncode"],
            "stdout": pull_res["stdout"] + "\n" + up_res["stdout"],
            "stderr": pull_res["stderr"] + "\n" + up_res["stderr"],
            "command": "pull + up -d",
        }
    finally:
        clear_busy(stack_id)


async def logs(
    stack_id: str,
    service: Optional[str] = None,
    tail: int = 200,
    override: Optional[str] = None,
) -> dict:
    args = ["logs", "--no-color", "--tail", str(tail), "--timestamps"]
    if service:
        args.append(service)
    return await run_compose(stack_id, args, override, timeout=60.0)


def now() -> datetime:
    return datetime.now(timezone.utc)