#!/bin/sh
# Clawdie Shell — ZFS Pool Detection Module
# Purpose: Detect existing clawdie pool, present boot mode menu
# POSIX-compliant (no bash-isms)
#
# Sets CLAWDIE_BOOT_MODE to one of:
#   install     — fresh install (destroy or no existing pool)
#   upgrade     — upgrade existing system in-place
#   maintenance — hand off to maintenance-mode.sh (never returns)

set -eu

POOL_NAME="${POOL_NAME:-clawdie}"
SHARE="${SHARE:-/usr/local/share/clawdie-iso}"
LOG_FILE="${LOG_FILE:-/var/log/clawdie-firstboot.log}"

# ============================================================================
# MAIN ENTRY POINT
# ============================================================================

clawdie_shell_zfs_detect() {
  log_msg "[zfs] Detecting existing pools"

  kldload zfs 2>/dev/null || true

  _existing_pool=""
  _pool_info=""

  # Check if a clawdie pool is importable (not yet imported)
  _importable=$(zpool import 2>/dev/null | awk '/pool:/ { print $2 }' || true)
  for _p in $_importable; do
    if [ "$_p" = "$POOL_NAME" ]; then
      _existing_pool="$_p"
      break
    fi
  done

  # Also check if already imported (e.g. booted from HDD)
  if [ -z "$_existing_pool" ] && zpool list "$POOL_NAME" >/dev/null 2>&1; then
    _existing_pool="$POOL_NAME"
  fi

  if [ -n "$_existing_pool" ]; then
    log_msg "[zfs] Found existing pool: $_existing_pool"
    _pool_info=$(_zfs_get_pool_info "$_existing_pool")
    _zfs_handle_existing_pool "$_existing_pool" "$_pool_info"
  else
    log_msg "[zfs] No existing pool named ${POOL_NAME} found"
    _zfs_handle_fresh_target
  fi

  log_msg "[zfs] Boot mode: ${CLAWDIE_BOOT_MODE:-install}"
  export CLAWDIE_BOOT_MODE
}

_zfs_handle_existing_pool() {
  _pool="$1"
  _info="$2"

  case "${CLAWDIE_BOOT_MODE_PRESET:-0}:${CLAWDIE_BOOT_MODE:-install}" in
    1:upgrade)
      log_msg "[zfs] setup-import preset upgrade mode"
      _zfs_prepare_upgrade "$_pool"
      ;;
    1:maintenance)
      log_msg "[zfs] setup-import preset maintenance mode"
      exec "${SHARE}/firstboot/maintenance-mode.sh"
      ;;
    1:install)
      log_msg "[zfs] setup-import requested fresh install but pool ${_pool} exists — showing menu for safety"
      _zfs_show_existing_menu "$_pool" "$_info"
      ;;
    *)
      _zfs_show_existing_menu "$_pool" "$_info"
      ;;
  esac
}

_zfs_handle_fresh_target() {
  case "${CLAWDIE_BOOT_MODE_PRESET:-0}:${CLAWDIE_BOOT_MODE:-install}" in
    1:install)
      log_msg "[zfs] setup-import preset fresh install"
      CLAWDIE_BOOT_MODE="install"
      ;;
    1:upgrade|1:maintenance)
      log_msg "[zfs] setup-import requested ${CLAWDIE_BOOT_MODE} but no matching pool was found — falling back to fresh-install menu"
      _zfs_show_fresh_menu
      ;;
    *)
      _zfs_show_fresh_menu
      ;;
  esac
}

# ============================================================================
# POOL INFO GATHERING
# ============================================================================

_zfs_get_pool_info() {
  _pool="$1"
  _info=""
  _imported=0

  # Temporarily import if not already
  if ! zpool list "$_pool" >/dev/null 2>&1; then
    zpool import -o altroot=/tmp/pool-inspect "$_pool" 2>/dev/null || true
    _imported=1
  fi

  _size=$(zpool list -Hp -o size "$_pool" 2>/dev/null || echo "0")
  _used=$(zpool list -Hp -o allocated "$_pool" 2>/dev/null || echo "0")
  _size_gb=$((_size / 1073741824))
  _used_gb=$((_used / 1073741824))

  _state=$(zpool status "$_pool" 2>/dev/null | awk '/state:/ { print $2; exit }')
  _disk_count=$(zpool status "$_pool" 2>/dev/null | grep -cE '^\s+(ada|da|nvd|nda)' || echo "0")

  _info="Pool: ${_pool}  State: ${_state}
Disks: ${_disk_count}  Used: ${_used_gb} GB / ${_size_gb} GB"

  # Check for existing clawdie-ai install
  _altroot="/tmp/pool-inspect"
  if [ "$_imported" -eq 0 ]; then
    _altroot=""
  fi
  _ai_env="${_altroot}/home/clawdie/clawdie-ai/.env"
  _ai_pkg="${_altroot}/home/clawdie/clawdie-ai/package.json"
  if [ -f "$_ai_env" ] || [ -f "$_ai_pkg" ]; then
    _info="${_info}
Clawdie-AI: detected"
  fi

  # Export if we imported it
  if [ "$_imported" -eq 1 ]; then
    zpool export "$_pool" 2>/dev/null || true
  fi

  echo "$_info"
}

# ============================================================================
# BOOT MENUS
# ============================================================================

_zfs_show_existing_menu() {
  _pool="$1"
  _info="$2"

  _choice=$(_dialog --menu "\
Existing Clawdie Pool Detected

${_info}

Select action:" 18 70 4 \
    "install"     "Fresh install (destroy existing pool)" \
    "upgrade"     "Upgrade existing system" \
    "maintenance" "Maintenance mode (repair, migrate)" \
    "shell"       "Drop to shell") || _choice="install"

  case "$_choice" in
    install)
      if _dialog --yesno "\
WARNING: This will DESTROY all data on pool '${_pool}'.

Are you absolutely sure?" 10 60; then
        CLAWDIE_BOOT_MODE="install"
        # Clear stale ZFS labels from every device that was in this pool.
        # Without this, the boot loader finds old pool GUIDs on the raw devices
        # after the new install writes fresh labels — same failure mode that
        # made PC-BSD upgrades unreliable. Do this before bsdinstall rewrites
        # the pool so the new labels are unambiguous.
        log_msg "[zfs] Clearing stale ZFS labels before fresh install"
        _zfs_labelclear_pool_devices "$_pool"
      else
        # User cancelled — re-show menu
        _zfs_show_existing_menu "$_pool" "$_info"
        return
      fi
      ;;
    upgrade)
      _zfs_prepare_upgrade "$_pool"
      ;;
    maintenance)
      CLAWDIE_BOOT_MODE="maintenance"
      log_msg "[zfs] Handing off to maintenance mode"
      exec "${SHARE}/firstboot/maintenance-mode.sh"
      # exec never returns
      ;;
    shell)
      log_msg "[zfs] Dropping to shell"
      exec /bin/sh
      ;;
    *)
      CLAWDIE_BOOT_MODE="install"
      ;;
  esac
}

_zfs_prepare_upgrade() {
  _pool="$1"
  CLAWDIE_BOOT_MODE="upgrade"
  _snap_ts=$(date +%s)
  log_msg "[zfs] Taking pre-upgrade snapshot @pre-upgrade-${_snap_ts}"
  if ! zpool list "$_pool" >/dev/null 2>&1; then
    zpool import "$_pool" 2>/dev/null || true
  fi
  if zfs snapshot -r "${_pool}@pre-upgrade-${_snap_ts}" 2>/dev/null; then
    log_msg "[zfs] Snapshot created: ${_pool}@pre-upgrade-${_snap_ts}"
    export UPGRADE_SNAPSHOT="${_pool}@pre-upgrade-${_snap_ts}"
  else
    log_msg "[zfs] WARNING: pre-upgrade snapshot failed — continuing anyway"
  fi
}

_zfs_show_fresh_menu() {
  _choice=$(_dialog --menu "\
No Existing Clawdie Pool

This USB will create a new ZFS pool and install Clawdie-AI.

Select action:" 14 60 3 \
    "install" "Fresh install (create new pool)" \
    "shell"   "Drop to shell" \
    "reboot"  "Reboot") || _choice="install"

  case "$_choice" in
    install)
      CLAWDIE_BOOT_MODE="install"
      ;;
    shell)
      exec /bin/sh
      ;;
    reboot)
      reboot
      ;;
    *)
      CLAWDIE_BOOT_MODE="install"
      ;;
  esac
}

# ============================================================================
# ZFS LABEL CLEANUP
# ============================================================================

_zfs_labelclear_pool_devices() {
  # Clear ZFS labels from every device/partition that was part of the named pool.
  # This prevents "can't find pool by GUID" boot failures after a fresh install
  # overwrites a disk that previously held a ZFS pool — the root cause of most
  # PC-BSD upgrade/reinstall instability.
  _pool="$1"
  log_msg "[zfs] Clearing ZFS labels from pool '${_pool}' devices"
  zpool status "$_pool" 2>/dev/null \
    | awk '/^\s+(ada|da|nvd|nda|vtbd|sd|cd)[0-9]/ { print $1 }' \
    | while read -r _dev; do
        zpool labelclear -f "/dev/$_dev" 2>/dev/null \
          && log_msg "[zfs] Cleared label: /dev/$_dev" \
          || log_msg "[zfs] Note: labelclear /dev/$_dev (may already be clear)"
      done
  log_msg "[zfs] ZFS label cleanup complete"
}

# ============================================================================
# LOGGING HELPER (may be overridden by firstboot.sh)
# ============================================================================

if ! command -v log_msg >/dev/null 2>&1; then
  log_msg() { echo "$(date '+%H:%M:%S') $1" | tee -a "$LOG_FILE" 2>/dev/null || true; }
fi

# Only run if executed directly (not sourced by firstboot.sh)
if [ "${SHELL_ZFS_TEST:-0}" -eq 0 ] && [ "$(basename "$0")" = "shell-zfs.sh" ]; then
  _dialog() { bsddialog --backtitle "Clawdie-AI Setup" "$@" 2>&1; }
  clawdie_shell_zfs_detect
fi
