#!/bin/sh
set -eu

ONE_SHOT=${ONE_SHOT:-false}
SWEEP_INTERVAL_SEC=${SWEEP_INTERVAL_SEC:-21600}
BASE_THRESHOLD_DAYS=${BASE_THRESHOLD_DAYS:-14}
HIGH_USAGE_THRESHOLD_DAYS=${HIGH_USAGE_THRESHOLD_DAYS:-3}
HIGH_USAGE_PERCENT=${HIGH_USAGE_PERCENT:-70}
EMERGENCY_USAGE_PERCENT=${EMERGENCY_USAGE_PERCENT:-85}
LOG_RETENTION_DAYS=${LOG_RETENTION_DAYS:-7}
JOURNAL_MAX_SIZE=${JOURNAL_MAX_SIZE:-200M}
SKIP="registry.k8s.io/pause k8s.gcr.io/pause rancher/mirrored-pause"

sweep_once() {
  usage=$(df -P /host | awk 'NR==2 {gsub(/%/,"",$5); print $5}') || usage=""
  threshold_days="${BASE_THRESHOLD_DAYS}"
  if [ -n "${usage}" ] && [ "${usage}" -ge "${HIGH_USAGE_PERCENT}" ]; then
    threshold_days="${HIGH_USAGE_THRESHOLD_DAYS}"
  fi

  cutoff=$(THRESHOLD_DAYS="${threshold_days}" python3 - <<'PY'
import os
import time

days = int(os.environ.get("THRESHOLD_DAYS", "14"))
print(int(time.time()) - days * 86400)
PY
)

  RUNNING=$(chroot /host /bin/sh -c "crictl ps -a --quiet 2>/dev/null" | tr -s ' ' '\n' | sort -u | tr '\n' ' ')
  IMAGES_JSON=$(chroot /host /bin/sh -c "crictl images -o json 2>/dev/null" || echo '{}')

  prune_list=$(printf "%s" "${IMAGES_JSON}" | CUTOFF="${cutoff}" RUNNING="${RUNNING}" SKIP="${SKIP}" python3 - <<'PY'
import json
import os
import sys
import time

try:
    data = json.load(sys.stdin)
except Exception:
    print("", end="")
    sys.exit(0)

cutoff = int(os.environ.get("CUTOFF", "0"))
running = set(os.environ.get("RUNNING", "").split())
skip = os.environ.get("SKIP", "").split()
now = int(time.time())
prune = []


def is_skip(tags):
    if not tags:
        return False
    for t in tags:
        for prefix in skip:
            if prefix and t.startswith(prefix):
                return True
    return False


for img in data.get("images", []):
    image_id = img.get("id", "")
    if not image_id:
        continue
    if image_id in running:
        continue
    tags = img.get("repoTags") or []
    if is_skip(tags):
        continue
    created = img.get("createdAt") or 0
    try:
        created = int(str(created)) // 1000000000
    except Exception:
        created = 0
    if created and created > now:
        created = now
    if cutoff and created and created < cutoff:
        prune.append(image_id)

seen = set()
for p in prune:
    if p in seen:
        continue
    seen.add(p)
    print(p)
PY
)

  if [ -n "${prune_list}" ]; then
    printf "%s" "${prune_list}" | while read -r image_id; do
      if [ -n "${image_id}" ]; then
        chroot /host /bin/sh -c "crictl rmi --prune ${image_id}" || true
      fi
    done
  fi

  find /host/var/lib/rancher/k3s/agent/images -type f -name "*.tar" -mtime +7 -print -delete 2>/dev/null || true
  find /host/var/lib/rancher/k3s/agent/containerd -maxdepth 1 -type f -mtime +7 -print -delete 2>/dev/null || true

  if [ -n "${usage}" ] && [ "${usage}" -ge "${EMERGENCY_USAGE_PERCENT}" ]; then
    # Emergency pass for rootfs pressure on SD-backed nodes.
    chroot /host /bin/sh -c "journalctl --vacuum-size='${JOURNAL_MAX_SIZE}' >/dev/null 2>&1 || true"
    find /host/var/log -type f -name "*.gz" -mtime +"${LOG_RETENTION_DAYS}" -print -delete 2>/dev/null || true
    find /host/var/log/pods -type f -name "*.log" -mtime +"${LOG_RETENTION_DAYS}" -print -delete 2>/dev/null || true
    chroot /host /bin/sh -c "if command -v apt-get >/dev/null 2>&1; then apt-get clean >/dev/null 2>&1 || true; fi"
  fi
}

sweep_once

if [ "${ONE_SHOT}" = "true" ]; then
  exit 0
fi

while true; do
  sleep "${SWEEP_INTERVAL_SEC}"
  sweep_once
done