#!/usr/bin/env python3
"""Careful Docker pruning for this compose/Sablier host.

The important guardrail is that networks are not pruned by default. Sablier
keeps some containers intentionally stopped, and deleting their compose-created
networks can leave those stopped containers unable to start later.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from dataclasses import dataclass


PROTECT_LABEL = "prune.protect"
PROTECT_VALUE = "true"
NETWORK_NEVER_REMOVE = {"bridge", "host", "none", "proxy", "mail_internal"}


@dataclass
class DockerObject:
    id: str
    name: str
    labels: dict[str, str]
    state: str = ""

    @property
    def protected(self) -> bool:
        return self.labels.get(PROTECT_LABEL, "").lower() == PROTECT_VALUE


def docker(args: list[str], *, check: bool = True) -> str:
    result = subprocess.run(
        ["docker", *args],
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if check and result.returncode != 0:
        sys.stderr.write(result.stderr)
        raise SystemExit(result.returncode)
    return result.stdout


def parse_labels(raw: str | None) -> dict[str, str]:
    labels: dict[str, str] = {}
    if not raw:
        return labels
    for item in raw.split(","):
        if not item:
            continue
        key, sep, value = item.partition("=")
        labels[key] = value if sep else ""
    return labels


def list_containers() -> list[DockerObject]:
    rows = docker(["container", "ls", "-a", "--format", "{{json .}}"]).splitlines()
    containers: list[DockerObject] = []
    for row in rows:
        data = json.loads(row)
        containers.append(
            DockerObject(
                id=data["ID"],
                name=data["Names"],
                labels=parse_labels(data.get("Labels")),
                state=data.get("State", ""),
            )
        )
    return containers


def list_networks() -> list[DockerObject]:
    rows = docker(["network", "ls", "--format", "{{json .}}"]).splitlines()
    networks: list[DockerObject] = []
    for row in rows:
        data = json.loads(row)
        networks.append(
            DockerObject(
                id=data["ID"],
                name=data["Name"],
                labels=parse_labels(data.get("Labels")),
            )
        )
    return networks


def network_has_containers(network_id: str) -> bool:
    raw = docker(["network", "inspect", network_id, "--format", "{{json .Containers}}"])
    containers = json.loads(raw)
    return bool(containers)


def remove_containers(dry_run: bool) -> int:
    removable_states = {"created", "exited", "dead"}
    candidates = [
        c
        for c in list_containers()
        if c.state.lower() in removable_states and not c.protected
    ]
    if not candidates:
        print("containers: none")
        return 0
    print("containers:")
    for container in candidates:
        print(f"  remove {container.name} ({container.state})")
        if not dry_run:
            docker(["container", "rm", container.id])
    return len(candidates)


def prune_images(dry_run: bool) -> None:
    cmd = ["image", "prune", "-af", "--filter", f"label!={PROTECT_LABEL}={PROTECT_VALUE}"]
    print("images:")
    if dry_run:
        print("  would run: docker " + " ".join(cmd))
        return
    print(docker(cmd).rstrip() or "  none")


def prune_build_cache(dry_run: bool) -> None:
    cmd = ["builder", "prune", "-af"]
    print("builder-cache:")
    if dry_run:
        print("  would run: docker " + " ".join(cmd))
        return
    print(docker(cmd).rstrip() or "  none")


def prune_networks(dry_run: bool) -> int:
    candidates: list[DockerObject] = []
    for network in list_networks():
        if network.name in NETWORK_NEVER_REMOVE or network.protected:
            continue
        if network_has_containers(network.id):
            continue
        candidates.append(network)

    if not candidates:
        print("networks: none")
        return 0

    print("networks:")
    for network in candidates:
        print(f"  remove {network.name}")
        if not dry_run:
            docker(["network", "rm", network.id])
    return len(candidates)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true", help="show what would be removed")
    parser.add_argument(
        "--prune-networks",
        action="store_true",
        help="also remove unused, unprotected networks",
    )
    args = parser.parse_args()

    print(f"mode: {'dry-run' if args.dry_run else 'execute'}")
    removed_containers = remove_containers(args.dry_run)
    prune_images(args.dry_run)
    prune_build_cache(args.dry_run)
    if args.prune_networks:
        prune_networks(args.dry_run)
    else:
        print("networks: skipped")
    print(f"summary: removable_containers={removed_containers}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())