#!/usr/bin/env python3
"""Steward's swarm prune (maintenance/swarm_prune.yml), run on each node.

Removes the images on this node that no swarm service is pinned to and no
container uses (by id, or by the reference it was made from), and keeps every image a service is pinned to even when its
task runs on another node. A node then holds what it needs to take over the
other's services without a pull: the registry may be behind the node that
failed, and a digest it has since deleted can't be pulled at all.

The node must be a swarm manager, the only kind that can list services. On
one that isn't, or when Docker doesn't answer, it removes nothing and fails,
so a prune never runs blind. With --dry-run it says what it would remove.
Prints one JSON line: kept, removed, failed, freed_bytes, dry_run.

Stdlib only; shipped beside the playbook in Steward's built-in source.
"""
from __future__ import annotations

import http.client
import json
import socket
import subprocess
import sys

DEADLINE_SECONDS = 300
DOCKER_SOCKET = "/var/run/docker.sock"


def docker(*args: str) -> str:
    done = subprocess.run(["docker", *args], capture_output=True, text=True,
                          timeout=DEADLINE_SECONDS)
    if done.returncode != 0:
        raise RuntimeError(f"docker {' '.join(args[:2])}: {done.stderr.strip()[-300:]}")
    return done.stdout


class _UnixConnection(http.client.HTTPConnection):
    def __init__(self, path: str):
        super().__init__("localhost", timeout=DEADLINE_SECONDS)
        self.unix_path = path

    def connect(self):
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.settimeout(self.timeout)
        sock.connect(self.unix_path)
        self.sock = sock


def api(path: str):
    """A GET on Docker's own API. The image list is read here, not from
    `docker image ls`, which leaves out images with no tag: every image a
    swarm task pulled by digest."""
    conn = _UnixConnection(DOCKER_SOCKET)
    try:
        conn.request("GET", path)
        resp = conn.getresponse()
        body = resp.read()
    finally:
        conn.close()
    if resp.status != 200:
        raise RuntimeError(f"docker API {path}: HTTP {resp.status}")
    return json.loads(body)


def short(repo: str) -> str:
    """A repository the way Docker lists it locally: no docker.io/ or library/."""
    for prefix in ("docker.io/", "index.docker.io/"):
        if repo.startswith(prefix):
            repo = repo[len(prefix):]
    return repo[len("library/"):] if repo.startswith("library/") else repo


def wanted(service_images) -> tuple[set, set]:
    """(digests, "repo:tag" refs) the services are pinned to. A spec with a
    digest is matched by it; one with a tag alone by that tag."""
    digests, refs = set(), set()
    for image in service_images:
        ref, _, digest = (image or "").partition("@")
        if digest:
            digests.add(digest)
        elif ref:
            name = ref.rsplit("/", 1)[-1]
            refs.add(short(ref if ":" in name else ref + ":latest"))
    return digests, refs


def removable(images: list, digests: set, refs: set, used: set) -> list:
    """The images (Docker API /images/json rows) no container uses and no
    service is pinned to."""
    out = []
    for im in images:
        if im.get("Id") in used:
            continue
        held = {d.split("@", 1)[1] for d in im.get("RepoDigests") or [] if "@" in d}
        tags = {short(t) for t in im.get("RepoTags") or []}
        if held & digests or tags & refs:
            continue
        out.append(im)
    return out


def label(im: dict) -> str:
    names = (im.get("RepoTags") or []) + (im.get("RepoDigests") or [])
    return names[0] if names else im.get("Id", "")[7:19]


def main(argv: list) -> int:
    dry = "--dry-run" in argv
    try:
        ids = docker("service", "ls", "-q").split()
        services = json.loads(docker("service", "inspect", *ids)) if ids else []
        ids = docker("ps", "-aq", "--no-trunc").split()
        containers = json.loads(docker("container", "inspect", *ids)) if ids else []
        images = api("/images/json")
    except (RuntimeError, ValueError, OSError, subprocess.TimeoutExpired, http.client.HTTPException) as exc:
        print(json.dumps({"error": f"nothing removed: {exc}"}))
        return 2
    # A container's image is kept by its id and by the reference it was made
    # from: on Docker's containerd store the id alone hasn't held a tag.
    used = {c.get("Image") for c in containers}
    digests, refs = wanted([s.get("Spec", {}).get("TaskTemplate", {}).get("ContainerSpec", {}).get("Image")
                            for s in services]
                           + [c.get("Config", {}).get("Image") for c in containers])
    gone = removable(images, digests, refs, used)
    removed, failed, freed = [], [], 0
    for im in gone:
        if not dry:
            try:
                docker("image", "rm", "-f", im["Id"])
            except (RuntimeError, subprocess.TimeoutExpired) as exc:
                failed.append({"image": label(im), "error": str(exc)[-200:]})
                continue
        removed.append(label(im))
        freed += int(im.get("Size") or 0)
    print(json.dumps({"kept": len(images) - len(gone), "removed": removed, "failed": failed,
                      "freed_bytes": freed, "dry_run": dry}))
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))