#!/bin/bash
# Pi-Kit DietPi image prep + check script
# Cleans host-unique data and optionally verifies the image state.

set -euo pipefail

SCRIPT_PATH="$(readlink -f "${BASH_SOURCE[0]}")"
SCRIPT_NAME="$(basename "$SCRIPT_PATH")"

PIKIT_HOST="${PIKIT_HOST:-pikit.local}"
PIKIT_USER="${PIKIT_USER:-dietpi}"
PIKIT_SSH_KEY="${PIKIT_SSH_KEY:-$HOME/.ssh/pikit}"
PIKIT_SSH_OPTS="${PIKIT_SSH_OPTS:-}"
PIKIT_REMOTE_TMP="${PIKIT_REMOTE_TMP:-/tmp/pikit-prep.sh}"
PIKIT_SELF_DELETE="${PIKIT_SELF_DELETE:-0}"

MODE="both"
LOCAL_ONLY=0

ERRORS=0
WARNINGS=0

usage() {
  cat <<'USAGE'
Usage: pikit-prep.sh [--prep-only|--check-only] [--local]

Defaults to prep + check (combined). When run on a non-DietPi host, it will
copy itself to the Pi (/tmp) and run with sudo, then clean up.

Options:
  --prep-only   Run prep only (no check)
  --check-only  Run checks only (no prep)
  --local       Force local execution (no SSH copy)
  --help        Show this help
USAGE
}

status() {
  local level="$1"
  shift
  printf '[%s] %s\n' "$level" "$*"
  case "$level" in
    FAIL) ERRORS=$((ERRORS + 1)) ;;
    WARN) WARNINGS=$((WARNINGS + 1)) ;;
  esac
}

section() {
  printf '\n== %s ==\n' "$1"
}

is_dietpi() {
  grep -qi "dietpi" /etc/os-release 2>/dev/null
}

require_root() {
  if [ "${EUID:-$(id -u)}" -ne 0 ]; then
    echo "[FAIL] This script must run as root (use sudo)." >&2
    exit 1
  fi
}

parse_args() {
  for arg in "$@"; do
    case "$arg" in
      --prep-only) MODE="prep" ;;
      --check-only) MODE="check" ;;
      --local) LOCAL_ONLY=1 ;;
      --help|-h) usage; exit 0 ;;
      *)
        echo "[FAIL] Unknown argument: $arg" >&2
        usage
        exit 1
        ;;
    esac
  done
}

run_remote() {
  local forward=()
  for arg in "$@"; do
    [ "$arg" = "--local" ] && continue
    forward+=("$arg")
  done
  if ! command -v scp >/dev/null 2>&1 || ! command -v ssh >/dev/null 2>&1; then
    echo "[FAIL] ssh/scp not available for remote prep" >&2
    exit 1
  fi
  scp -i "$PIKIT_SSH_KEY" $PIKIT_SSH_OPTS "$SCRIPT_PATH" "${PIKIT_USER}@${PIKIT_HOST}:${PIKIT_REMOTE_TMP}"
  ssh -i "$PIKIT_SSH_KEY" $PIKIT_SSH_OPTS "${PIKIT_USER}@${PIKIT_HOST}" \
    "sudo PIKIT_SELF_DELETE=1 bash ${PIKIT_REMOTE_TMP} --local ${forward[*]}; rc=\$?; rm -f ${PIKIT_REMOTE_TMP}; exit \$rc"
  exit $?
}

clean_logs_dir() {
  local dir="$1" pattern="${2:-*}"
  if [ -d "$dir" ]; then
    find "$dir" -type f -name "$pattern" -print -delete 2>/dev/null | sed 's/^/[CLEANED] /' || true
    status CLEANED "logs $pattern in $dir"
  else
    status SKIP "$dir (missing)"
  fi
}

truncate_file() {
  local file="$1"
  if [ -e "$file" ]; then
    :> "$file" && status CLEANED "truncated $file" || status FAIL "truncate $file"
  else
    status SKIP "$file (missing)"
  fi
}

clean_file() {
  local path="$1"
  if [ -e "$path" ]; then
    rm -f "$path" && status CLEANED "$path" || status FAIL "$path"
  else
    status SKIP "$path (missing)"
  fi
}

clean_dir_files() {
  local dir="$1" pattern="$2"
  if [ -d "$dir" ]; then
    find "$dir" -type f -name "$pattern" -print -delete 2>/dev/null | sed 's/^/[CLEANED] /' || true
    status CLEANED "files $pattern in $dir"
  else
    status SKIP "$dir (missing)"
  fi
}

truncate_dir() {
  local dir="$1"
  if [ -d "$dir" ]; then
    # keep systemd-private dirs intact while services run
    find "$dir" -mindepth 1 -maxdepth 1 ! -path "$dir/systemd-private-*" -exec rm -rf {} + 2>/dev/null
    status CLEANED "$dir"
  else
    status SKIP "$dir (missing)"
  fi
}

clean_backups() {
  local dir="$1"
  if [ -d "$dir" ]; then
    find "$dir" -type f \( -name '*~' -o -name '*.bak*' -o -name '*.orig*' \) -print -delete 2>/dev/null | sed 's/^/[CLEANED] /' || true
    status CLEANED "backup/editor files in $dir"
  else
    status SKIP "$dir (missing)"
  fi
}

clean_home_dir() {
  local dir="$1"
  shift
  local keep=("$@")
  if [ -d "$dir" ]; then
    shopt -s dotglob nullglob
    for entry in "$dir"/*; do
      local base
      base="$(basename "$entry")"
      case "$base" in
        .|..) continue ;;
      esac
      local keep_it=0
      for k in "${keep[@]}"; do
        if [ "$base" = "$k" ]; then
          keep_it=1
          break
        fi
      done
      if [ "$keep_it" -eq 0 ]; then
        rm -rf "$entry" && status CLEANED "$dir/$base" || status FAIL "$dir/$base"
      fi
    done
    shopt -u dotglob nullglob
  else
    status SKIP "$dir (missing)"
  fi
}

prep_image() {
  section "Prep"

  # --- Identity ---
  truncate -s 0 /etc/machine-id && status CLEANED /etc/machine-id || status FAIL /etc/machine-id
  mkdir -p /var/lib/dbus || true
  rm -f /var/lib/dbus/machine-id
  ln -s /etc/machine-id /var/lib/dbus/machine-id && status CLEANED "/var/lib/dbus/machine-id -> /etc/machine-id" || status FAIL "/var/lib/dbus/machine-id"
  clean_file /var/lib/systemd/random-seed

  # --- SSH host keys ---
  if ls /etc/ssh/ssh_host_* >/dev/null 2>&1; then
    rm -f /etc/ssh/ssh_host_* && status CLEANED "SSH host keys" || status FAIL "SSH host keys"
  else
    status SKIP "SSH host keys (none)"
  fi

  # --- SSH client traces ---
  :> /root/.ssh/known_hosts 2>/dev/null && status CLEANED "/root/.ssh/known_hosts" || status SKIP "/root/.ssh/known_hosts"
  :> /home/dietpi/.ssh/known_hosts 2>/dev/null && status CLEANED "/home/dietpi/.ssh/known_hosts" || status SKIP "/home/dietpi/.ssh/known_hosts"
  :> /home/dietpi/.ssh/authorized_keys 2>/dev/null && status CLEANED "/home/dietpi/.ssh/authorized_keys" || status SKIP "/home/dietpi/.ssh/authorized_keys"
  :> /root/.ssh/authorized_keys 2>/dev/null && status CLEANED "/root/.ssh/authorized_keys" || status SKIP "/root/.ssh/authorized_keys"

  # --- Default login ---
  if id -u dietpi >/dev/null 2>&1; then
    echo "dietpi:pikit" | chpasswd && status CLEANED "reset dietpi password" || status FAIL "reset dietpi password"
  else
    status SKIP "dietpi user missing"
  fi

  # --- Shell history ---
  :> /root/.bash_history 2>/dev/null && status CLEANED "/root/.bash_history" || status SKIP "/root/.bash_history"
  :> /home/dietpi/.bash_history 2>/dev/null && status CLEANED "/home/dietpi/.bash_history" || status SKIP "/home/dietpi/.bash_history"

  # --- Home directories ---
  clean_home_dir /home/dietpi .bashrc .bash_logout .profile .ssh RESCUE.md
  clean_home_dir /root .bashrc .bash_logout .profile .ssh RESCUE.md

  # --- Ready flag ---
  clean_file /var/run/pikit-ready

  # --- First-boot state + TLS ---
  clean_dir_files /var/lib/pikit/firstboot "*"
  clean_file /var/lib/pikit/firstboot/firstboot.done
  clean_file /var/lib/pikit/firstboot/state.json
  clean_file /var/lib/pikit/firstboot/firstboot.log
  clean_file /var/lib/pikit/firstboot/firstboot.error
  clean_file /var/lib/pikit/firstboot/firstboot.lock
  clean_file /etc/pikit/certs/pikit-ca.crt
  clean_file /etc/pikit/certs/pikit-ca.key
  clean_file /etc/pikit/certs/pikit-ca.srl
  clean_file /etc/pikit/certs/pikit.local.crt
  clean_file /etc/pikit/certs/pikit.local.key
  clean_file /etc/pikit/certs/pikit.local.csr
  clean_file /var/www/pikit-web/assets/pikit-ca.crt

  # --- Backup/editor cruft ---
  clean_backups /var/www/pikit-web
  clean_backups /usr/local/bin

  # --- Logs ---
  clean_dir_files /var/log "*"
  clean_dir_files /var/log/nginx "*"
  if [ -d /var/log/journal ]; then
    find /var/log/journal -type f -print -delete 2>/dev/null | sed 's/^/[CLEANED] /' || true
    status CLEANED "/var/log/journal"
  else
    status SKIP "/var/log/journal (missing)"
  fi
  if [ -d /var/crash ]; then
    find /var/crash -type f -print -delete 2>/dev/null | sed 's/^/[CLEANED] /' || true
    status CLEANED "/var/crash"
  else
    status SKIP "/var/crash (missing)"
  fi

  if command -v pihole >/dev/null 2>&1; then
    pihole -f >/dev/null 2>&1 && status CLEANED "pihole logs via pihole -f" || status FAIL "pihole -f"
    clean_logs_dir /var/log/pihole '*'
    clean_file /etc/pihole/pihole-FTL.db
  fi

  if [ -x /opt/AdGuardHome/AdGuardHome ]; then
    clean_logs_dir /var/opt/AdGuardHome/data/logs '*'
    clean_file /opt/AdGuardHome/data/querylog.db
  fi

  if command -v ufw >/dev/null 2>&1; then
    truncate_file /var/log/ufw.log
  fi

  if command -v fail2ban-client >/dev/null 2>&1; then
    truncate_file /var/log/fail2ban.log
  fi

  clean_logs_dir /var/log/unbound '*'
  clean_logs_dir /var/log/dnsmasq '*'
  clean_logs_dir /var/log/powerdns '*'
  clean_logs_dir /var/lib/technitium-dns/Logs '*'

  clean_logs_dir /var/log/jellyfin '*'
  clean_logs_dir /var/lib/jellyfin/log '*'
  clean_logs_dir /var/log/jellyseerr '*'
  clean_logs_dir /opt/jellyseerr/logs '*'
  clean_logs_dir /var/log/ustreamer '*'
  clean_logs_dir /var/log/gitea '*'
  clean_logs_dir /var/lib/gitea/log '*'
  clean_logs_dir /var/log/fmd '*'
  clean_logs_dir /var/log/uptime-kuma '*'
  clean_logs_dir /opt/uptime-kuma/data/logs '*'
  clean_logs_dir /var/log/romm '*'
  clean_logs_dir /var/log/privatebin '*'
  clean_logs_dir /var/log/crafty '*'
  clean_logs_dir /var/log/rustdesk '*'
  clean_logs_dir /var/log/memos '*'
  clean_logs_dir /var/lib/memos/logs '*'
  clean_logs_dir /var/log/traccar '*'
  clean_logs_dir /var/log/webmin '*'
  clean_logs_dir /var/log/homarr '*'
  clean_logs_dir /var/log/termix '*'
  clean_logs_dir /var/log/syncthing '*'
  clean_logs_dir /var/log/netdata '*'
  clean_logs_dir /var/lib/netdata/dbengine '*'
  clean_logs_dir /var/log/AdGuardHome '*'

  clean_logs_dir /var/log/mysql '*'
  clean_logs_dir /var/log/mariadb '*'
  clean_logs_dir /var/log/postgresql '*'
  truncate_file /var/log/redis/redis-server.log
  clean_logs_dir /var/log/influxdb '*'
  clean_logs_dir /var/log/prometheus '*'
  clean_logs_dir /var/log/grafana '*'
  clean_logs_dir /var/log/loki '*'
  clean_logs_dir /var/log/caddy '*'
  clean_logs_dir /var/log/apache2 '*'
  clean_logs_dir /var/log/lighttpd '*'
  clean_logs_dir /var/log/samba '*'
  clean_logs_dir /var/log/mosquitto '*'
  clean_logs_dir /var/log/openvpn '*'
  clean_logs_dir /var/log/wireguard '*'
  clean_logs_dir /var/log/node-red '*'
  truncate_file /var/log/nodered-install.log
  clean_logs_dir /var/log/transmission-daemon '*'
  clean_logs_dir /var/log/deluge '*'
  clean_logs_dir /var/log/qbittorrent '*'
  clean_logs_dir /var/log/paperless-ngx '*'
  clean_logs_dir /var/log/photoprism '*'
  clean_logs_dir /var/log/navidrome '*'
  clean_logs_dir /var/log/minio '*'
  clean_logs_dir /var/log/nzbget '*'
  clean_logs_dir /var/log/sabnzbd '*'
  clean_logs_dir /var/log/jackett '*'
  clean_logs_dir /var/log/radarr '*'
  clean_logs_dir /var/log/sonarr '*'
  clean_logs_dir /var/log/lidarr '*'
  clean_logs_dir /var/log/prowlarr '*'
  clean_logs_dir /var/log/bazarr '*'
  clean_logs_dir /var/log/overseerr '*'
  clean_logs_dir /var/log/emby-server '*'

  truncate_file /home/homeassistant/.homeassistant/home-assistant.log
  clean_logs_dir /home/homeassistant/.homeassistant/logs '*'
  truncate_file /var/www/nextcloud/data/nextcloud.log
  truncate_file /var/www/owncloud/data/owncloud.log

  if [ -d /var/lib/docker/containers ]; then
    find /var/lib/docker/containers -type f -name '*-json.log' -print0 2>/dev/null | while IFS= read -r -d '' f; do
      :> "$f" && status CLEANED "truncated $f" || status FAIL "truncate $f"
    done
  else
    status SKIP "/var/lib/docker/containers (missing)"
  fi
  clean_file /var/log/wtmp.db
  clean_dir_files /var/tmp/dietpi/logs "*"

  # --- Caches ---
  apt-get clean >/dev/null 2>&1 && status CLEANED "apt cache" || status FAIL "apt cache"
  rm -rf /var/lib/apt/lists/* /var/cache/apt/* 2>/dev/null && status CLEANED "apt lists/cache" || status FAIL "apt lists/cache"
  find /var/cache/debconf -type f -print -delete 2>/dev/null | sed 's/^/[CLEANED] /' || true
  status CLEANED "/var/cache/debconf files"

  # --- Temp directories ---
  truncate_dir /tmp
  truncate_dir /var/tmp

  # --- DietPi RAMlog store (preserve log dir structure) ---
  if [ -d /var/log ]; then
    local log_group="adm"
    if ! getent group "$log_group" >/dev/null 2>&1; then
      log_group="root"
    fi
    install -d -m 0755 -o root -g "$log_group" /var/log/nginx \
      && status CLEANED "/var/log/nginx (ensured)" \
      || status FAIL "/var/log/nginx"
  else
    status SKIP "/var/log (missing)"
  fi
  if [ -x /boot/dietpi/func/dietpi-ramlog ]; then
    /boot/dietpi/func/dietpi-ramlog 1 >/dev/null 2>&1 \
      && status CLEANED "DietPi RAMlog store" \
      || status FAIL "DietPi RAMlog store"
  else
    status SKIP "DietPi RAMlog (missing)"
  fi

  # --- DHCP leases ---
  clean_file /var/lib/dhcp/dhclient.eth0.leases

  # --- Nginx caches ---
  if [ -d /var/lib/nginx ]; then
    find /var/lib/nginx -mindepth 1 -maxdepth 1 -type d -exec rm -rf {} + 2>/dev/null
    status CLEANED "/var/lib/nginx/*"
  else
    status SKIP "/var/lib/nginx"
  fi

  status DONE "Prep complete"
}

check_home_dir() {
  local dir="$1"
  shift
  local keep=("$@")
  if [ -d "$dir" ]; then
    local extra=()
    shopt -s dotglob nullglob
    for entry in "$dir"/*; do
      local base
      base="$(basename "$entry")"
      case "$base" in
        .|..) continue ;;
      esac
      local keep_it=0
      for k in "${keep[@]}"; do
        if [ "$base" = "$k" ]; then
          keep_it=1
          break
        fi
      done
      if [ "$keep_it" -eq 0 ]; then
        extra+=("$base")
      fi
    done
    shopt -u dotglob nullglob
    if [ "${#extra[@]}" -gt 0 ]; then
      status WARN "extra files in $dir: ${extra[*]}"
    else
      status OK "home clean: $dir"
    fi
  else
    status FAIL "missing home dir: $dir"
  fi
}

check_file_missing() {
  local path="$1"
  if [ -e "$path" ]; then
    status FAIL "unexpected file exists: $path"
  else
    status OK "missing as expected: $path"
  fi
}

check_file_empty_or_missing() {
  local path="$1"
  if [ -e "$path" ]; then
    if [ -s "$path" ]; then
      status FAIL "file not empty: $path"
    else
      status OK "empty file: $path"
    fi
  else
    status OK "missing (ok): $path"
  fi
}

check_image() {
  section "Check"

  section "Identity files"
  if [ -e /etc/machine-id ]; then
    if [ -s /etc/machine-id ]; then
      status FAIL "/etc/machine-id not empty"
    else
      status OK "/etc/machine-id empty"
    fi
  else
    status FAIL "/etc/machine-id missing"
  fi
  if [ -L /var/lib/dbus/machine-id ]; then
    status OK "/var/lib/dbus/machine-id symlink present"
  else
    status WARN "dbus machine-id missing or not a symlink"
  fi
  check_file_missing /var/lib/systemd/random-seed

  section "SSH host keys"
  if ls /etc/ssh/ssh_host_* >/dev/null 2>&1; then
    status FAIL "SSH host keys still present"
  else
    status OK "no SSH host keys"
  fi

  section "SSH client traces"
  check_file_empty_or_missing /root/.ssh/known_hosts
  check_file_empty_or_missing /home/dietpi/.ssh/known_hosts
  check_file_empty_or_missing /home/dietpi/.ssh/authorized_keys
  check_file_empty_or_missing /root/.ssh/authorized_keys

  section "Home directories"
  check_home_dir /home/dietpi .bashrc .bash_logout .profile .ssh RESCUE.md
  check_home_dir /root .bashrc .bash_logout .profile .ssh RESCUE.md

  section "Ready flag"
  check_file_missing /var/run/pikit-ready

  section "Firstboot state"
  if [ -d /var/lib/pikit/firstboot ]; then
    local count
    count="$(find /var/lib/pikit/firstboot -type f 2>/dev/null | wc -l | tr -d ' ')"
    if [ "$count" -gt 0 ]; then
      status WARN "firstboot files still present: $count"
    else
      status OK "firstboot dir empty"
    fi
  else
    status OK "firstboot dir missing (ok)"
  fi

  section "TLS certs"
  check_file_missing /etc/pikit/certs/pikit-ca.crt
  check_file_missing /etc/pikit/certs/pikit-ca.key
  check_file_missing /etc/pikit/certs/pikit.local.crt
  check_file_missing /etc/pikit/certs/pikit.local.key
  check_file_missing /var/www/pikit-web/assets/pikit-ca.crt

  section "Logs"
  if [ -d /var/log ]; then
    local nonempty
    nonempty="$(find /var/log -type f -size +0c 2>/dev/null | wc -l | tr -d ' ')"
    if [ "$nonempty" -gt 0 ]; then
      status WARN "/var/log has non-empty files: $nonempty"
    else
      status OK "/var/log empty"
    fi
  else
    status WARN "/var/log missing"
  fi
  if [ -d /var/log/nginx ]; then
    local nginx_nonempty
    nginx_nonempty="$(find /var/log/nginx -type f -size +0c 2>/dev/null | wc -l | tr -d ' ')"
    if [ "$nginx_nonempty" -gt 0 ]; then
      status WARN "/var/log/nginx has non-empty files: $nginx_nonempty"
    else
      status OK "/var/log/nginx empty"
    fi
  else
    status WARN "/var/log/nginx missing"
  fi

  section "DietPi RAM logs"
  if [ -d /var/tmp/dietpi/logs/dietpi-ramlog_store ]; then
    status OK "DietPi RAMlog store present"
  else
    status FAIL "DietPi RAMlog store missing"
  fi

  section "Caches"
  du -sh /var/cache/apt /var/lib/apt/lists /var/cache/debconf 2>/dev/null || true

  section "Temp dirs"
  local tmp_extra
  tmp_extra="$(find /tmp /var/tmp -maxdepth 1 -mindepth 1 ! -name 'systemd-private-*' ! -path '/var/tmp/dietpi' 2>/dev/null | wc -l | tr -d ' ')"
  if [ "$tmp_extra" -gt 0 ]; then
    status WARN "temp dirs not empty"
  else
    status OK "temp dirs empty"
  fi

  section "DHCP lease"
  check_file_missing /var/lib/dhcp/dhclient.eth0.leases

  section "Nginx cache dirs"
  if [ -d /var/lib/nginx ]; then
    local nginx_cache
    nginx_cache="$(find /var/lib/nginx -mindepth 1 -maxdepth 1 -type d 2>/dev/null | wc -l | tr -d ' ')"
    if [ "$nginx_cache" -gt 0 ]; then
      status WARN "nginx cache dirs present: $nginx_cache"
    else
      status OK "/var/lib/nginx empty"
    fi
  else
    status OK "/var/lib/nginx missing (ok)"
  fi
}

finalize() {
  section "Summary"
  status OK "warnings: $WARNINGS"
  status OK "errors: $ERRORS"
  if [ "$ERRORS" -gt 0 ]; then
    echo "[FAIL] Prep/check completed with errors."
    exit 1
  fi
  echo "[OK] Prep/check completed."
}

maybe_self_delete() {
  if [ "$PIKIT_SELF_DELETE" -eq 1 ] && [[ "$SCRIPT_PATH" == /tmp/* ]]; then
    rm -f "$SCRIPT_PATH" || true
  fi
}

main() {
  parse_args "$@"

  if [ "$LOCAL_ONLY" -eq 0 ] && ! is_dietpi; then
    run_remote "$@"
  fi

  require_root

  case "$MODE" in
    prep) prep_image ;;
    check) check_image ;;
    both)
      prep_image
      check_image
      ;;
  esac

  finalize
  maybe_self_delete
}

main "$@"