#!/usr/bin/env python3

"""This script replaces `podman auto-update`. If a digest in a registry has been updated, `podman
pull` may not necessarily pull the image for quite some time. However, `auto-update` will still try
every day, but since it doesn't check if the new digest has actually been pulled it will restart the
service again and again. This script attempts to solve the problem by explicitly checking the digest
after the pull.

"""

import collections
import functools
import json
import subprocess


def podman_ps():
    out = subprocess.run(["podman", "ps", "--format", "json"], capture_output=True, check=True)
    return json.loads(out.stdout)


def podman_image_inspect(image):
    out = subprocess.run(["podman", "image", "inspect", "--format", "json", image],
                         capture_output=True, check=True)
    inspect = json.loads(out.stdout)
    assert inspect
    if len(inspect) > 1:
        raise ValueError("podman image inspect returned multiple entries")
    return inspect[0]


def podman_pull(image):
    subprocess.run(["podman", "pull", image], capture_output=True, check=True)


if __name__ == "__main__":
    containers = podman_ps()

    units = collections.defaultdict(list)
    for container in containers:
        labels = container.get("Labels", None)
        if labels is None:
            continue

        autoupdate = labels.get("io.containers.autoupdate", "disabled")
        if autoupdate == "disabled":
            continue
        if autoupdate != "image":
            raise ValueError(f"unrecognised autopdate label: {autoupdate}")

        if "PODMAN_SYSTEMD_UNIT" not in labels:
            raise ValueError(
                f"container {container['Names'][0]} does not have \"PODMAN_SYSTEMD_UNIT\" label")
        units[container["Image"]].append(labels["PODMAN_SYSTEMD_UNIT"])

    updated = set()
    for image in units.keys():
        inspect = podman_image_inspect(image)
        original_digest = inspect["Digest"]

        podman_pull(image)

        inspect = podman_image_inspect(image)
        new_digest = inspect["Digest"]

        if new_digest != original_digest:
            updated.add(image)

    if updated:
        print(f"The following images have been updated: {', '.join(updated)}")
        restart = functools.reduce(lambda x, y: x + y, [units[image] for image in updated])

        print(f"The following units will be restarted: {', '.join(restart)}")
        subprocess.run(["systemctl", "--user", "restart"] + restart, check=True)