#!/bin/bash
# -----------------------------------------------------------------------------------------------
# --------------------------------- Docker Watchdog --------------------------------------------
# -----------------------------------------------------------------------------------------------
# Self-healing watchdog for Docker containers — monitors memory, CPU and HTTP responsiveness.
# Restarts containers that exceed configured thresholds using a strike system for CPU and
# responsiveness checks to avoid restarting on brief spikes.
#
# Behaviour:
#   Memory  — immediate restart if hard limit is exceeded
#   CPU     — strike system, restarts after CPU_FAIL_LIMIT consecutive over-threshold checks
#   HTTP    — strike system, restarts after RESP_FAIL_LIMIT consecutive failed curl checks
#
# Strike system:
#   Strikes persist between runs via WATCHDOG_STATE_FILE (/tmp — resets on reboot)
#   Strike cadence depends on cron schedule:
#     Every 15min + 2 strikes = 30min sustained abuse before restart
#     Every 10min + 2 strikes = 20min sustained abuse before restart
#     Every  5min + 2 strikes = 10min sustained abuse before restart
#
# All configuration lives in Master.conf under the Docker Watchdog section.
# Supports --dry-run to show what would be restarted without taking any action.
# -----------------------------------------------------------------------------------------------

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

source "$SCRIPT_DIR/../Master.conf"
source "$SCRIPT_DIR/../common.sh"

parse_args "$@"

# Auto-detect total CPU cores for normalisation
TOTAL_CORES=$(nproc)

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_GEAR Setup ━━━
# -----------------------------------------------------------------------------------------------
echo ""
echo "━━━ $ICON_GEAR Setup ━━━"

# ROOT CHECK
if [[ "$EUID" -ne 0 ]]; then
    error "Must be run as root"
    exit 1
fi

success "Running as root"
info "$ICON_WATCHDOG Watchdog initialising — $TOTAL_CORES cores detected"

# Ensure state file exists
touch "$WATCHDOG_STATE_FILE" 2>/dev/null || {
    error "Cannot create state file: $WATCHDOG_STATE_FILE"
    exit 1
}

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_SUMMARY Status ━━━
# -----------------------------------------------------------------------------------------------
if [[ "$SHOW_STATUS" == true ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY STATUS ━━━━━"
    echo "$ICON_WATCHDOG Containers monitored: ${!WATCHDOG_CONTAINERS[*]}"
    echo "$ICON_MEM      Soft mem threshold:   ${SOFT_MEM_THRESHOLD}% of per-container limit"
    echo "$ICON_ZFS      CPU soft threshold:   ${SOFT_CPU_THRESHOLD}%"
    echo "$ICON_ZFS      CPU hard threshold:   ${HARD_CPU_THRESHOLD}%"
    echo "$ICON_RETRY    CPU fail limit:       ${CPU_FAIL_LIMIT} strikes"
    echo "$ICON_PING     Resp fail limit:      ${RESP_FAIL_LIMIT} strikes"
    echo "$ICON_TIME     Curl timeout:         ${CURL_TIMEOUT}s"
    echo "$ICON_GEAR     Dry Run:              $DRY_RUN"
    echo "━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

[[ "$DRY_RUN" == true ]] && warn "DRY RUN — no containers will be restarted"

# -----------------------------------------------------------------------------------------------
# STATE HELPERS
# Reads and writes per-container strike counts to the state file.
# State file format: container:metric:count
# -----------------------------------------------------------------------------------------------

# Returns current strike count for a container/metric pair.
# Usage: get_strikes "Emby" "CPU"
get_strikes() {
    local container="$1" metric="$2"
    grep -E "^${container}:${metric}:" "$WATCHDOG_STATE_FILE" 2>/dev/null | cut -d':' -f3
}

# Sets strike count for a container/metric pair.
# Usage: set_strikes "Emby" "CPU" 2
set_strikes() {
    local container="$1" metric="$2" count="$3"
    grep -vE "^${container}:${metric}:" "$WATCHDOG_STATE_FILE" 2>/dev/null > "${WATCHDOG_STATE_FILE}.tmp"
    echo "${container}:${metric}:${count}" >> "${WATCHDOG_STATE_FILE}.tmp"
    mv "${WATCHDOG_STATE_FILE}.tmp" "$WATCHDOG_STATE_FILE"
}

# -----------------------------------------------------------------------------------------------
# HELPERS
# -----------------------------------------------------------------------------------------------

# Fetches memory and CPU stats for a container in a single docker stats call.
# Returns: MEM_USAGE|CPU_PERCENT
parse_stats() {
    local container="$1"
    docker stats --no-stream --format "{{.MemUsage}}|{{.CPUPerc}}" "$container"
}

# Converts a memory value and unit to MB.
# Supports KiB, MiB, GiB — returns UNKNOWN for unrecognised units.
convert_to_mb() {
    local value="$1" unit="$2"
    case "$unit" in
        KiB) awk "BEGIN {print $value / 1024}" ;;
        MiB) echo "$value" ;;
        GiB) awk "BEGIN {print $value * 1024}" ;;
        *) echo "UNKNOWN" ;;
    esac
}

# Restarts a container locally.
# In dry run mode reports what would happen without acting.
restart_container() {
    local container="$1" reason="$2"

    if [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would restart $container ($reason)"
        return
    fi

    info "Restarting $container ($reason)..."
    if docker restart "$container" >/dev/null 2>&1; then
        echo "$ICON_STARTED $container restarted"
        log "Restarted $container — reason: $reason"
    else
        error "Failed to restart $container"
    fi
}

# -----------------------------------------------------------------------------------------------
# MEMORY CHECK
# Compares current container memory usage against its configured hard limit.
# Restarts immediately if at or above 100% of limit.
# Warns if at or above SOFT_MEM_THRESHOLD % of limit.
# Usage: check_memory "Emby" 16384
# -----------------------------------------------------------------------------------------------
check_memory() {
    local container="$1" limit_mb="$2"
    local stats mem_raw mem_val mem_unit mem_mb usage_pct

    stats=$(parse_stats "$container")
    mem_raw=$(echo "$stats" | awk -F'|' '{print $1}' | awk '{print $1}')
    mem_val=$(echo "$mem_raw" | sed -E 's/([0-9.]+).*/\1/')
    mem_unit=$(echo "$mem_raw" | sed -E 's/[0-9.]+([a-zA-Z]+).*/\1/')
    mem_mb=$(convert_to_mb "$mem_val" "$mem_unit")
    local mem_int
    mem_int=$(printf "%.0f" "$mem_mb")
    usage_pct=$(( (mem_int * 100) / limit_mb ))

    if (( usage_pct >= 100 )); then
        error "$ICON_MEM $container memory ${mem_int}MB (${usage_pct}%) — exceeded ${limit_mb}MB hard limit"
        restart_container "$container" "memory hard limit"
        set_strikes "$container" "CPU" 0
        set_strikes "$container" "RESP" 0
    elif (( usage_pct >= SOFT_MEM_THRESHOLD )); then
        warn "$ICON_MEM $container memory ${mem_int}MB (${usage_pct}% of ${limit_mb}MB limit)"
    else
        success "$ICON_MEM $container memory ${mem_int}MB (${usage_pct}% of ${limit_mb}MB limit)"
    fi
}

# -----------------------------------------------------------------------------------------------
# CPU CHECK
# Normalises CPU usage against total core count and applies the strike system.
# Warns at SOFT_CPU_THRESHOLD, strikes at HARD_CPU_THRESHOLD.
# Restarts after CPU_FAIL_LIMIT consecutive strikes — resets strikes on restart or recovery.
# Usage: check_cpu "Emby"
# -----------------------------------------------------------------------------------------------
check_cpu() {
    local container="$1"
    local stats cpu_raw cpu_norm cpu_int violations

    stats=$(parse_stats "$container")
    cpu_raw=$(echo "$stats" | awk -F'|' '{print $2}' | tr -d '%')
    cpu_norm=$(awk "BEGIN {print $cpu_raw / $TOTAL_CORES}")
    cpu_int=$(printf "%.0f" "$cpu_norm")

    violations=$(get_strikes "$container" "CPU")
    [[ -z "$violations" ]] && violations=0

    if (( cpu_int >= HARD_CPU_THRESHOLD )); then
        ((violations++))
        error "$ICON_ZFS $container CPU ${cpu_int}% — hard threshold hit ($violations/$CPU_FAIL_LIMIT strikes)"
        set_strikes "$container" "CPU" "$violations"
    elif (( cpu_int >= SOFT_CPU_THRESHOLD )); then
        ((violations++))
        warn "$ICON_ZFS $container CPU ${cpu_int}% — soft threshold hit ($violations/$CPU_FAIL_LIMIT strikes)"
        set_strikes "$container" "CPU" "$violations"
    else
        if (( violations > 0 )); then
            info "$ICON_ZFS $container CPU ${cpu_int}% — recovered, resetting strikes"
        else
            success "$ICON_ZFS $container CPU ${cpu_int}%"
        fi
        set_strikes "$container" "CPU" 0
        violations=0
    fi

    if (( violations >= CPU_FAIL_LIMIT )); then
        error "$ICON_ZFS $container hit CPU limit for $CPU_FAIL_LIMIT consecutive checks"
        restart_container "$container" "sustained CPU abuse"
        set_strikes "$container" "CPU" 0
    fi
}

# -----------------------------------------------------------------------------------------------
# RESPONSIVENESS CHECK
# Sends an HTTP request to the container's configured URL.
# Skips containers with no URL defined in WATCHDOG_CONTAINER_URLS.
# Applies the same strike system as CPU — restarts after RESP_FAIL_LIMIT consecutive failures.
# Usage: check_responsiveness "Emby"
# -----------------------------------------------------------------------------------------------
check_responsiveness() {
    local container="$1"
    local url="${WATCHDOG_CONTAINER_URLS[$container]:-}"

    [[ -z "$url" ]] && return

    local fails
    fails=$(get_strikes "$container" "RESP")
    [[ -z "$fails" ]] && fails=0

    if ! curl -s --max-time "$CURL_TIMEOUT" "$url" >/dev/null 2>&1; then
        ((fails++))
        warn "$ICON_PING $container unresponsive at $url ($fails/$RESP_FAIL_LIMIT strikes)"
        set_strikes "$container" "RESP" "$fails"
    else
        if (( fails > 0 )); then
            info "$ICON_PING $container responsive again — resetting strikes"
        else
            success "$ICON_PING $container responsive at $url"
        fi
        set_strikes "$container" "RESP" 0
        fails=0
    fi

    if (( fails >= RESP_FAIL_LIMIT )); then
        error "$ICON_PING $container unresponsive for $RESP_FAIL_LIMIT consecutive checks"
        restart_container "$container" "HTTP unresponsive"
        set_strikes "$container" "RESP" 0
    fi
}

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_WATCHDOG Watchdog Check ━━━
# -----------------------------------------------------------------------------------------------
echo ""
echo "━━━ $ICON_WATCHDOG Watchdog Check — $(date '+%Y-%m-%d %H:%M:%S') ━━━"
echo ""

RESTARTED=()
SKIPPED=()

START=$(date +%s)

for container in "${!WATCHDOG_CONTAINERS[@]}"; do
    echo "━━━ $ICON_CONTAINERS $container ━━━"

    # Verify container exists
    if ! docker inspect "$container" &>/dev/null; then
        warn "$container not found on this host — skipping"
        SKIPPED+=("$container")
        echo ""
        continue
    fi

    # Verify container is running
    if ! docker ps --filter "name=^/${container}$" --format "{{.Names}}" | grep -qw "$container"; then
        warn "$ICON_NOT_RUNNING $container is not running — skipping"
        SKIPPED+=("$container")
        echo ""
        continue
    fi

    check_memory "$container" "${WATCHDOG_CONTAINERS[$container]}"
    check_cpu "$container"
    check_responsiveness "$container"

    echo ""
done

END=$(date +%s)

# -----------------------------------------------------------------------------------------------
# ━━━ $ICON_SUMMARY Summary ━━━
# -----------------------------------------------------------------------------------------------
echo "━━━━━ $ICON_SUMMARY WATCHDOG SUMMARY ━━━━━"
echo "$ICON_TIME     $(date '+%Y-%m-%d %H:%M:%S')"
echo "$ICON_TIME     Duration:  $(format_duration $((END - START)))"
echo "$ICON_WATCHDOG Monitored: ${#WATCHDOG_CONTAINERS[@]} containers"
echo "$ICON_NOT_RUNNING Skipped: ${#SKIPPED[@]} containers"
if [[ "$DRY_RUN" == true ]]; then
    echo "$ICON_WARN     Dry Run:   no restarts executed"
fi
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"