#!/usr/bin/env bash

#########################################################################################
#### container-update-checker.sh — Check Docker/Podman containers for image updates  ####
#### Compares local image digests against remote registry digests                     ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./container-update-checker.sh                                                 ####
####   ./container-update-checker.sh --docker --filter nginx                         ####
####   ./container-update-checker.sh --json --quiet                                  ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
RUNTIME="${CONTAINER_RUNTIME:-auto}"
TIMEOUT="${REGISTRY_TIMEOUT:-10}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
JSON_OUTPUT="false"
QUIET="false"
FILTER=""
LABEL=""
TEXTFILE_DIR="/var/lib/node_exporter"
PROM_FILE=""

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
COUNT_CURRENT=0
COUNT_UPDATE=0
COUNT_ERROR=0
COUNT_TOTAL=0
JSON_ITEMS=""
PROM_LINES=""

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m' GREEN='\033[0;32m' YELLOW='\033[0;33m'
        BOLD='\033[1m' DIM='\033[2m' RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*" >&2; fi; }

# ── Runtime Detection ─────────────────────────────────────────────────
detect_runtime() {
    if [[ "$RUNTIME" == "docker" || "$RUNTIME" == "podman" ]]; then
        if ! command -v "$RUNTIME" &>/dev/null; then
            err "${RUNTIME^} not found"; exit 2
        fi
        return
    fi
    if command -v docker &>/dev/null && docker info &>/dev/null; then
        RUNTIME="docker"
    elif command -v podman &>/dev/null; then
        RUNTIME="podman"
    else
        err "Neither Docker nor Podman found"; exit 2
    fi
    verbose "Auto-detected runtime: ${RUNTIME}"
}

# ── Auth Helper ───────────────────────────────────────────────────────
get_auth_header() {
    local registry="$1" config_file=""
    if [[ "$RUNTIME" == "podman" ]]; then
        config_file="${XDG_RUNTIME_DIR:-/run/user/$(id -u)}/containers/auth.json"
        [[ -f "$config_file" ]] || config_file="${HOME}/.config/containers/auth.json"
    fi
    [[ -f "${config_file:-}" ]] || config_file="${HOME}/.docker/config.json"
    [[ -f "$config_file" ]] || return 0
    local auth
    auth=$(grep -A1 "\"${registry}\"" "$config_file" 2>/dev/null \
        | grep '"auth"' | head -1 | sed 's/.*"auth"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/') || true
    if [[ -n "$auth" ]]; then
        echo "Authorization: Basic ${auth}"
    fi
}

# ── Parse Image Reference ────────────────────────────────────────────
parse_image_ref() {
    local image="$1" registry="" path="" tag=""
    local without_tag="${image%%@*}"
    if [[ "$without_tag" == *:* && "${without_tag##*:}" != */* ]]; then
        tag="${without_tag##*:}"
        without_tag="${without_tag%:*}"
    fi
    [[ -z "$tag" ]] && tag="latest"
    if [[ "$without_tag" == *"."*"/"* ]] || [[ "$without_tag" == *":"*"/"* ]] || [[ "$without_tag" == "localhost/"* ]]; then
        registry="${without_tag%%/*}"
        path="${without_tag#*/}"
    else
        registry="docker.io"
        [[ "$without_tag" == *"/"* ]] && path="$without_tag" || path="library/${without_tag}"
    fi
    echo "${registry}" "${path}" "${tag}"
}

# ── Get Local Digest ─────────────────────────────────────────────────
get_local_digest() {
    local image="$1" digest
    digest=$($RUNTIME image inspect "$image" --format '{{index .RepoDigests 0}}' 2>/dev/null) || true
    if [[ -n "$digest" && "$digest" == *"@"* ]]; then
        echo "${digest##*@}"; return
    fi
    digest=$($RUNTIME image inspect "$image" --format '{{.Id}}' 2>/dev/null) || true
    echo "${digest:-}"
}

# ── Extract JSON Value (pure bash, no python/jq) ─────────────────────
json_value() {
    local key="$1"
    sed -n "s/.*\"${key}\"[[:space:]]*:[[:space:]]*\"\([^\"]*\)\".*/\1/p" | head -1
}

# ── Get Remote Digest via Skopeo ──────────────────────────────────────
get_remote_digest_skopeo() {
    local registry="$1" path="$2" tag="$3"
    local digest
    digest=$(timeout "$TIMEOUT" skopeo inspect --no-tags "docker://${registry}/${path}:${tag}" 2>/dev/null \
        | json_value "Digest") || true
    echo "${digest:-}"
}

# ── Get Remote Digest via Curl ────────────────────────────────────────
get_remote_digest_curl() {
    local registry="$1" path="$2" tag="$3"
    local token="" digest=""
    if [[ "$registry" == "docker.io" || "$registry" == "registry-1.docker.io" ]]; then
        token=$(curl -sf --max-time "$TIMEOUT" \
            "https://auth.docker.io/token?service=registry.docker.io&scope=repository:${path}:pull" \
            | json_value "token") || true
        [[ -z "$token" ]] && return
        digest=$(curl -sf --max-time "$TIMEOUT" \
            -H "Accept: application/vnd.docker.distribution.manifest.v2+json" \
            -H "Accept: application/vnd.oci.image.index.v1+json" \
            -H "Authorization: Bearer ${token}" \
            "https://registry-1.docker.io/v2/${path}/manifests/${tag}" \
            -o /dev/null -D - 2>/dev/null \
            | grep -i "docker-content-digest" | tr -d '\r' | awk '{print $2}') || true
    else
        local auth_hdr auth_args=()
        auth_hdr=$(get_auth_header "$registry")
        [[ -n "$auth_hdr" ]] && auth_args=(-H "$auth_hdr")
        digest=$(curl -sf --max-time "$TIMEOUT" \
            -H "Accept: application/vnd.docker.distribution.manifest.v2+json" \
            -H "Accept: application/vnd.oci.image.index.v1+json" \
            "${auth_args[@]+"${auth_args[@]}"}" \
            "https://${registry}/v2/${path}/manifests/${tag}" \
            -o /dev/null -D - 2>/dev/null \
            | grep -i "docker-content-digest" | tr -d '\r' | awk '{print $2}') || true
    fi
    echo "${digest:-}"
}
# ── Get Remote Digest (skopeo then curl fallback) ─────────────────────
get_remote_digest() {
    local registry="$1" path="$2" tag="$3" digest=""
    if command -v skopeo &>/dev/null; then
        verbose "Trying skopeo for ${registry}/${path}:${tag}"
        digest=$(get_remote_digest_skopeo "$registry" "$path" "$tag")
    fi
    if [[ -z "$digest" ]]; then
        verbose "Trying curl fallback for ${registry}/${path}:${tag}"
        digest=$(get_remote_digest_curl "$registry" "$path" "$tag")
    fi
    echo "${digest:-}"
}

# ── Check Single Container ────────────────────────────────────────────
check_container() {
    local name="$1" image="$2"
    local status="" local_digest="" remote_digest="" registry path tag
    read -r registry path tag <<< "$(parse_image_ref "$image")"
    verbose "Container=${name} image=${image} registry=${registry} path=${path} tag=${tag}"
    local_digest=$(get_local_digest "$image")
    verbose "Local digest: ${local_digest:-none}"
    if [[ -z "$local_digest" ]]; then
        status="error"
    else
        remote_digest=$(get_remote_digest "$registry" "$path" "$tag")
        verbose "Remote digest: ${remote_digest:-none}"
        if [[ -z "$remote_digest" ]]; then
            status="error"
        elif [[ "$local_digest" == "$remote_digest" ]]; then
            status="current"
        else
            status="update"
        fi
    fi
    COUNT_TOTAL=$((COUNT_TOTAL + 1))
    case "$status" in
        current) COUNT_CURRENT=$((COUNT_CURRENT + 1)) ;;
        update)  COUNT_UPDATE=$((COUNT_UPDATE + 1)) ;;
        error)   COUNT_ERROR=$((COUNT_ERROR + 1)) ;;
    esac
    if [[ -n "$PROM_FILE" ]]; then
        local val=1; [[ "$status" == "update" ]] && val=0
        PROM_LINES+="container_image_up_to_date{name=\"${name}\",image=\"${image}\"} ${val}"$'\n'
    fi
    [[ "$QUIET" == "true" && "$status" != "update" ]] && return
    if [[ "$JSON_OUTPUT" == "true" ]]; then
        local item
        item=$(printf '{"container":"%s","image":"%s","status":"%s"}' "$name" "$image" "$status")
        [[ -n "$JSON_ITEMS" ]] && JSON_ITEMS="${JSON_ITEMS},${item}" || JSON_ITEMS="${item}"
    else
        local color symbol
        case "$status" in
            current) color="$GREEN";  symbol="up-to-date" ;;
            update)  color="$YELLOW"; symbol="update available" ;;
            error)   color="$RED";    symbol="check failed" ;;
            *)       color="";        symbol="?" ;;
        esac
        printf "  %-30s %-40s %b%s%b\n" "$name" "$image" "$color" "$symbol" "$RESET"
    fi
}

# ── List Containers ───────────────────────────────────────────────────
list_containers() {
    local filter_args=()
    [[ -n "$LABEL" ]] && filter_args+=(--filter "label=${LABEL}")
    $RUNTIME ps --format '{{.Names}}\t{{.Image}}' "${filter_args[@]}" 2>/dev/null
}

# ── Write Prometheus Metrics ──────────────────────────────────────────
write_prom_metrics() {
    local file="$1"
    local output_dir
    output_dir="$(dirname "$file")"
    mkdir -p "$output_dir"
    local tmp
    tmp=$(mktemp "${output_dir}/.container_updates.XXXXXX")
    {
        echo "# HELP container_image_up_to_date Whether the container image is up to date (1=yes, 0=no)"
        echo "# TYPE container_image_up_to_date gauge"
        printf '%s' "$PROM_LINES"
        echo "# HELP container_update_check_timestamp Unix timestamp of last update check"
        echo "# TYPE container_update_check_timestamp gauge"
        echo "container_update_check_timestamp $(date +%s)"
        echo "# HELP container_update_check_total Total containers checked"
        echo "# TYPE container_update_check_total gauge"
        echo "container_update_check_total ${COUNT_TOTAL}"
        echo "# HELP container_update_available_total Containers with updates available"
        echo "# TYPE container_update_available_total gauge"
        echo "container_update_available_total ${COUNT_UPDATE}"
    } > "$tmp"
    chmod 644 "$tmp"
    mv -f "$tmp" "$file"
    verbose "Metrics written to ${file}"
}

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Check Docker/Podman containers for available image updates

USAGE:
  ${SCRIPT_NAME} [OPTIONS]

OPTIONS:
  --docker              Force Docker as runtime
  --podman              Force Podman as runtime
  --filter NAME         Only check containers matching name pattern
  --label KEY=VALUE     Only check containers with matching label
  --json                Output results in JSON format
  --quiet               Only show containers with available updates
  --no-color            Disable colored output
  --textfile            Write metrics to node_exporter textfile collector
  -o, --output PATH     Write metrics to custom file path
  --timeout SECONDS     Registry check timeout (default: ${TIMEOUT})
  --verbose             Enable debug output
  --help                Show this help

ENVIRONMENT VARIABLES:
  CONTAINER_RUNTIME     Force runtime: docker or podman (default: auto)
  REGISTRY_TIMEOUT      Registry check timeout in seconds (default: 10)
  COLOR                 Color mode: auto, always, never (default: auto)
  VERBOSE               Enable debug output (true/false)

EXAMPLES:
  # Check all running containers
  ./${SCRIPT_NAME}

  # Force Docker, filter by name
  ./${SCRIPT_NAME} --docker --filter nginx

  # JSON output, quiet mode
  ./${SCRIPT_NAME} --json --quiet

  # Write Prometheus metrics to textfile collector
  ./${SCRIPT_NAME} --textfile

  # Write Prometheus metrics to custom path
  ./${SCRIPT_NAME} -o /tmp/container_updates.prom
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --docker)   RUNTIME="docker"; shift ;;
            --podman)   RUNTIME="podman"; shift ;;
            --filter)   FILTER="$2"; shift 2 ;;
            --label)    LABEL="$2"; shift 2 ;;
            --json)     JSON_OUTPUT="true"; shift ;;
            --quiet)    QUIET="true"; shift ;;
            --no-color) COLOR="never"; shift ;;
            --textfile) PROM_FILE="$TEXTFILE_DIR/container_updates.prom"; shift ;;
            -o|--output) PROM_FILE="$2"; shift 2 ;;
            --timeout)  TIMEOUT="$2"; shift 2 ;;
            --verbose)  VERBOSE="true"; shift ;;
            --help|-h)  setup_colors; usage; exit 0 ;;
            -*)
                err "Unknown option: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 2 ;;
            *)
                err "Unexpected argument: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 2 ;;
        esac
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    parse_args "$@"
    setup_colors
    detect_runtime

    local containers=()
    while IFS=$'\t' read -r name image; do
        [[ -z "$name" ]] && continue
        [[ -n "$FILTER" && "$name" != *"${FILTER}"* ]] && continue
        containers+=("${name}	${image}")
    done < <(list_containers)

    if [[ ${#containers[@]} -eq 0 ]]; then
        if [[ "$JSON_OUTPUT" == "true" ]]; then
            echo '{"results":[],"summary":{"total":0,"current":0,"update_available":0,"errors":0}}'
        else
            warn "No running containers found"
        fi
        exit 0
    fi

    verbose "Found ${#containers[@]} containers to check"

    if [[ "$JSON_OUTPUT" != "true" ]]; then
        echo ""
        echo -e "${BOLD}Container Update Checker${RESET}"
        echo -e "${DIM}Runtime: ${RUNTIME} | Timeout: ${TIMEOUT}s${RESET}"
        echo ""
        printf "  ${BOLD}%-30s %-40s %s${RESET}\n" "CONTAINER" "IMAGE" "STATUS"
        printf "  %s\n" "$(printf '%.0s─' {1..82})"
    fi

    for entry in "${containers[@]}"; do
        check_container "${entry%%	*}" "${entry#*	}"
    done

    if [[ "$JSON_OUTPUT" == "true" ]]; then
        printf '{"results":[%s],"summary":{"total":%d,"current":%d,"update_available":%d,"errors":%d}}\n' \
            "$JSON_ITEMS" "$COUNT_TOTAL" "$COUNT_CURRENT" "$COUNT_UPDATE" "$COUNT_ERROR"
    else
        echo ""
        echo -e "  ${BOLD}Summary${RESET}"
        printf "  %-20s %d\n" "Total checked:" "$COUNT_TOTAL"
        printf "  %-20s %b%d%b\n" "Up-to-date:" "$GREEN" "$COUNT_CURRENT" "$RESET"
        printf "  %-20s %b%d%b\n" "Update available:" "$YELLOW" "$COUNT_UPDATE" "$RESET"
        printf "  %-20s %b%d%b\n" "Errors:" "$RED" "$COUNT_ERROR" "$RESET"
        echo ""
    fi

    [[ -n "$PROM_FILE" ]] && write_prom_metrics "$PROM_FILE"

    if [[ "$COUNT_ERROR" -gt 0 ]]; then exit 2
    elif [[ "$COUNT_UPDATE" -gt 0 ]]; then exit 1
    fi
    exit 0
}

main "$@"