# Image cache janitor script.
# Periodically checks container image cache size via crictl and removes
# unused images (largest first) when the cache exceeds ~500 GiB.
apiVersion: v1
kind: ConfigMap
metadata:
  name: image-cache-janitor-script
  namespace: kube-system
  labels:
    app: image-cache-janitor
data:
  janitor.py: |
    #!/usr/bin/env python3
    """Image cache janitor DaemonSet agent.

    Periodically checks the container image cache size on each node via
    crictl and removes unused images (largest first) when the cache
    exceeds the configured limit. This keeps image cache usage within the
    ~500 GiB EBS budget without relying on kubelet's imageGC thresholds
    (which operate on total disk usage and would misfire on pod ephemeral
    storage spikes).

    Host access: crictl is bundled in the container image and talks to the
    node's containerd via a mounted CRI socket.
    Safety: crictl rmi refuses to remove images that are in use by running
    containers, so eviction is always safe.
    """

    import json
    import logging
    import os
    import socket
    import subprocess
    import sys
    import threading
    import time
    from dataclasses import dataclass, field
    from http.server import HTTPServer, BaseHTTPRequestHandler

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        datefmt="%Y-%m-%dT%H:%M:%S",
        stream=sys.stdout,
    )
    log = logging.getLogger("janitor")

    # Configuration from environment
    IMAGE_CACHE_LIMIT_GI = int(os.environ.get("IMAGE_CACHE_LIMIT_GI", "500"))
    IMAGE_CACHE_TARGET_GI = int(os.environ.get("IMAGE_CACHE_TARGET_GI", "400"))
    CHECK_INTERVAL_SECONDS = int(os.environ.get("CHECK_INTERVAL_SECONDS", "300"))
    METRICS_PORT = 9103

    # Convert GiB to bytes
    GI = 1024 ** 3
    IMAGE_CACHE_LIMIT_BYTES = IMAGE_CACHE_LIMIT_GI * GI
    IMAGE_CACHE_TARGET_BYTES = IMAGE_CACHE_TARGET_GI * GI


    # ---------------------------------------------------------------------------
    # Image info
    # ---------------------------------------------------------------------------

    @dataclass
    class ImageInfo:
        """Parsed container image metadata from crictl."""

        id: str
        repo_tags: list[str] = field(default_factory=list)
        repo_digests: list[str] = field(default_factory=list)
        size: int = 0
        pinned: bool = False


    def parse_crictl_images(json_str: str) -> list[ImageInfo]:
        """Parse ``crictl images -o json`` output into ImageInfo objects."""
        data = json.loads(json_str)
        images_raw = data.get("images", [])
        result: list[ImageInfo] = []
        for img in images_raw:
            size_val = img.get("size", 0)
            if isinstance(size_val, str):
                size_val = int(size_val)
            result.append(
                ImageInfo(
                    id=img.get("id", ""),
                    repo_tags=img.get("repoTags") or [],
                    repo_digests=img.get("repoDigests") or [],
                    size=size_val,
                    pinned=bool(img.get("pinned", False)),
                )
            )
        return result


    def calculate_total_cache_size(images: list[ImageInfo]) -> int:
        """Sum the size of all images in bytes."""
        return sum(img.size for img in images)


    def select_images_to_remove(
        images: list[ImageInfo],
        total_size: int,
        limit_bytes: int,
        target_bytes: int,
    ) -> list[ImageInfo]:
        """Select unused images to remove to bring cache under target.

        Strategy: filter out pinned images, sort remaining by size
        descending (largest first), accumulate until total drops below
        target. Returns empty list if total is under the limit.
        """
        if total_size <= limit_bytes:
            return []

        removable = [img for img in images if not img.pinned]
        removable.sort(key=lambda img: img.size, reverse=True)

        to_remove: list[ImageInfo] = []
        projected = total_size
        for img in removable:
            if projected <= target_bytes:
                break
            to_remove.append(img)
            projected -= img.size

        return to_remove


    # ---------------------------------------------------------------------------
    # Prometheus metrics server (stdlib only, no pip dependencies)
    # ---------------------------------------------------------------------------

    class DualStackHTTPServer(HTTPServer):
        # Bind ::0 with IPV6_V6ONLY=0 so the socket accepts both IPv6 and
        # IPv4-mapped IPv6 connections (single socket serves both families).
        # Required because runner pods on IPv6-only EKS have no IPv4 address
        # but PodMonitor scrapes can originate from either family.
        address_family = socket.AF_INET6

        def server_bind(self) -> None:
            self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
            super().server_bind()


    class MetricsServer:
        """Thread-safe Prometheus metrics with an HTTP endpoint on /metrics."""

        _METRIC_DEFS = [
            ("image_cache_janitor_cache_size_bytes", "gauge",
             "Total size of container image cache in bytes"),
            ("image_cache_janitor_cache_image_count", "gauge",
             "Number of container images in the cache"),
            ("image_cache_janitor_gc_cycles_total", "counter",
             "Total number of GC check cycles run"),
            ("image_cache_janitor_gc_evictions_total", "counter",
             "Total number of images evicted"),
            ("image_cache_janitor_gc_eviction_errors_total", "counter",
             "Total number of image eviction failures (in-use images)"),
            ("image_cache_janitor_gc_evicted_bytes_total", "counter",
             "Total bytes freed by eviction"),
            ("image_cache_janitor_last_cycle_timestamp", "gauge",
             "Unix timestamp of the last GC cycle"),
            ("image_cache_janitor_cache_limit_bytes", "gauge",
             "Configured cache size limit in bytes"),
            ("image_cache_janitor_cache_target_bytes", "gauge",
             "Configured cache target size in bytes"),
        ]

        def __init__(self) -> None:
            self._lock = threading.Lock()
            self._values: dict[str, float] = {}
            for name, _mtype, _mhelp in self._METRIC_DEFS:
                self._values[name] = 0.0

        def set(self, name: str, value: float) -> None:
            with self._lock:
                self._values[name] = value

        def inc(self, name: str, delta: float = 1.0) -> None:
            with self._lock:
                self._values[name] = self._values.get(name, 0.0) + delta

        def format(self) -> str:
            lines: list[str] = []
            with self._lock:
                snapshot = dict(self._values)
            for name, mtype, mhelp in self._METRIC_DEFS:
                value = snapshot.get(name, 0.0)
                lines.append(f"# HELP {name} {mhelp}")
                lines.append(f"# TYPE {name} {mtype}")
                lines.append(f"{name} {value}")
            lines.append("")  # trailing newline
            return "\n".join(lines)

        def start(self) -> None:
            """Start the HTTP server in a daemon thread."""
            metrics_ref = self

            class Handler(BaseHTTPRequestHandler):
                def do_GET(self) -> None:
                    if self.path == "/metrics":
                        body = metrics_ref.format().encode()
                        self.send_response(200)
                        self.send_header(
                            "Content-Type",
                            "text/plain; version=0.0.4; charset=utf-8",
                        )
                        self.send_header("Content-Length", str(len(body)))
                        self.end_headers()
                        self.wfile.write(body)
                    else:
                        self.send_response(404)
                        self.end_headers()

                def log_message(self, fmt, *args) -> None:  # noqa: ARG002
                    pass

            server = DualStackHTTPServer(("::", METRICS_PORT), Handler)
            t = threading.Thread(target=server.serve_forever, daemon=True)
            t.start()
            log.info("Metrics server listening on :%d/metrics", METRICS_PORT)


    # Global metrics instance
    metrics = MetricsServer()


    # ---------------------------------------------------------------------------
    # crictl interaction (host access via nsenter)
    # ---------------------------------------------------------------------------

    def run_crictl(*args: str) -> str:
        """Run a crictl command against the node's containerd CRI socket.

        crictl is bundled in the container image; CONTAINER_RUNTIME_ENDPOINT
        env var tells it which socket to use.
        Returns stdout as a string. Raises CalledProcessError on failure.
        """
        cmd = ["crictl"] + list(args)
        log.info("$ %s", " ".join(cmd))
        result = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=120)
        return result.stdout


    def remove_image(image_id: str) -> bool:
        """Remove a single image via crictl rmi.

        Returns True if removed, False if the image is in use (or other error).
        crictl rmi is safe — it refuses to remove images used by running containers.
        """
        try:
            run_crictl("rmi", image_id)
            return True
        except subprocess.CalledProcessError as exc:
            log.warning("Failed to remove image %s: %s", image_id, exc.stderr.strip())
            return False


    # ---------------------------------------------------------------------------
    # GC cycle
    # ---------------------------------------------------------------------------

    def _run_gc_cycle() -> None:
        """Run one garbage collection cycle."""
        metrics.inc("image_cache_janitor_gc_cycles_total")
        metrics.set("image_cache_janitor_last_cycle_timestamp", time.time())

        # List all images
        try:
            output = run_crictl("images", "-o", "json")
        except subprocess.CalledProcessError as exc:
            log.error("Failed to list images: %s", exc.stderr.strip())
            return

        images = parse_crictl_images(output)
        total_size = calculate_total_cache_size(images)

        metrics.set("image_cache_janitor_cache_size_bytes", float(total_size))
        metrics.set("image_cache_janitor_cache_image_count", float(len(images)))

        total_gi = total_size / GI
        log.info(
            "Image cache: %.1f GiB (%d images), limit: %d GiB, target: %d GiB",
            total_gi, len(images), IMAGE_CACHE_LIMIT_GI, IMAGE_CACHE_TARGET_GI,
        )

        if total_size <= IMAGE_CACHE_LIMIT_BYTES:
            log.info("Cache within limit, no eviction needed")
            return

        to_remove = select_images_to_remove(
            images, total_size, IMAGE_CACHE_LIMIT_BYTES, IMAGE_CACHE_TARGET_BYTES,
        )

        if not to_remove:
            log.info("No removable images found (all pinned)")
            return

        log.info(
            "Evicting %d images (%.1f GiB) to reach target",
            len(to_remove),
            sum(img.size for img in to_remove) / GI,
        )

        for img in to_remove:
            tags = ", ".join(img.repo_tags) if img.repo_tags else img.id[:12]
            log.info("Removing: %s (%.1f MiB)", tags, img.size / (1024 ** 2))
            if remove_image(img.id):
                metrics.inc("image_cache_janitor_gc_evictions_total")
                metrics.inc("image_cache_janitor_gc_evicted_bytes_total", float(img.size))
            else:
                metrics.inc("image_cache_janitor_gc_eviction_errors_total")


    # ---------------------------------------------------------------------------
    # Main
    # ---------------------------------------------------------------------------

    def main() -> None:
        log.info("Image cache janitor starting")
        log.info("Limit: %d GiB, Target: %d GiB, Interval: %ds",
                 IMAGE_CACHE_LIMIT_GI, IMAGE_CACHE_TARGET_GI, CHECK_INTERVAL_SECONDS)

        # Validate crictl is accessible
        try:
            version_output = run_crictl("version")
            log.info("crictl version: %s", version_output.strip().split("\n")[0])
        except (subprocess.CalledProcessError, FileNotFoundError) as exc:
            log.error("Cannot access crictl: %s", exc)
            log.error("Ensure containerd socket is mounted and CONTAINER_RUNTIME_ENDPOINT is set")
            sys.exit(1)

        # Start metrics server
        metrics.start()
        metrics.set("image_cache_janitor_cache_limit_bytes", float(IMAGE_CACHE_LIMIT_BYTES))
        metrics.set("image_cache_janitor_cache_target_bytes", float(IMAGE_CACHE_TARGET_BYTES))

        # Main loop
        while True:
            try:
                _run_gc_cycle()
            except Exception:
                log.exception("Unexpected error in GC cycle")
            time.sleep(CHECK_INTERVAL_SECONDS)


    if __name__ == "__main__":
        main()
