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
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),
                image=(c.image.tags[0] if c.image and c.image.tags else attrs.get("Config", {}).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 compute_status(stack_id: str) -> str:
    if stack_id in _BUSY:
        return "updating"
    try:
        containers = containers_for_stack(stack_id)
    except DockerError:
        return "unknown"
    if not containers:
        return "stopped"
    states = [c.state for c in containers]
    if any(s in ("dead",) for s in states):
        return "error"
    if any(
        c.state == "exited" and _nonzero_exit(c) for c in containers
    ):
        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 _nonzero_exit(c: ContainerInfo) -> bool:
    # We only have the textual state here; treat plain "exited" as stopped, not
    # an error unless health says otherwise. Detailed exit codes handled in detail view.
    return False


# --------------------------------------------------------------------------- #
# 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 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()


# 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)