#!/usr/bin/env bash
set -euo pipefail
IFS=$'\n\t'

SCRIPT_NAME="$(basename "$0")"
TASK_INDEX=0
TASK_TOTAL=0
TEMP_FILES=()
STORAGE_DETECTED=0

ROOT_SOURCE=""
VG_NAME=""
LV_NAME=""
PV_NAME=""
LUKS_DEV=""
LUKS_NAME=""
LUKS_PART=""
DISK_DEV=""
PART_NUM=""

if [ -t 1 ]; then
  BOLD="$(tput bold 2>/dev/null || true)"
  DIM="$(tput dim 2>/dev/null || true)"
  RED="$(tput setaf 1 2>/dev/null || true)"
  GREEN="$(tput setaf 2 2>/dev/null || true)"
  YELLOW="$(tput setaf 3 2>/dev/null || true)"
  BLUE="$(tput setaf 4 2>/dev/null || true)"
  RESET="$(tput sgr0 2>/dev/null || true)"
else
  BOLD=""
  DIM=""
  RED=""
  GREEN=""
  YELLOW=""
  BLUE=""
  RESET=""
fi

die() {
  echo "${RED}Error:${RESET} $*" >&2
  exit 1
}

log_info() {
  echo "${BLUE}INFO:${RESET} $*"
}

log_ok() {
  echo "${GREEN}OK:${RESET} $*"
}

log_warn() {
  echo "${YELLOW}WARN:${RESET} $*"
}

section() {
  local title="$1"
  echo
  echo "${BOLD}${title}${RESET}"
  echo "${DIM}------------------------------------------------------------${RESET}"
}

add_temp_file() {
  TEMP_FILES+=("$1")
}

cleanup() {
  local file
  for file in "${TEMP_FILES[@]:-}"; do
    [ -f "$file" ] && rm -f "$file"
  done
}
trap cleanup EXIT

ensure_tty() {
  if [ ! -t 0 ] || [ ! -t 1 ]; then
    die "This setup must run interactively in a TTY."
  fi
}

ensure_root() {
  if [ "${EUID:-$(id -u)}" -ne 0 ]; then
    if command -v sudo >/dev/null 2>&1; then
      log_info "Re-running with sudo..."
      exec sudo -E "$0" "$@"
    fi
    die "Must be root or have sudo privileges to run this setup."
  fi
}

require_cmd() {
  local cmd="$1"
  if ! command -v "$cmd" >/dev/null 2>&1; then
    log_warn "Missing command: $cmd"
    return 1
  fi
  return 0
}

prompt_input() {
  local label="$1"
  local default="${2:-}"
  local value=""
  if [ -n "$default" ]; then
    read -r -p "${label} [${default}]: " value </dev/tty
    value="${value:-$default}"
  else
    read -r -p "${label}: " value </dev/tty
  fi
  printf '%s' "$value"
}

prompt_secret() {
  local label="$1"
  local value=""
  while true; do
    read -r -s -p "${label}: " value </dev/tty
    printf '\n' >/dev/tty
    if [ -n "$value" ]; then
      printf '%s' "$value"
      return 0
    fi
    log_warn "Value cannot be empty."
  done
}

prompt_secret_confirm() {
  local label="$1"
  local confirm_label="$2"
  local a=""
  local b=""
  while true; do
    read -r -s -p "${label}: " a </dev/tty
    printf '\n' >/dev/tty
    read -r -s -p "${confirm_label}: " b </dev/tty
    printf '\n' >/dev/tty
    if [ -z "$a" ]; then
      log_warn "Value cannot be empty."
      continue
    fi
    if [ "$a" != "$b" ]; then
      log_warn "Values do not match. Please try again."
      continue
    fi
    printf '%s' "$a"
    return 0
  done
}

confirm() {
  local label="$1"
  local default="${2:-yes}"
  local prompt=""
  local answer=""

  if [ "$default" = "yes" ]; then
    prompt="[Y/n]"
  else
    prompt="[y/N]"
  fi

  while true; do
    read -r -p "${label} ${prompt} " answer </dev/tty
    if [ -z "$answer" ]; then
      answer="$default"
    fi
    case "${answer,,}" in
      y|yes) return 0 ;;
      n|no) return 1 ;;
      *) log_warn "Please answer yes or no." ;;
    esac
  done
}

human_bytes() {
  local bytes="$1"
  if command -v numfmt >/dev/null 2>&1; then
    numfmt --to=iec --suffix=B "$bytes"
  else
    awk -v b="$bytes" 'BEGIN {
      split("B KiB MiB GiB TiB", u, " ");
      i=1;
      while (b>=1024 && i<5) { b/=1024; i++; }
      printf "%.1f %s", b, u[i];
    }'
  fi
}

mib_to_human() {
  local mib="$1"
  awk -v m="$mib" 'BEGIN { printf "%.1f GiB", m/1024 }'
}

add_task() {
  TASK_TITLES+=("$1")
  TASK_FUNCS+=("$2")
  TASK_TOTAL=$((TASK_TOTAL + 1))
}

run_tasks() {
  local i
  local title
  local func
  for i in "${!TASK_FUNCS[@]}"; do
    TASK_INDEX=$((TASK_INDEX + 1))
    title="${TASK_TITLES[$i]}"
    func="${TASK_FUNCS[$i]}"
    section "Task ${TASK_INDEX}/${TASK_TOTAL}: ${title}"
    "$func"
  done
}

lsblk_attr() {
  local path="$1"
  local attr="$2"
  lsblk -dn -o "$attr" "$path" 2>/dev/null | head -n1 | xargs || true
}

infer_disk_part_from_partition() {
  local part_path="$1"
  local resolved
  local base

  resolved="$(readlink -f "$part_path" 2>/dev/null || printf '%s' "$part_path")"
  base="$(basename "$resolved")"

  if [[ "$base" =~ ^(nvme[0-9]+n[0-9]+)p([0-9]+)$ ]]; then
    printf '/dev/%s\n%s\n' "${BASH_REMATCH[1]}" "${BASH_REMATCH[2]}"
    return 0
  fi
  if [[ "$base" =~ ^(mmcblk[0-9]+)p([0-9]+)$ ]]; then
    printf '/dev/%s\n%s\n' "${BASH_REMATCH[1]}" "${BASH_REMATCH[2]}"
    return 0
  fi
  if [[ "$base" =~ ^(md[0-9]+)p([0-9]+)$ ]]; then
    printf '/dev/%s\n%s\n' "${BASH_REMATCH[1]}" "${BASH_REMATCH[2]}"
    return 0
  fi
  if [[ "$base" =~ ^((sd|vd|xvd|hd)[a-z]+)([0-9]+)$ ]]; then
    printf '/dev/%s\n%s\n' "${BASH_REMATCH[1]}" "${BASH_REMATCH[3]}"
    return 0
  fi
  return 1
}

resolve_luks_backing_partition() {
  local mapper_path="$1"
  local mapper_name="$2"
  local mapper_name_alt="$3"
  local part=""

  if command -v cryptsetup >/dev/null 2>&1; then
    part="$(cryptsetup status "$mapper_name" 2>/dev/null | awk '/^[[:space:]]*device:/ {print $2; exit}' | xargs || true)"
    if [ -z "$part" ] && [ -n "$mapper_name_alt" ] && [ "$mapper_name_alt" != "$mapper_name" ]; then
      part="$(cryptsetup status "$mapper_name_alt" 2>/dev/null | awk '/^[[:space:]]*device:/ {print $2; exit}' | xargs || true)"
    fi
    if [ -z "$part" ]; then
      part="$(cryptsetup status "$mapper_path" 2>/dev/null | awk '/^[[:space:]]*device:/ {print $2; exit}' | xargs || true)"
    fi
  fi

  if [ -z "$part" ]; then
    part="$(lsblk -nro PATH,TYPE -s "$mapper_path" 2>/dev/null | awk '$2=="part" {print $1; exit}' | xargs || true)"
  fi

  printf '%s' "$part"
}

refresh_partition_table() {
  local disk="$1"
  if command -v partprobe >/dev/null 2>&1; then
    partprobe "$disk" || true
  fi
  if command -v partx >/dev/null 2>&1; then
    partx -u "$disk" || true
  fi
  if command -v udevadm >/dev/null 2>&1; then
    udevadm settle || true
  fi
}

wait_for_partition_growth() {
  local part="$1"
  local old_bytes="$2"
  local new_bytes=0
  local i=0

  for ((i = 0; i < 12; i++)); do
    new_bytes="$(blockdev --getsize64 "$part" 2>/dev/null || echo 0)"
    if [ "$new_bytes" -gt "$old_bytes" ]; then
      return 0
    fi
    sleep 1
  done
  return 1
}

is_last_partition_on_disk() {
  local disk="$1"
  local part_num="$2"
  local last_part

  last_part="$(parted -ms "$disk" unit s print 2>/dev/null | awk -F: '$1 ~ /^[0-9]+$/ {last=$1} END {print last}')"
  [ -n "$last_part" ] && [ "$part_num" = "$last_part" ]
}

get_trailing_free_bytes() {
  local disk="$1"
  local part_num="$2"
  local disk_bytes
  local part_end_bytes
  local free_bytes

  disk_bytes="$(blockdev --getsize64 "$disk" 2>/dev/null || true)"
  part_end_bytes="$(parted -ms "$disk" unit B print 2>/dev/null | awk -F: -v p="$part_num" '$1==p {gsub("B","",$3); print $3; exit}')"
  if [ -z "$disk_bytes" ] || [ -z "$part_end_bytes" ]; then
    return 1
  fi

  free_bytes=$((disk_bytes - part_end_bytes - 1))
  if [ "$free_bytes" -lt 0 ]; then
    free_bytes=0
  fi
  echo "$free_bytes"
}

resize_open_luks_mapping() {
  if [ -n "$LUKS_NAME" ] && cryptsetup resize "$LUKS_NAME" >/dev/null 2>&1; then
    return 0
  fi
  cryptsetup resize "$LUKS_DEV" >/dev/null 2>&1
}

detect_storage_stack() {
  STORAGE_DETECTED=0
  LUKS_NAME=""
  LUKS_PART=""
  DISK_DEV=""
  PART_NUM=""

  require_cmd findmnt || return 1
  require_cmd lvs || return 1
  require_cmd pvs || return 1
  require_cmd lsblk || return 1
  require_cmd cryptsetup || return 1

  ROOT_SOURCE="$(findmnt -n -o SOURCE / || true)"
  if [ -z "$ROOT_SOURCE" ] || [ ! -e "$ROOT_SOURCE" ]; then
    log_warn "Unable to detect root device."
    return 1
  fi

  VG_NAME="$(lvs --noheadings -o vg_name "$ROOT_SOURCE" 2>/dev/null | xargs || true)"
  LV_NAME="$(lvs --noheadings -o lv_name "$ROOT_SOURCE" 2>/dev/null | xargs || true)"
  if [ -z "$VG_NAME" ]; then
    log_warn "Root does not appear to be on LVM."
    return 1
  fi

  PV_NAME="$(pvs --noheadings -o pv_name --select "vg_name=${VG_NAME}" 2>/dev/null | head -n1 | xargs || true)"
  if [ -z "$PV_NAME" ]; then
    log_warn "Unable to detect LVM physical volume."
    return 1
  fi

  LUKS_DEV="$PV_NAME"
  LUKS_NAME="$(basename "$LUKS_DEV")"
  local mapper_name_alt
  mapper_name_alt="$(lsblk_attr "$LUKS_DEV" NAME)"

  local luks_type
  luks_type="$(lsblk_attr "$LUKS_DEV" TYPE)"
  if [ "$luks_type" != "crypt" ]; then
    log_warn "LVM PV is not on a LUKS device (type: ${luks_type:-unknown}). LUKS resize will be skipped."
    return 1
  fi

  LUKS_PART="$(resolve_luks_backing_partition "$LUKS_DEV" "$LUKS_NAME" "$mapper_name_alt")"
  if [ -z "$LUKS_PART" ] || [ ! -b "$LUKS_PART" ]; then
    log_warn "Unable to detect LUKS backing partition."
    return 1
  fi

  local inferred
  local inferred_disk=""
  local inferred_part=""
  inferred="$(infer_disk_part_from_partition "$LUKS_PART" || true)"
  if [ -n "$inferred" ]; then
    inferred_disk="$(printf '%s\n' "$inferred" | sed -n '1p')"
    inferred_part="$(printf '%s\n' "$inferred" | sed -n '2p')"
  fi

  DISK_DEV="$inferred_disk"
  PART_NUM="$inferred_part"

  if [ -z "$DISK_DEV" ]; then
    local disk_parent
    disk_parent="$(lsblk_attr "$LUKS_PART" PKNAME)"
    if [ -n "$disk_parent" ]; then
      DISK_DEV="/dev/${disk_parent}"
    else
      DISK_DEV="$(lsblk -nro PATH,TYPE -s "$LUKS_PART" 2>/dev/null | awk '$2=="disk" {print $1; exit}' | xargs || true)"
    fi
  fi

  if [ -z "$PART_NUM" ]; then
    PART_NUM="$(lsblk_attr "$LUKS_PART" PARTNUM)"
  fi

  if [ -z "$DISK_DEV" ] || [ ! -b "$DISK_DEV" ] || [ -z "$PART_NUM" ] || ! [[ "$PART_NUM" =~ ^[0-9]+$ ]]; then
    log_warn "Unable to detect disk device or partition number."
    log_warn "Detected values: LUKS_PART=${LUKS_PART:-<empty>} DISK_DEV=${DISK_DEV:-<empty>} PART_NUM=${PART_NUM:-<empty>}"
    return 1
  fi

  STORAGE_DETECTED=1
  return 0
}

resize_lvm_on_luks() {
  if [ "$STORAGE_DETECTED" -ne 1 ]; then
    log_warn "Storage layout not detected; skipping resize."
    return 0
  fi

  require_cmd parted || return 0
  require_cmd blockdev || return 0
  require_cmd cryptsetup || return 0
  require_cmd pvresize || return 0
  require_cmd lvextend || return 0

  log_info "Root LV: ${ROOT_SOURCE}"
  log_info "VG: ${VG_NAME} | LV: ${LV_NAME}"
  log_info "LUKS device: ${LUKS_DEV}"
  log_info "LUKS partition: ${LUKS_PART}"
  log_info "Disk: ${DISK_DEV} | Partition number: ${PART_NUM}"
  [ -n "$LUKS_NAME" ] && log_info "LUKS mapper name: ${LUKS_NAME}"

  if ! is_last_partition_on_disk "$DISK_DEV" "$PART_NUM"; then
    log_warn "Partition ${PART_NUM} is not the last partition on ${DISK_DEV}. Automatic growth is skipped."
    return 0
  fi

  local part_size_before
  local free_bytes

  part_size_before="$(blockdev --getsize64 "$LUKS_PART" 2>/dev/null || echo 0)"
  free_bytes="$(get_trailing_free_bytes "$DISK_DEV" "$PART_NUM" || true)"
  if [ -z "$free_bytes" ]; then
    log_warn "Unable to determine free disk space."
    return 0
  fi

  if [ "$free_bytes" -lt $((1024 * 1024)) ]; then
    log_ok "No significant free space detected after the root partition."
    return 0
  fi

  log_info "Unallocated space available: $(human_bytes "$free_bytes")"
  if ! confirm "Extend partition, LUKS device, VG, and root LV now?" "yes"; then
    log_warn "Skipped resize."
    return 0
  fi

  log_info "Extending partition ${LUKS_PART} to 100% of disk..."
  if ! parted -s "$DISK_DEV" resizepart "$PART_NUM" 100%; then
    log_warn "Partition resize failed."
    return 0
  fi

  refresh_partition_table "$DISK_DEV"
  if [ "$part_size_before" -gt 0 ] && ! wait_for_partition_growth "$LUKS_PART" "$part_size_before"; then
    log_warn "Kernel has not reported the new partition size yet; continuing anyway."
  fi

  log_info "Resizing LUKS device..."
  if ! resize_open_luks_mapping; then
    log_warn "LUKS resize failed."
    return 0
  fi

  log_info "Resizing LVM physical volume..."
  if ! pvresize "$LUKS_DEV"; then
    log_warn "pvresize failed."
    return 0
  fi

  log_info "Extending root LV to use all free space..."
  if ! lvextend -l +100%FREE -r "$ROOT_SOURCE"; then
    log_warn "lvextend failed."
    return 0
  fi

  log_ok "Resize complete."
}

parse_size_to_mib() {
  local input="$1"
  local normalized
  normalized="$(echo "$input" | tr '[:upper:]' '[:lower:]' | xargs)"
  if [[ "$normalized" =~ ^[0-9]+$ ]]; then
    echo "$normalized"
    return 0
  fi
  if [[ "$normalized" =~ ^([0-9]+)(g|gb)$ ]]; then
    echo $((BASH_REMATCH[1] * 1024))
    return 0
  fi
  if [[ "$normalized" =~ ^([0-9]+)(m|mb)$ ]]; then
    echo "${BASH_REMATCH[1]}"
    return 0
  fi
  return 1
}

setup_swap() {
  require_cmd awk || return 0
  require_cmd fallocate || return 0
  require_cmd mkswap || return 0
  require_cmd swapon || return 0

  local mem_kib
  local mem_mib
  local swap_mib

  mem_kib="$(awk '/MemTotal/ {print $2}' /proc/meminfo)"
  mem_mib=$((mem_kib / 1024))

  if [ "$mem_mib" -lt 2048 ]; then
    swap_mib=$((mem_mib * 2))
  elif [ "$mem_mib" -lt 4096 ]; then
    swap_mib=$mem_mib
  else
    swap_mib=$((mem_mib / 5))
  fi

  log_info "Detected RAM: $(mib_to_human "$mem_mib")"
  log_info "Recommended swap size: $(mib_to_human "$swap_mib")"

  if ! confirm "Use recommended swap size?" "yes"; then
    local custom
    while true; do
      custom="$(prompt_input "Enter custom swap size (MiB or GiB, e.g. 2048 or 4G)")"
      if swap_mib="$(parse_size_to_mib "$custom")"; then
        if [ "$swap_mib" -gt 0 ]; then
          break
        fi
      fi
      log_warn "Invalid size. Try again."
    done
  fi

  local swapfile="/swapfile"
  if [ -e "$swapfile" ]; then
    log_warn "Swap file already exists at ${swapfile}."
    if ! confirm "Replace existing swap file?" "no"; then
      log_warn "Skipped swap setup."
      return 0
    fi
    if swapon --show=NAME --noheadings | grep -qx "$swapfile"; then
      swapoff "$swapfile" || true
    fi
    rm -f "$swapfile"
  fi

  log_info "Creating swap file (${swap_mib} MiB) at ${swapfile}..."
  fallocate -l "${swap_mib}M" "$swapfile"
  chmod 600 "$swapfile"
  mkswap "$swapfile" >/dev/null
  swapon "$swapfile"

  if ! grep -qE "^[[:space:]]*${swapfile}[[:space:]]" /etc/fstab; then
    echo "${swapfile} none swap sw 0 0" >> /etc/fstab
  fi

  log_ok "Swap enabled."
}

change_luks_passphrase() {
  if [ "$STORAGE_DETECTED" -ne 1 ]; then
    log_warn "Storage layout not detected; skipping LUKS passphrase change."
    return 0
  fi
  if [ -z "$LUKS_PART" ] || [ ! -e "$LUKS_PART" ]; then
    log_warn "LUKS partition not found; skipping."
    return 0
  fi
  require_cmd cryptsetup || return 0

  if ! confirm "Change LUKS passphrase in slot 0 now?" "yes"; then
    log_warn "Skipped LUKS passphrase change."
    return 0
  fi

  local old_pass
  local new_pass
  local tmp_old
  local tmp_new

  old_pass="$(prompt_secret "Enter current LUKS passphrase")"
  new_pass="$(prompt_secret_confirm "Enter new LUKS passphrase" "Confirm new LUKS passphrase")"

  tmp_old="$(mktemp)"
  tmp_new="$(mktemp)"
  add_temp_file "$tmp_old"
  add_temp_file "$tmp_new"

  printf '%s' "$old_pass" >"$tmp_old"
  printf '%s' "$new_pass" >"$tmp_new"

  log_info "Updating LUKS passphrase in slot 0..."
  if cryptsetup luksChangeKey --batch-mode --key-slot 0 --key-file "$tmp_old" "$LUKS_PART" "$tmp_new"; then
    log_ok "LUKS passphrase updated."
  else
    log_warn "Failed to update LUKS passphrase."
  fi
}

setup_clevis() {
  log_info "Clevis/Tang setup is not implemented in this template yet."
  if confirm "Would you like to configure Clevis with a Tang server now? (will be skipped)" "no"; then
    local tang
    tang="$(prompt_input "Tang server URL" "http://tang.int.r3w.de")"
    log_warn "Clevis setup for ${tang} is not implemented yet. Skipping."
  else
    log_info "Skipping Clevis setup."
  fi
}

setup_tailscale() {
  if ! require_cmd tailscale; then
    log_warn "Tailscale is not installed; skipping."
    return 0
  fi

  if ! confirm "Set up Tailscale now?" "yes"; then
    log_warn "Skipped Tailscale setup."
    return 0
  fi

  local server
  local tags
  local key
  local tag_list=""
  local t

  server="$(prompt_input "Tailscale/Headscale server URL" "https://vpn.s1q.dev")"
  tags="$(prompt_input "Client tags (comma-separated)" "server")"
  key="$(prompt_secret "Pre-authentication key")"

  if [ -n "$tags" ]; then
    IFS=',' read -r -a tag_array <<<"$tags"
    for t in "${tag_array[@]:-}"; do
      t="$(echo "$t" | xargs)"
      if [ -n "$t" ]; then
        if [ -n "$tag_list" ]; then
          tag_list+=","
        fi
        tag_list+="tag:${t}"
      fi
    done
  fi

  log_info "Bringing up Tailscale..."
  if [ -n "$tag_list" ]; then
    tailscale up --login-server "$server" --authkey "$key" --ssh --advertise-tags "$tag_list"
  else
    tailscale up --login-server "$server" --authkey "$key" --ssh
  fi

  log_ok "Tailscale setup complete."
}

setup_crowdsec() {
  if ! require_cmd cscli; then
    log_warn "CrowdSec (cscli) not installed; skipping."
    return 0
  fi

  if ! confirm "Set up CrowdSec now?" "yes"; then
    log_warn "Skipped CrowdSec setup."
    return 0
  fi

  local key
  key="$(prompt_secret "Enrollment key")"

  log_info "Enrolling CrowdSec..."
  if cscli console enroll -e "$key"; then
    log_info "Restarting CrowdSec service..."
    systemctl restart crowdsec || log_warn "Failed to restart crowdsec service."
    log_ok "CrowdSec enrollment complete."
  else
    log_warn "CrowdSec enrollment failed."
  fi
}

prompt_reboot() {
  log_warn "A reboot is strongly recommended after partition or swap changes."
  if confirm "Reboot now?" "yes"; then
    log_info "Rebooting..."
    reboot
  else
    log_info "Please reboot later to ensure changes take effect."
  fi
}

welcome() {
  section "Initial VM Setup"
  log_info "This setup runs once and is fully interactive."
  log_info "Hostname: $(hostname)"
}

main() {
  ensure_tty
  ensure_root "$@"
  welcome

  if detect_storage_stack; then
    log_ok "Detected LVM on LUKS storage layout."
  else
    log_warn "Storage layout detection incomplete; some steps may be skipped."
  fi

  TASK_TITLES=()
  TASK_FUNCS=()

  add_task "Resize LVM on LUKS (if free space exists)" resize_lvm_on_luks
  add_task "Configure swap file" setup_swap
  add_task "Change LUKS passphrase (slot 0)" change_luks_passphrase
  add_task "Clevis/Tang setup (placeholder)" setup_clevis
  add_task "Configure Tailscale" setup_tailscale
  add_task "Configure CrowdSec" setup_crowdsec
  add_task "Reboot recommendation" prompt_reboot

  run_tasks
  log_ok "Initial setup finished."
}

main "$@"