#!/bin/bash
# preflight.txt — System checks + launcher for preflight138
# Unified log: /var/log/securiti/diagnostics.log

set -euo pipefail

# ---- require root (simplifies logging + avoids sudo TTY weirdness) ----
if [[ $EUID -ne 0 ]]; then
  echo "Please run as root. Try: sudo -i ; then: ./preflight.txt"
  exit 1
fi

# ---- logging setup ----
LOG_DIR="/var/log/securiti"
LOG="$LOG_DIR/diagnostics.log"
TS="$(date +'%Y-%m-%d_%H-%M-%S')"
RETAIN=5
declare -a CREATED_FILES=()
declare -a CREATED_DIRS=()

ensure_logdir() {
  if [[ ! -d "$LOG_DIR" ]]; then
    mkdir -p "$LOG_DIR"
    chmod 755 "$LOG_DIR"
    CREATED_DIRS+=("$LOG_DIR")
  fi
  if [[ ! -e "$LOG" ]]; then
    : > "$LOG"
    chmod 644 "$LOG"
    CREATED_FILES+=("$LOG")
  fi
}

rotate_logs() {
  ensure_logdir

  # rotate current log if non-empty
  if [ -s "$LOG" ]; then
    local rolled="${LOG_DIR}/diagnostics.${TS}.log"
    mv -f "$LOG" "$rolled"
    CREATED_FILES+=("$rolled")
    : > "$LOG"
    chmod 644 "$LOG"
  fi

  # trim older rolled logs (safe even if none exist)
  set +o pipefail
  find "$LOG_DIR" -maxdepth 1 -type f -name 'diagnostics.*.log' -printf '%T@ %p\n' 2>/dev/null \
    | sort -nr \
    | awk -v keep="$RETAIN" 'NR>keep {print $2}' \
    | xargs -r rm -f
  set -o pipefail
}

rotate_logs

# mirror stdout+stderr to console + log (line-buffer if stdbuf present)
if command -v stdbuf >/dev/null 2>&1; then
  exec > >(stdbuf -oL tee -a "$LOG") 2>&1
else
  exec > >(tee -a "$LOG") 2>&1
fi

banner() {
  local phase="$1"
  printf '\n\n%s\n' "======================================================================"
  printf '>>> %s\n' "$phase"
  printf '%s\n\n' "======================================================================"
}

# ---- space requirements ----
declare -A MIN
MIN["/mnt/installation"]="200G"
MIN["/mnt/rancher"]="300G"
MIN["/mnt/securiti-app"]="500G"
MIN["/usr"]="5G"
MIN["/run"]="60G"
MIN["/var"]="30G"
MIN["/etc"]="5G"
MIN["/tmp"]="50G"
ROOT_OPS_MIN="20G"   # baseline for /

to_bytes() {
  local s="${1^^}" num unit
  if [[ "$s" =~ ^([0-9]+)([KMGTP]?)(B)?$ ]]; then
    num="${BASH_REMATCH[1]}"; unit="${BASH_REMATCH[2]}"
  else
    echo "0"; return 1
  fi
  case "$unit" in
    "")  echo "$num" ;;
    K)   echo $(( num * 1024 )) ;;
    M)   echo $(( num * 1024 * 1024 )) ;;
    G)   echo $(( num * 1024 * 1024 * 1024 )) ;;
    T)   echo $(( num * 1024 * 1024 * 1024 * 1024 )) ;;
    P)   echo $(( num * 1024 * 1024 * 1024 * 1024 * 1024 )) ;;
    *)   echo "0"; return 1 ;;
  esac
}

bytes_to_gib() { awk -v b="$1" 'BEGIN { printf "%.1f", (b/1024/1024/1024) }'; }
print_line() { printf "%-18s %-12s %-12s %-10s %-s\n" "$@"; }

check_space() {
  local path="$1" need_h="$2"
  local need_b avail_b status note
  need_b="$(to_bytes "$need_h")" || need_b=0
  avail_b="$(df -B1 --output=avail "$path" 2>/dev/null | tail -n 1 | tr -d '[:space:]')"
  if [[ -z "$avail_b" ]]; then print_line "$path" "$need_h" "--" "WARNING" "path/mount not found"; return; fi
  if ! [[ "$avail_b" =~ ^[0-9]+$ ]]; then print_line "$path" "$need_h" "--" "WARNING" "unreadable avail"; return; fi
  if (( avail_b >= need_b )); then status="GOOD"; note=""; else status="WARNING"; note="insufficient space"; fi
  print_line "$path" "$need_h" "$(bytes_to_gib "$avail_b")GiB" "$status" "$note"
}

check_space_root() {
  local need_h="$1"
  local need_b avail_b status note
  need_b="$(to_bytes "$need_h")" || need_b=0
  avail_b="$(df -B1 --output=avail / 2>/dev/null | tail -n 1 | tr -d '[:space:]')"
  if [[ -z "$avail_b" || ! "$avail_b" =~ ^[0-9]+$ ]]; then print_line "/" "$need_h" "--" "WARNING" "unreadable root avail"; return; fi
  if (( avail_b >= need_b )); then status="GOOD"; note=""; else status="WARNING"; note="insufficient space"; fi
  print_line "/" "$need_h" "$(bytes_to_gib "$avail_b")GiB" "$status" "$note"
}

run_disk_precheck() {
  echo "--- Disk Space Preconditions ---"
  print_line "Path" "Min" "Avail" "Status" "Notes"
  print_line "----" "----" "-----" "------" "-----"

  local ROOT_SUBS=("/usr" "/run" "/var" "/etc" "/tmp")
  local MNT_SUBS=("/mnt/installation" "/mnt/rancher" "/mnt/securiti-app")

  # Check root-impacting subdirs
  for d in "${ROOT_SUBS[@]}"; do
    if [[ -d "$d" ]]; then
      check_space "$d" "${MIN["$d"]}"
    else
      print_line "$d" "${MIN["$d"]}" "--" "INFO" "directory missing (counted toward /)"
    fi
  done

  # /mnt subtree aggregation
  local missing_mnt_bytes=0
  local all_mnt_exist="yes"
  for d in "${MNT_SUBS[@]}"; do
    if [[ -d "$d" ]]; then
      check_space "$d" "${MIN["$d"]}"
    else
      all_mnt_exist="no"
      missing_mnt_bytes=$(( missing_mnt_bytes + $(to_bytes "${MIN["$d"]}") ))
      print_line "$d" "${MIN["$d"]}" "--" "INFO" "directory missing (aggregated)"
    fi
  done

  if [[ "$all_mnt_exist" == "no" ]] && (( missing_mnt_bytes > 0 )); then
    local missing_mnt_h
    missing_mnt_h="$(bytes_to_gib "$missing_mnt_bytes")GiB"
    if [[ -d /mnt ]]; then
      check_space "/mnt" "$missing_mnt_h"
    else
      print_line "/mnt" "$missing_mnt_h" "--" "INFO" "/mnt not present; requirement added to /"
    fi
  fi

  # Root needs: baseline + any missing root subdirs + possibly missing /mnt sum (if /mnt absent)
  local root_need_bytes
  root_need_bytes=$(to_bytes "$ROOT_OPS_MIN")
  for d in "${ROOT_SUBS[@]}"; do
    if [[ ! -d "$d" ]]; then
      root_need_bytes=$(( root_need_bytes + $(to_bytes "${MIN["$d"]}") ))
    fi
  done
  if [[ "$all_mnt_exist" == "no" && ! -d /mnt ]]; then
    root_need_bytes=$(( root_need_bytes + missing_mnt_bytes ))
  fi
  check_space_root "$(bytes_to_gib "$root_need_bytes")GiB"
  echo ""
}

run_os_validation() {
  echo "--- OS Validation ---"
  source /etc/os-release 2>/dev/null || true
  local os_id="${ID:-unknown}" os_ver="${VERSION_ID:-unknown}" pn="${PRETTY_NAME:-$os_id $os_ver}"
  case "$os_id:$os_ver" in
    ubuntu:20.*|ubuntu:22.*|ubuntu:24.*|\
    centos:8*|centos:9*|\
    rhel:8*|rhel:9*|\
    amzn:2|amzn:2023|\
    sles:15*|suse:15*|\
    ol:9*)
      echo "GOOD: Supported OS detected ($pn)"
      ;;
    *)
      echo "WARNING: Unsupported OS ($pn)"
      ;;
  esac
  echo ""
}

run_selinux_check() {
  echo "--- SELinux Check ---"
  if command -v getenforce >/dev/null 2>&1; then
    mode="$(getenforce)"
    if [[ "$mode" == "Enforcing" ]]; then
      echo "INFO: SELINUX enforcing is enabled"
    else
      echo "INFO: SELINUX is not found or disabled (mode=$mode)"
    fi
  else
    if grep -q '^SELINUX=enforcing' /etc/selinux/config 2>/dev/null; then
      echo "INFO: SELINUX enforcing is enabled (config)"
    else
      echo "INFO: SELINUX is not found or disabled"
    fi
  fi
  echo ""
}

run_cis_check() {
  echo "--- CIS Image Check ---"
  if grep -qi "CIS" /etc/issue /etc/motd 2>/dev/null || \
     find /etc -maxdepth 1 -iname "*cis*" 2>/dev/null | grep -q .; then
    echo "Possible ERROR: CIS indicators found. Check CIS is NOT used."
  else
    echo "GOOD: No CIS indicators detected"
  fi
  echo ""
}

run_storage_media_check() {
  echo "--- Storage Media Classification ---"
  printf "%-10s %-6s %-6s %-4s %-10s %-20s %-28s %-6s %-10s %-s\n" \
         "Device" "Type" "Tran" "Rota" "Vendor" "Model" "By-ID" "ID_SSD" "Class" "Reason / Notes"
  
  local warn_hdd=0 info_uncertain=0

  # Build map of first by-id symlink per device (nounset-safe)
  declare -A BYID=()
  if [[ -d /dev/disk/by-id ]]; then
    for link in /dev/disk/by-id/*; do
      [[ -L "$link" ]] || continue
      tgt="$(readlink -f "$link" 2>/dev/null || true)" || true
      [[ -n "$tgt" ]] || continue
      dev="${tgt##*/}"
      # With `set -u`, never expand an unset array element without a default.
      # ${BYID[$dev]+x} expands to non-empty if the key exists, empty otherwise.
      if [[ -z ${BYID[$dev]+x} ]]; then
        BYID[$dev]="${link##*/}"
      fi
    done
  fi

  mapfile -t devs < <(lsblk -dn -e7 -o NAME,TYPE,ROTA,TRAN,MODEL 2>/dev/null | awk '$2=="disk" {print $0}')

  for line in "${devs[@]}"; do
    name=$(awk '{print $1}' <<<"$line")
    type=$(awk '{print $2}' <<<"$line")
    rota=$(awk '{print $3}' <<<"$line")
    tran=$(awk '{print $4}' <<<"$line")
    model=$(cut -d' ' -f5- <<<"$line")

    # prefer kernel flag
    if [[ -r "/sys/block/$name/queue/rotational" ]]; then
      rota_sys=$(cat "/sys/block/$name/queue/rotational" 2>/dev/null || echo "$rota")
      [[ "$rota_sys" =~ ^[01]$ ]] && rota="$rota_sys"
    fi

    vendor="--"
    [[ -r "/sys/block/$name/device/vendor" ]] && vendor="$(tr -d '\t' </sys/block/$name/device/vendor 2>/dev/null | sed 's/  */ /g')"
    [[ -z "$model" && -r "/sys/block/$name/device/model" ]] && model="$(tr -d '\t' </sys/block/$name/device/model 2>/dev/null | sed 's/  */ /g')"

    id_ssd="--"
    if command -v udevadm >/dev/null 2>&1; then
      if props="$(udevadm info --query=property --name=/dev/$name 2>/dev/null)"; then
        val="$(printf "%s\n" "$props" | awk -F= '/^ID_SSD=/{print $2}' | tail -n 1)"
        [[ -n "$val" ]] && id_ssd="$val"
      fi
    fi

    shopt -s nocasematch
    hint=""
    if   [[ $model =~ (amazon|ebs) ]]; then
      hint="AWS EBS"
    elif [[ $model =~ google ]]; then
      hint="GCP PD"
    elif [[ $model =~ (msft|microsoft).*virtual.*disk ]]; then
      hint="Azure Managed Disk"
    elif [[ $model =~ (oracle|block[[:space:]]?volume) ]]; then
      hint="OCI Block Volume"
    elif [[ $tran == "nvme" || $name == nvme* ]]; then
      hint="NVMe"
    elif [[ $tran == "virtio" || $name == vd* || $name == xvd* ]]; then
      hint="virtio (VM)"
    elif [[ $tran == "sata" || $tran == "ata" || $tran == "sas" ]]; then
      hint="SATA/SAS"
    fi
    shopt -u nocasematch

    class="SSD"; reason="rotational=0"
    if   [[ $name == nvme* || $tran == "nvme" ]]; then
      class="SSD"; reason="nvme transport"
    elif [[ $id_ssd == "1" ]]; then
      class="SSD"; reason="udev ID_SSD=1"
    elif [[ $tran == "virtio" || $name == vd* || $name == xvd* ]]; then
      class="UNCERTAIN"; reason="virtio device; verify cloud volume type"; info_uncertain=1
    elif [[ $rota == "1" ]]; then
      if compgen -G "/sys/block/$name/slaves/*" >/dev/null 2>&1; then
        local rot_slave=0
        for s in /sys/block/$name/slaves/*; do
          r=$(cat "$s/queue/rotational" 2>/dev/null || echo 1)
          [[ "$r" == "1" ]] && rot_slave=1
        done
        if (( rot_slave==1 )); then class="HDD"; reason="rotational=1 via slaves"
        else class="SSD"; reason="non-rotational slaves"; fi
      else
        class="HDD"; reason="rotational=1"
      fi
    fi

    byid="${BYID[$name]:---}"
    printf "%-10s %-6s %-6s %-4s %-10s %-20s %-28s %-6s %-10s %-s\n" \
           "$name" "$type" "${tran:---}" "${rota:-?}" "${vendor:---}" "${model:---}" "$byid" "${id_ssd:--}" "$class" "$reason${hint:+; $hint}"

    [[ "$class" == "HDD" ]] && warn_hdd=1
    [[ "$class" == "UNCERTAIN" ]] && info_uncertain=1
  done

  echo ""
  (( warn_hdd )) && echo "WARNING: One or more devices classified as HDD (rotational). SSD/NVMe recommended."
  (( info_uncertain )) && echo "INFO: Some devices appear paravirtual (virtio/xvd). ROTA may be misreported; confirm cloud volume class (gp3/io2/premium)."
  echo ""
}

# ---- launcher helpers ----
LAST_RC=0
run_phase() {
  local title="$1"; shift
  banner "$title"
  echo "Command: $*"
  if "$@"; then LAST_RC=0; else LAST_RC=$?; echo "NOTE: '$title' exited with code $LAST_RC (continuing)"; fi
  echo ""
}

write_summary() {
  echo ""
  echo "============================== SUMMARY =============================="

  ph1="${PH1_RC:-99}"
  ph2="${PH2_RC:-99}"
  ph3="${PH3_RC:-99}"

  if [ "$ph1" -eq 0 ] 2>/dev/null; then s1="PASS"; else s1="FAIL($ph1)"; fi
  if [ "$ph2" -eq 0 ] 2>/dev/null; then s2="PASS"; else s2="FAIL($ph2)"; fi
  if [ "$ph3" -eq 0 ] 2>/dev/null; then s3="PASS"; else s3="FAIL($ph3)"; fi

  printf "Preflight138 --check (initial): %s\n" "$s1"
  printf "Preflight138 --fix            : %s\n" "$s2"
  printf "Preflight138 --check (final)  : %s\n" "$s3"
  echo "----------------------------------------------------------------------"

  warn_cnt=$(grep -i -c '\bWARNING\b' "$LOG" 2>/dev/null || echo 0)
  err_cnt=$(grep -i -c '\bERROR\b' "$LOG" 2>/dev/null || echo 0)
  printf "Total WARNING lines: %s\n" "$warn_cnt"
  printf "Total ERROR   lines: %s\n" "$err_cnt"

  echo "======================================================================"
  echo "Output saved to: $LOG"
}

# Patch preflight138 targets to our LOG (output.txt / diagnostics*.txt -> LOG)
patch_outputs_to_log() {
  local f="$1"
  local LOG_ESC="${LOG//\//\\/}"
  sed -i -E "s/\\boutput\\.txt\\b/${LOG_ESC}/gI" "$f" || true
  sed -i -E "s/\\bdiagnostic(s)?\\.txt\\b/${LOG_ESC}/gI" "$f" || true
}

# ---- main ----
banner "System Information"
echo "--- Disk Usage (df -hT) ---"
df -hT
echo ""
echo "--- Block Devices (lsblk -f) ---"
lsblk -f
echo ""

run_disk_precheck
run_os_validation
run_selinux_check
run_cis_check
run_storage_media_check

banner "preflight138 bootstrap"

# Locate or fetch preflight138, make executable, and patch log outputs
PRE138="./preflight138.sh"
if [[ ! -x "$PRE138" && -f ./preflight138.txt ]]; then
  cp -f ./preflight138.txt ./preflight138.sh
  chmod +x ./preflight138.sh
  patch_outputs_to_log ./preflight138.sh
  CREATED_FILES+=("$(readlink -f ./preflight138.sh)")
fi

if [[ ! -x "$PRE138" ]]; then
  echo "INFO: preflight138.sh not found locally; downloading..."
  curl -fsSL https://www.two14.org/preflight/preflight138.txt -o preflight138.sh
  chmod +x preflight138.sh
  patch_outputs_to_log ./preflight138.sh
  CREATED_FILES+=("$(readlink -f ./preflight138.sh)")
fi

# ---- Interactive Confirmation Helper ----
prompt_fix_confirmation() {
  echo ""
  echo "======================================================================"
  echo "NOTICE: About to run PHASE 2 (--fix)"
  echo "This will alter environment variables and system configuration."
  echo ""
  echo "Do you want to continue? (Yes/No) [Default: No]"
  echo ""
  
  local answer=""
  local timeout_seconds=30
  local elapsed=0
  local pid_bg=""
  
  # Background process to show countdown without interfering with input
  (
    while [ $elapsed -lt $timeout_seconds ]; do
      local remaining=$((timeout_seconds - elapsed))
      
      # Show countdown every 5 seconds
      if [ $((elapsed % 5)) -eq 0 ]; then
        echo "  [${remaining}s remaining...]" >&2
      fi
      
      sleep 1
      elapsed=$((elapsed + 1))
    done
    
    # Timeout reached - send a signal to parent
    echo "  [TIMEOUT]" >&2
  ) &
  pid_bg=$!
  
  # Main prompt (user can type freely)
  echo -n "Enter your choice (Yes/No): "
  
  # Read with full timeout, no interference
  if read -t $timeout_seconds -r answer 2>/dev/null; then
    # User provided input - kill background countdown
    kill $pid_bg 2>/dev/null
    wait $pid_bg 2>/dev/null
    echo ""
  else
    # Timeout occurred
    echo ""
    wait $pid_bg 2>/dev/null
    echo "INFO: No response received within 30 seconds. Defaulting to 'No'."
    answer="No"
  fi
  
  # Normalize answer (case-insensitive, trim whitespace)
  answer=$(echo "$answer" | tr '[:upper:]' '[:lower:]' | xargs)
  
  case "$answer" in
    yes|y)
      echo "User confirmed: Proceeding with --fix"
      return 0
      ;;
    *)
      echo "User declined or timed out: Skipping --fix"
      return 1
      ;;
  esac
}

# ---- PHASES ----
run_phase "PHASE 1: Preflight138 — INSTALL --CHECK" ./preflight138.sh install --check
PH1_RC=$LAST_RC

# Interactive confirmation for --fix
if prompt_fix_confirmation; then
  run_phase "PHASE 2: Preflight138 — INSTALL --FIX"   ./preflight138.sh install --fix
  PH2_RC=$LAST_RC
  
  run_phase "PHASE 3: Preflight138 — INSTALL --CHECK (FINAL)" ./preflight138.sh install --check
  PH3_RC=$LAST_RC
else
  # User declined --fix
  PH2_RC=99  # Mark as skipped
  PH3_RC=99  # Mark as skipped
  echo "" | tee -a "$LOG"
  echo "INFO: --fix was not run (user declined or timeout)" | tee -a "$LOG"
  echo "" | tee -a "$LOG"
fi

banner "Complete"
write_summary

# Exit if --fix was skipped
if [ "${PH2_RC:-99}" -eq 99 ]; then
  echo ""
  echo "Exiting: --fix phase was not executed."
  exit 0
fi

# Cleanup temporary files in current directory
echo ""
echo "======================================================================"
echo "Cleaning up temporary files in current directory..."
if [ -f ./diagnostics.log ]; then
  rm -f ./diagnostics.log
  echo "  Removed: ./diagnostics.log"
fi
if [ -f ./preflight138.sh ]; then
  rm -f ./preflight138.sh
  echo "  Removed: ./preflight138.sh"
fi
echo "Cleanup complete."
echo ""

# Show user where logs are stored
echo "======================================================================"
echo "Log files are stored in: $LOG_DIR"
echo ""
echo "Current log files:"
if [ -d "$LOG_DIR" ]; then
  ls -lh "$LOG_DIR"/diagnostics.* 2>/dev/null || echo "  No diagnostic logs found"
else
  echo "  Directory not found: $LOG_DIR"
fi
echo ""
echo "You can delete these logs manually if desired:"
echo "  rm -rf $LOG_DIR"
echo "======================================================================"