#!/usr/bin/env bash
# BigBrother Operator Setup Wizard
#
# Single interactive script for prepping Orange Pi Zero 3 (or compatible Debian) SD cards.
# Operator runs this on their machine with an SD card inserted.
#
# What this does:
#   1. Detects block devices and has operator pick SD card target
#   2. Downloads/caches Armbian minimal image for Orange Pi Zero 3
#   3. Flashes image to card with verification
#   4. Mounts card and collects configuration (WiFi, SSH, VPN, alerter, etc.)
#   5. Applies all config to card in a single batch after operator confirms summary
#   6. Sets up first-boot systemd service for autonomous installation
#   7. Unmounts card — ready to insert and power on
#
# Style: bash only, ANSI colors, numbered menus, c2itall-style (no preamble, direct action)
# Root required for: mount, unmount, dd, partprobe, rsync to /mnt, chown, chmod on mounted fs
#
set -euo pipefail

# ══════════════════════════════════════════════════════════════════════════════
# COLORS (bash ANSI only)
# ══════════════════════════════════════════════════════════════════════════════
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
BOLD='\033[1m'
NC='\033[0m'

# ══════════════════════════════════════════════════════════════════════════════
# HELPERS
# ══════════════════════════════════════════════════════════════════════════════
info()  { echo -e "${GREEN}[+]${NC} $*"; }
warn()  { echo -e "${YELLOW}[!]${NC} $*"; }
error() { echo -e "${RED}[-]${NC} $*"; exit 1; }
step()  { echo -e "${CYAN}[*]${NC} $*"; }
section() {
    echo ""
    echo -e "${BOLD}${BLUE}─────────────────────────────────────────────${NC}"
    echo -e "${BOLD}$*${NC}"
    echo -e "${BOLD}${BLUE}─────────────────────────────────────────────${NC}"
    echo ""
}
banner() {
    echo -e "${BOLD}${MAGENTA}╔════════════════════════════════════════════╗${NC}"
    echo -e "${BOLD}${MAGENTA}║  BIGBROTHER // OPERATOR SETUP              ║${NC}"
    echo -e "${BOLD}${MAGENTA}╚════════════════════════════════════════════╝${NC}"
    echo ""
}

prompt_default() {
    local var="$1" prompt="$2" default="$3"
    echo -n -e "${CYAN}${prompt}${NC}"
    [[ -n "$default" ]] && echo -n -e " ${BOLD}[${default}]${NC}"
    echo -n ": "
    read -r input
    if [[ -z "$input" ]]; then
        printf -v "$var" '%s' "$default"
    else
        printf -v "$var" '%s' "$input"
    fi
}

# ══════════════════════════════════════════════════════════════════════════════
# PREFLIGHT
# ══════════════════════════════════════════════════════════════════════════════
[[ $EUID -eq 0 ]] || error "Must run as root (use: sudo bash operator_setup.sh)"

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
# When run via sudo, use the invoking user's home — not /root
REAL_HOME=$(getent passwd "${SUDO_USER:-$(whoami)}" | cut -d: -f6)
CACHE_DIR="${REAL_HOME}/.bigbrother/images"
CREDS="${REAL_HOME}/.local/bin/creds"
MOUNT_POINT="/mnt/opi"

mkdir -p "$CACHE_DIR"

# ══════════════════════════════════════════════════════════════════════════════
# 1. BANNER
# ══════════════════════════════════════════════════════════════════════════════
clear
banner

# ══════════════════════════════════════════════════════════════════════════════
# 2. SD CARD DETECTION
# ══════════════════════════════════════════════════════════════════════════════
section "SD Card Detection"
info "Scanning block devices..."
echo ""

lsblk -ndo NAME,SIZE,LABEL,TYPE | grep -E "^(sd|mmc)" | nl -v 1 -w1 -s') ' || error "No block devices found. Insert SD card and try again."

echo ""
echo -n -e "${CYAN}Enter device number (1-9):${NC} "
read -r device_num
[[ "$device_num" =~ ^[1-9]$ ]] || error "Invalid selection"

device_row=$(lsblk -ndo NAME,SIZE,LABEL,TYPE | grep -E "^(sd|mmc)" | sed -n "${device_num}p" || true)
DEVICE=$(echo "$device_row" | awk '{print $1}')
DEVICE_SIZE=$(echo "$device_row" | awk '{print $2}')
[[ -n "$DEVICE" ]] || error "Invalid selection — no device at that number"

info "Selected: /dev/${DEVICE} (${DEVICE_SIZE})"
warn "This will OVERWRITE all data on /dev/${DEVICE}"
echo -n -e "${YELLOW}Type device name to confirm (${DEVICE}):${NC} "
read -r confirm_dev
[[ "$confirm_dev" == "$DEVICE" ]] || error "Device name mismatch. Aborting."

info "Device confirmed: /dev/${DEVICE}"

# ══════════════════════════════════════════════════════════════════════════════
# 3. IMAGE MANAGEMENT
# ══════════════════════════════════════════════════════════════════════════════
section "Image Management"

# Look for existing cached image
CACHED_IMAGE=$(find "$CACHE_DIR" -maxdepth 1 -name 'Armbian_*_Orangepizero3_bookworm_current_*_minimal.img.xz' 2>/dev/null | head -1 || true)

if [[ -n "$CACHED_IMAGE" ]]; then
    info "Found cached image: $(basename $CACHED_IMAGE)"
    echo -n -e "${CYAN}Use cached image? (y/n):${NC} "
    read -r use_cached
    if [[ "$use_cached" =~ ^[yY]$ ]]; then
        IMAGE="$CACHED_IMAGE"
    else
        IMAGE=""
    fi
else
    IMAGE=""
fi

if [[ -z "$IMAGE" ]]; then
    step "Downloading Armbian image..."
    # Hardcoded known-good URL (operator can update)
    IMAGE_URL="https://dl.armbian.com/orangepizero3/Armbian_25.5.1_Orangepizero3_bookworm_current_6.12.23_minimal.img.xz"
    IMAGE_FILENAME=$(basename "$IMAGE_URL")
    IMAGE="${CACHE_DIR}/${IMAGE_FILENAME}"

    if [[ ! -f "$IMAGE" ]]; then
        curl -L --progress-bar "$IMAGE_URL" -o "$IMAGE" || error "Failed to download image"
        # Verify download: file exists and is > 50MB
        if [[ ! -s "$IMAGE" ]] || [[ $(stat -c%s "$IMAGE" 2>/dev/null || echo 0) -lt 52428800 ]]; then
            rm -f "$IMAGE"
            error "Download failed or file too small (expected > 50MB)"
        fi
        info "Image downloaded: $IMAGE_FILENAME"
    else
        info "Image already cached: $IMAGE_FILENAME"
    fi
fi

# SHA256 verification (optional, skip if no .sha file)
SHA_FILE="${IMAGE}.sha"
if [[ -f "$SHA_FILE" ]]; then
    step "Verifying SHA256..."
    cd "$CACHE_DIR"
    sha256sum -c "$(basename $SHA_FILE)" >/dev/null 2>&1 && info "SHA256 verified" || warn "SHA256 mismatch (proceeding anyway)"
    cd - >/dev/null
fi

# ══════════════════════════════════════════════════════════════════════════════
# 4. CONFIGURATION WIZARD
# ══════════════════════════════════════════════════════════════════════════════
section "Configuration Wizard"

# WiFi (optional)
step "WiFi Configuration"
prompt_default enable_wifi "Configure WiFi? (y/n)" "n"
if [[ "$enable_wifi" =~ ^[yY]$ ]]; then
    prompt_default wifi_ssid "WiFi SSID" ""
    while true; do
        echo -n -e "${CYAN}WiFi Passphrase${NC}: "
        read -s -r wifi_pass; echo
        echo -n -e "${CYAN}Confirm Passphrase${NC}: "
        read -s -r wifi_pass2; echo
        if [[ "$wifi_pass" == "$wifi_pass2" ]]; then
            break
        else
            warn "Passphrases do not match. Try again."
        fi
    done
    WIFI_SSID="$wifi_ssid"
    WIFI_PASS="$wifi_pass"
    unset wifi_pass wifi_pass2
else
    WIFI_SSID=""
    WIFI_PASS=""
fi

# Root password (required)
step "Root Password"
DEFAULT_ROOT_PASS=$(openssl rand -base64 16 | tr -d '=+/')
ROOT_PASS_IS_AUTO=false
echo -e "${YELLOW}Auto-generated password: ${BOLD}${DEFAULT_ROOT_PASS}${NC}"
prompt_default root_pass "Enter custom password or press Enter to use auto-generated" ""
if [[ -z "$root_pass" ]]; then
    ROOT_PASS="$DEFAULT_ROOT_PASS"
    ROOT_PASS_IS_AUTO=true
else
    ROOT_PASS="$root_pass"
fi
unset root_pass

# SSH Key
step "SSH Key Configuration"
echo "  [1] Generate new keypair (stored in Infisical)"
echo "  [2] Use existing public key file"
prompt_default ssh_opt "SSH option (1-2)" "1"

if [[ "$ssh_opt" == "1" ]]; then
    SSH_MODE="generate"
elif [[ "$ssh_opt" == "2" ]]; then
    SSH_MODE="existing"
    prompt_default ssh_pubkey_path "Path to public key file" ""
    [[ -f "$ssh_pubkey_path" ]] || error "File not found: $ssh_pubkey_path"
    SSH_PUBKEY=$(cat "$ssh_pubkey_path")
else
    error "Invalid SSH option"
fi

# Device identity — pick profile first, derive hostname + device ID from same OUI
step "Device Identity"
PROFILE_NAMES=(
    "Fire TV Stick 4K|amazon-firetv|FC:65:DE"
    "Fire TV Stick Lite|amazon-lite|FC:65:DE"
    "Fire TV Cube|amazon-cube|FC:65:DE"
    "Roku Ultra|roku-ultra|B8:3E:59"
    "Roku Express|roku-express|B8:3E:59"
    "Chromecast|chromecast|54:60:09"
    "Chromecast with Google TV|chromecast-gtv|54:60:09"
    "Apple TV 4K|apple-tv-4k|F0:B4:79"
    "Apple TV HD|apple-tv-hd|F0:B4:79"
    "Samsung Smart TV|samsung-tv|8C:79:F0"
    "Samsung QLED TV|samsung-qled|8C:79:F0"
    "LG webOS TV|lg-tv-webos|A8:23:FE"
    "LG OLED TV|lg-tv-oled|A8:23:FE"
    "Vizio SmartCast TV|vizio-tv|00:E0:36"
    "TCL Roku TV|tcl-tv|54:B7:B2"
    "Sony Bravia Android TV|sony-tv|F0:D7:AA"
    "Sonos One|sonos-one|94:9F:3E"
    "Sonos Beam|sonos-beam|94:9F:3E"
    "Google Home Mini|google-home|54:60:09"
    "Google Nest Hub|google-nest|54:60:09"
    "Echo Dot|echo-dot|FC:65:DE"
    "Echo Show|echo-show|FC:65:DE"
    "Ring Doorbell|ring-doorbell|B0:09:DA"
    "Nest Thermostat|nest-thermo|18:B4:30"
    "Nest Protect|nest-protect|18:B4:30"
    "Philips Hue Bridge|philips-hue|00:17:88"
    "TP-Link Kasa Smart Plug|tplink-kasa|50:C7:BF"
    "Wyze Cam v3|wyze-cam|2C:AA:8E"
    "Ecobee Thermostat|ecobee-thermo|44:61:32"
    "Roomba i7|roomba-i7|80:91:33"
    "Roomba j7|roomba-j7|80:91:33"
    "Raspberry Pi IoT Sensor|raspberrypi-sensor|DC:A6:32"
    "Raspberry Pi IoT Gateway|raspberrypi-gateway|DC:A6:32"
    "Raspberry Pi IoT Controller|raspberrypi-ctrl|DC:A6:32"
    "HP LaserJet Pro|hp-laserjet|3C:D9:2B"
    "HP OfficeJet Pro|hp-officejet|3C:D9:2B"
    "Brother HL-L2350DW|brother-hl|00:80:77"
    "Brother MFC-L2750DW|brother-mfc|00:80:77"
    "Brother DCP-L2550DW|brother-dcp|00:80:77"
    "Epson EcoTank ET-4760|epson-ecotank|00:26:AB"
    "Canon PIXMA G6020|canon-pixma|00:1E:8F"
    "Xbox Series X|xbox-series-x|7C:1E:52"
    "PlayStation 5|ps5|BC:60:A7"
    "PlayStation 4|ps4|BC:60:A7"
    "Nintendo Switch|nintendo-switch|98:B6:E9"
    "iPhone 15 Pro|iphone-15|3C:22:FB"
    "iPad Air|ipad-air|3C:22:FB"
    "Galaxy Tab S9|galaxy-tab|8C:BE:BE"
    "Galaxy S24|galaxy-s24|8C:BE:BE"
    "UniFi AP|unifi-ap|24:A4:3C"
    "USG Gateway|usg-gateway|24:A4:3C"
    "TP-Link Archer AX50|archer-ax50|50:C7:BF"
    "TP-Link Deco M5|deco-m5|50:C7:BF"
    "Netgear Orbi|orbi|9C:3D:CF"
    "Netgear Nighthawk|nighthawk|9C:3D:CF"
)
PROFILE_IDX=$(( RANDOM % ${#PROFILE_NAMES[@]} ))
SELECTED_PROFILE="${PROFILE_NAMES[$PROFILE_IDX]}"
PROFILE_DISPLAY="${SELECTED_PROFILE%%|*}"
_prest="${SELECTED_PROFILE#*|}"
DEFAULT_HOSTNAME="${_prest%%|*}"
PROFILE_OUI="${SELECTED_PROFILE##*|}"

echo -e "  Profile: ${BOLD}${PROFILE_DISPLAY}${NC} (OUI: ${PROFILE_OUI})"
prompt_default bb_hostname "Hostname" "$DEFAULT_HOSTNAME"
BB_HOSTNAME="$bb_hostname"

# Device ID derived from same profile OUI — MAC-format, consistent with network identity
OUI_CLEAN=$(echo "$PROFILE_OUI" | tr -d ':' | tr '[:upper:]' '[:lower:]')
RAND_SUFFIX=$(printf '%02x%02x%02x' $(( RANDOM % 256 )) $(( RANDOM % 256 )) $(( RANDOM % 256 )))
DEFAULT_DEVICE_ID="${OUI_CLEAN}${RAND_SUFFIX}"
prompt_default device_id "Device ID" "$DEFAULT_DEVICE_ID"
DEVICE_ID="$device_id"

# Auto-start mode
step "Auto-start Mode on Boot"
echo "  [1] Passive only (capture + analysis, no active offense)"
echo "  [2] Full autonomous (all modules)"
echo "  [3] Manual (operator SSHs in to start)"
prompt_default mode_opt "Select mode (1-3)" "1"

if [[ "$mode_opt" == "1" ]]; then
    BB_MODE="passive"
elif [[ "$mode_opt" == "2" ]]; then
    BB_MODE="autonomous"
elif [[ "$mode_opt" == "3" ]]; then
    BB_MODE="manual"
else
    error "Invalid mode"
fi

# VPN tunnel
step "VPN Tunnel"
echo "  [1] None"
echo "  [2] Tailscale"
echo "  [3] WireGuard"
echo "  [4] Reverse SSH Tunnel (autossh) — no traffic until operator connects"
prompt_default vpn_opt "Select VPN (1-4)" "2"

BB_VPN="none"
TAILSCALE_KEY=""
WG_CONFIG=""
AUTOSSH_SERVER=""
AUTOSSH_USER="root"
AUTOSSH_PORT="22"
AUTOSSH_REMOTE_PORT=""
AUTOSSH_PRIVKEY=""
AUTOSSH_PUBKEY=""

if [[ "$vpn_opt" == "1" ]]; then
    BB_VPN="none"
elif [[ "$vpn_opt" == "2" ]]; then
    BB_VPN="tailscale"
    echo -e "  ${YELLOW}Use a reusable, non-expiring key from Tailscale admin → Settings → Auth keys${NC}"
    prompt_default tailscale_key "Tailscale auth key" ""
    TAILSCALE_KEY="$tailscale_key"
elif [[ "$vpn_opt" == "3" ]]; then
    BB_VPN="wireguard"
    prompt_default wg_config_path "Path to WireGuard config (wg0.conf)" ""
    [[ -f "$wg_config_path" ]] || error "File not found: $wg_config_path"
    WG_CONFIG=$(cat "$wg_config_path")
elif [[ "$vpn_opt" == "4" ]]; then
    BB_VPN="autossh"
    prompt_default AUTOSSH_SERVER "Operator SSH server (host or IP)" ""
    [[ -n "$AUTOSSH_SERVER" ]] || error "Server address required"
    prompt_default AUTOSSH_USER "SSH user on server" "root"
    prompt_default AUTOSSH_PORT "SSH port on server" "22"
    DEFAULT_REMOTE_PORT=$(( 20000 + RANDOM % 10000 ))
    prompt_default AUTOSSH_REMOTE_PORT "Remote tunnel port (device SSH appears here on server)" "$DEFAULT_REMOTE_PORT"

    # Generate dedicated tunnel keypair — keeps it separate from device management key
    _TMP_TDIR=$(mktemp -d)
    _TMP_TKEY="${_TMP_TDIR}/id_ed25519"
    ssh-keygen -t ed25519 -N "" -C "bb-tunnel-${DEVICE_ID}" -f "$_TMP_TKEY" -q
    AUTOSSH_PRIVKEY=$(cat "$_TMP_TKEY")
    AUTOSSH_PUBKEY=$(cat "${_TMP_TKEY}.pub")
    rm -rf "$_TMP_TDIR"

    echo ""
    echo -e "  ${YELLOW}Add this public key to ${AUTOSSH_USER}@${AUTOSSH_SERVER}:~/.ssh/authorized_keys before powering on:${NC}"
    echo ""
    echo -e "  ${BOLD}${AUTOSSH_PUBKEY}${NC}"
    echo ""
    echo -n -e "${CYAN}  Press Enter once the key is on the server:${NC} "
    read -r
else
    error "Invalid VPN option"
fi

# Network alerter webhook
step "Network Alerter"
prompt_default enable_alerter "Enable online alert? (y/n)" "n"
BB_ALERTER_TYPE=""
BB_ALERTER_WEBHOOK=""
BB_ALERTER_MATRIX_HS=""
BB_ALERTER_MATRIX_ROOM=""
BB_ALERTER_MATRIX_TOKEN=""
BB_ALERTER_INFRA_IPS=""
if [[ "$enable_alerter" =~ ^[yY]$ ]]; then
    echo "  [1] Matrix"
    echo "  [2] Slack / Discord (webhook URL)"
    echo "  [3] ntfy"
    prompt_default alerter_type "Alerter type (1-3)" "1"
    if [[ "$alerter_type" == "1" ]]; then
        BB_ALERTER_TYPE="matrix"
        prompt_default matrix_hs "Homeserver URL" ""
        prompt_default matrix_room "Room ID (e.g. !abc123:homeserver)" ""
        prompt_default matrix_token "Bot access token" ""
        BB_ALERTER_MATRIX_HS="$matrix_hs"
        BB_ALERTER_MATRIX_ROOM="$matrix_room"
        BB_ALERTER_MATRIX_TOKEN="$matrix_token"
    elif [[ "$alerter_type" == "2" ]]; then
        BB_ALERTER_TYPE="webhook"
        prompt_default webhook_url "Webhook URL" ""
        BB_ALERTER_WEBHOOK="$webhook_url"
    elif [[ "$alerter_type" == "3" ]]; then
        BB_ALERTER_TYPE="ntfy"
        prompt_default ntfy_url "ntfy topic URL" ""
        BB_ALERTER_WEBHOOK="$ntfy_url"
    fi

fi

# BLE Presence Detection
step "BLE Presence Detection"
prompt_default enable_ble "Enable BLE presence alerter? (y/n)" "n"
BB_BLE_MODE="open"
BB_BLE_DEVICES=""
BB_BLE_TIMEOUT="60"
BB_BLE_MATRIX_HS=""
BB_BLE_MATRIX_ROOM=""
BB_BLE_MATRIX_TOKEN=""
if [[ "$enable_ble" =~ ^[yY]$ ]]; then
    echo "  [1] Open mode (alert on any named device)"
    echo "  [2] Known mode (alert only for specific devices)"
    prompt_default ble_mode "BLE mode (1-2)" "1"
    if [[ "$ble_mode" == "1" ]]; then
        BB_BLE_MODE="open"
    elif [[ "$ble_mode" == "2" ]]; then
        BB_BLE_MODE="known"
        echo -e "  ${YELLOW}Enter device names one per line (empty line to finish)${NC}"
        _DEVICE_LIST=""
        while true; do
            echo -n -e "    ${CYAN}Device name:${NC} "
            read -r _DEV_NAME
            [[ -z "$_DEV_NAME" ]] && break
            _DEVICE_LIST="${_DEVICE_LIST}${_DEVICE_LIST:+,}${_DEV_NAME}"
        done
        BB_BLE_DEVICES="$_DEVICE_LIST"
    fi
    prompt_default ble_timeout "Departure timeout (seconds)" "60"
    BB_BLE_TIMEOUT="$ble_timeout"
    echo -e "  ${YELLOW}Use same Matrix credentials as network alerter${NC}"
    prompt_default ble_matrix_hs "Homeserver URL" "${BB_ALERTER_MATRIX_HS:-}"
    prompt_default ble_matrix_room "Room ID" "${BB_ALERTER_MATRIX_ROOM:-}"
    prompt_default ble_matrix_token "Bot access token" "${BB_ALERTER_MATRIX_TOKEN:-}"
    BB_BLE_MATRIX_HS="$ble_matrix_hs"
    BB_BLE_MATRIX_ROOM="$ble_matrix_room"
    BB_BLE_MATRIX_TOKEN="$ble_matrix_token"
fi

# ══════════════════════════════════════════════════════════════════════════════
# 5. CONFIG SUMMARY
# ══════════════════════════════════════════════════════════════════════════════
section "Configuration Summary"

echo -e "Device ID:          ${BOLD}${DEVICE_ID}${NC}"
echo -e "Hostname:           ${BOLD}${BB_HOSTNAME}${NC}"
echo -e "Target:             ${BOLD}/dev/${DEVICE}${NC}"
echo -e "Image:              ${BOLD}$(basename $IMAGE)${NC}"
echo ""
if [[ "$ROOT_PASS_IS_AUTO" == true ]]; then
    echo -e "Root Pass:          ${BOLD}${RED}AUTO-GENERATED — RECORD THIS: ${DEFAULT_ROOT_PASS}${NC}"
else
    echo -e "Root Pass:          ${BOLD}[custom]${NC}"
fi
echo -e "WiFi:               ${BOLD}$([ -n "$WIFI_SSID" ] && echo "$WIFI_SSID" || echo "disabled")${NC}"
echo -e "SSH Mode:           ${BOLD}${SSH_MODE}${NC}"
echo -e "Auto-start:         ${BOLD}${BB_MODE}${NC}"
if [[ "$BB_VPN" == "autossh" ]]; then
    echo -e "VPN:                ${BOLD}autossh → ${AUTOSSH_USER}@${AUTOSSH_SERVER}:${AUTOSSH_PORT} (remote port ${AUTOSSH_REMOTE_PORT})${NC}"
else
    echo -e "VPN:                ${BOLD}${BB_VPN}${NC}"
fi
echo -e "Alerter:            ${BOLD}$([ -n "$BB_ALERTER_TYPE" ] && echo "$BB_ALERTER_TYPE" || echo "disabled")${NC}"
echo -e "BLE Alerter:        ${BOLD}$([ -n "$BB_BLE_MODE" ] && echo "$BB_BLE_MODE" || echo "disabled")${NC}"
echo ""

echo -n -e "${CYAN}Proceed with flashing and configuration? (y/n):${NC} "
read -r confirm_proceed
[[ "$confirm_proceed" =~ ^[yY]$ ]] || { info "Aborted"; exit 0; }

# ══════════════════════════════════════════════════════════════════════════════
# 6. FLASH AND MOUNT
# ══════════════════════════════════════════════════════════════════════════════
section "Flashing SD Card"

step "Unmounting any existing mounts on /dev/${DEVICE}..."
for part in /dev/${DEVICE}*; do
    [[ -b "$part" ]] && umount "$part" 2>/dev/null || true
done

step "Flashing image to /dev/${DEVICE}..."
xzcat "$IMAGE" | dd of="/dev/${DEVICE}" bs=4M status=progress conv=fsync || error "Flash failed"
info "Flash complete"

step "Running partprobe..."
partprobe "/dev/${DEVICE}" || warn "partprobe failed (non-critical)"
sleep 3

step "Mounting rootfs to ${MOUNT_POINT}..."
mkdir -p "$MOUNT_POINT"
# Armbian images: p1=FAT32 boot, p2=ext4 rootfs. Mount the rootfs partition.
# Detect by filesystem type; fall back to p1 for single-partition images.
_ROOTFS_PART=""
for _try in "/dev/${DEVICE}2" "/dev/${DEVICE}p2"; do
    if [[ -b "$_try" ]] && blkid "$_try" 2>/dev/null | grep -qi ext; then
        _ROOTFS_PART="$_try"
        break
    fi
done
if [[ -n "$_ROOTFS_PART" ]]; then
    mount "$_ROOTFS_PART" "$MOUNT_POINT" || error "Mount of rootfs (${_ROOTFS_PART}) failed"
else
    # Single-partition image fallback
    mount "/dev/${DEVICE}1" "$MOUNT_POINT" 2>/dev/null || mount "/dev/${DEVICE}p1" "$MOUNT_POINT" || error "Mount failed — could not locate ext4 rootfs partition"
fi

info "Mounted at ${MOUNT_POINT} ($(blkid -o value -s TYPE "${_ROOTFS_PART:-/dev/${DEVICE}1}"))"

# ══════════════════════════════════════════════════════════════════════════════
# 7. APPLY CONFIGURATION TO CARD
# ══════════════════════════════════════════════════════════════════════════════
section "Applying Configuration"

# WiFi profile — Armbian uses Netplan + wpa_supplicant (not NetworkManager)
if [[ -n "$WIFI_SSID" ]]; then
    step "Writing WiFi profile..."

    # Remove any netplan wifi config — conflicts with wpa_supplicant@wlan0.service
    rm -f "${MOUNT_POINT}/etc/netplan/20-wifi.yaml"

    # wpa_supplicant direct config (backup — used by wpa_supplicant@wlan0.service)
    mkdir -p "${MOUNT_POINT}/etc/wpa_supplicant"
    cat > "${MOUNT_POINT}/etc/wpa_supplicant/wpa_supplicant-wlan0.conf" << EOF
ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
update_config=1
country=US

network={
    ssid="${WIFI_SSID}"
    psk="${WIFI_PASS}"
    key_mgmt=WPA-PSK
    scan_ssid=1
}
EOF
    chmod 600 "${MOUNT_POINT}/etc/wpa_supplicant/wpa_supplicant-wlan0.conf"

    # systemd-networkd config for wlan0 DHCP
    mkdir -p "${MOUNT_POINT}/etc/systemd/network"
    cat > "${MOUNT_POINT}/etc/systemd/network/10-wifi-wlan0.network" << EOF
[Match]
Name=wlan0

[Network]
DHCP=yes
IPv6PrivacyExtensions=yes
EOF

    # Disable the generic wpa_supplicant.service — it conflicts with @wlan0 instance
    # Both would fight over wlan0, causing WiFi instability and SSH drops
    rm -f "${MOUNT_POINT}/etc/systemd/system/multi-user.target.wants/wpa_supplicant.service"

    # Enable the per-interface instance — this is what actually drives wlan0
    mkdir -p "${MOUNT_POINT}/etc/systemd/system/multi-user.target.wants"
    ln -sf /lib/systemd/system/wpa_supplicant@.service \
        "${MOUNT_POINT}/etc/systemd/system/multi-user.target.wants/wpa_supplicant@wlan0.service"

    # Disable WiFi power saving — aw859a driver drops under load with power save on
    mkdir -p "${MOUNT_POINT}/etc/udev/rules.d"
    cat > "${MOUNT_POINT}/etc/udev/rules.d/70-wifi-pm.rules" << 'EOF'
ACTION=="add", SUBSYSTEM=="net", KERNEL=="wlan*", RUN+="/usr/sbin/iw dev %k set power_save off"
EOF

    info "WiFi profile written (SSID: ${WIFI_SSID})"
fi

# Patch ethernet netplan config to suppress default route
step "Patching ethernet netplan config..."
mkdir -p "${MOUNT_POINT}/etc/netplan"
cat > "${MOUNT_POINT}/etc/netplan/10-dhcp-all-interfaces.yaml" << 'EOF'
# Added by Armbian (modified: suppress default route on ethernet)
network:
  version: 2
  renderer: networkd
  ethernets:
    all-eth-interfaces:
      match:
        name: "e*"
      dhcp4: yes
      dhcp6: yes
      ipv6-privacy: yes
      dhcp4-overrides:
        use-routes: false
      dhcp6-overrides:
        use-routes: false
EOF
info "Ethernet netplan patched"

# Root password
step "Setting root password..."
ROOT_PASS_HASH=$(openssl passwd -6 "$ROOT_PASS")
sed -i "s|^root:[^:]*:|root:${ROOT_PASS_HASH}:|" "${MOUNT_POINT}/etc/shadow"
info "Root password set"
unset ROOT_PASS ROOT_PASS_HASH

# SSH keys
step "Configuring SSH..."
SSH_DIR="${MOUNT_POINT}/root/.ssh"
mkdir -p "$SSH_DIR"

if [[ "$SSH_MODE" == "generate" ]]; then
    TMP_DIR=$(mktemp -d)
    TMP_KEY="${TMP_DIR}/id_ed25519"
    ssh-keygen -t ed25519 -N "" -C "" -f "$TMP_KEY" -q
    PRIVKEY=$(cat "$TMP_KEY")
    PUBKEY=$(cat "${TMP_KEY}.pub")
    rm -rf "$TMP_DIR"

    echo "$PUBKEY" > "${SSH_DIR}/authorized_keys"

    INFISICAL_KEY="SSH_PRIVKEY_$(echo "$DEVICE_ID" | tr '[:lower:]-' '[:upper:]_')"
    if [[ -x "$CREDS" ]] && "$CREDS" set "$INFISICAL_KEY" "$PRIVKEY" 2>/dev/null; then
        info "Private key stored in Infisical: ${INFISICAL_KEY}"
    else
        warn "Could not store in Infisical. SAVE THIS PRIVATE KEY:"
        echo "---"
        echo "$PRIVKEY"
        echo "---"
    fi

    # Save to the invoking user's ~/.ssh (SUDO_USER if running via sudo, else HOME)
    REAL_HOME=$(getent passwd "${SUDO_USER:-$USER}" 2>/dev/null | cut -d: -f6)
    LOCAL_KEY_DIR="${REAL_HOME:-$HOME}/.ssh"
    LOCAL_KEY_PATH="${LOCAL_KEY_DIR}/bb-${DEVICE_ID}"
    mkdir -p "$LOCAL_KEY_DIR"
    echo "$PRIVKEY" > "$LOCAL_KEY_PATH"
    chmod 600 "$LOCAL_KEY_PATH"
    [[ -n "$SUDO_USER" ]] && chown "$SUDO_USER" "$LOCAL_KEY_PATH" || true
    info "Private key saved locally: ${LOCAL_KEY_PATH}"

    unset PRIVKEY
else
    echo "$SSH_PUBKEY" > "${SSH_DIR}/authorized_keys"
fi

chmod 700 "$SSH_DIR"
chmod 600 "${SSH_DIR}/authorized_keys"
info "SSH authorized_keys configured"

# Ensure SSH is allowed through nftables (Armbian default drops input)
step "Configuring nftables firewall..."
cat > "${MOUNT_POINT}/etc/nftables.conf" << 'NFTEOF'
#!/usr/sbin/nft -f
flush ruleset

table inet filter {
    chain input {
        type filter hook input priority 0; policy drop;
        ct state established,related accept
        iif lo accept
        ip protocol icmp accept
        ip6 nexthdr icmpv6 accept
        tcp dport 22 accept
    }
    chain forward {
        type filter hook forward priority 0; policy drop;
    }
    chain output {
        type filter hook output priority 0; policy accept;
    }
}
NFTEOF
info "nftables config written (SSH allowed)"

# Hostname
step "Setting hostname..."
echo "$BB_HOSTNAME" > "${MOUNT_POINT}/etc/hostname"
sed -i "s/orangepi/$BB_HOSTNAME/g" "${MOUNT_POINT}/etc/hosts" 2>/dev/null || true
info "Hostname set to: $BB_HOSTNAME"

# Disable first-login wizard
rm -f "${MOUNT_POINT}/root/.not_logged_in_yet"

# Copy BigBrother source
step "Copying BigBrother source to card..."
SRC_DEST="${MOUNT_POINT}/root/bb-src"
rm -rf "$SRC_DEST"
rsync -a \
    --exclude='.git' \
    --exclude='.venv' \
    --exclude='storage/' \
    --exclude='__pycache__' \
    --exclude='*.pyc' \
    --exclude='*.db' \
    "${SCRIPT_DIR}/" "$SRC_DEST/" || error "rsync failed"
info "Source copied to /root/bb-src/"

# Write bb-config.env
step "Writing configuration to card..."
cat > "${MOUNT_POINT}/root/bb-config.env" << EOF
BB_MODE="${BB_MODE}"
BB_VPN="${BB_VPN}"
BB_DEVICE_ID="${DEVICE_ID}"
EOF

if [[ -n "$BB_ALERTER_TYPE" ]]; then
    echo "BB_ALERTER_TYPE=\"${BB_ALERTER_TYPE}\"" >> "${MOUNT_POINT}/root/bb-config.env"
fi
if [[ -n "$BB_ALERTER_WEBHOOK" ]]; then
    echo "BB_ALERTER_WEBHOOK=\"${BB_ALERTER_WEBHOOK}\"" >> "${MOUNT_POINT}/root/bb-config.env"
fi
if [[ -n "$BB_ALERTER_MATRIX_HS" ]]; then
    echo "BB_ALERTER_MATRIX_HS=\"${BB_ALERTER_MATRIX_HS}\"" >> "${MOUNT_POINT}/root/bb-config.env"
    echo "BB_ALERTER_MATRIX_ROOM=\"${BB_ALERTER_MATRIX_ROOM}\"" >> "${MOUNT_POINT}/root/bb-config.env"
    echo "BB_ALERTER_MATRIX_TOKEN=\"${BB_ALERTER_MATRIX_TOKEN}\"" >> "${MOUNT_POINT}/root/bb-config.env"
fi

# Store Tailscale auth key if needed
if [[ "$BB_VPN" == "tailscale" && -n "$TAILSCALE_KEY" ]]; then
    echo "BB_TAILSCALE_KEY=\"${TAILSCALE_KEY}\"" >> "${MOUNT_POINT}/root/bb-config.env"
fi

# Store WireGuard config if needed
if [[ "$BB_VPN" == "wireguard" && -n "$WG_CONFIG" ]]; then
    echo "$WG_CONFIG" > "${MOUNT_POINT}/etc/wireguard/wg0.conf"
    chmod 600 "${MOUNT_POINT}/etc/wireguard/wg0.conf"
fi

# Store autossh config if needed
if [[ "$BB_VPN" == "autossh" ]]; then
    echo "BB_AUTOSSH_SERVER=\"${AUTOSSH_SERVER}\"" >> "${MOUNT_POINT}/root/bb-config.env"
    echo "BB_AUTOSSH_USER=\"${AUTOSSH_USER}\"" >> "${MOUNT_POINT}/root/bb-config.env"
    echo "BB_AUTOSSH_PORT=\"${AUTOSSH_PORT}\"" >> "${MOUNT_POINT}/root/bb-config.env"
    echo "BB_AUTOSSH_REMOTE_PORT=\"${AUTOSSH_REMOTE_PORT}\"" >> "${MOUNT_POINT}/root/bb-config.env"

    # Write tunnel private key — separate from device management key
    mkdir -p "${MOUNT_POINT}/root/.ssh"
    echo "$AUTOSSH_PRIVKEY" > "${MOUNT_POINT}/root/.ssh/bb-tunnel-key"
    chmod 600 "${MOUNT_POINT}/root/.ssh/bb-tunnel-key"
    info "Tunnel key written to /root/.ssh/bb-tunnel-key"
    unset AUTOSSH_PRIVKEY
fi

chmod 600 "${MOUNT_POINT}/root/bb-config.env"
info "Configuration written"

# First-boot script
step "Installing first-boot service..."
cat > "${MOUNT_POINT}/root/bb-firstboot-run.sh" << 'RUNEOF'
#!/usr/bin/env bash
# Runs once on first boot, then disables itself
LOG=/root/bb-firstboot.log
exec >> "$LOG" 2>&1
set -uo pipefail

echo "=== BigBrother first-boot setup: $(date) ==="

# Wait for internet connectivity
echo "[*] Waiting for internet..."
for i in $(seq 1 30); do
    if curl -sf --max-time 5 http://deb.debian.org/debian > /dev/null 2>&1; then
        echo "[+] Internet OK"
        break
    fi
    if [[ $i -eq 30 ]]; then
        echo "[-] No internet after 90s"
        exit 1
    fi
    sleep 3
done

source /root/bb-config.env

# Install BigBrother (skip if setup.sh not present — dev/test mode)
if [[ -f /root/bb-src/setup.sh ]]; then
    cd /root/bb-src
    bash setup.sh --enable-service || echo "[-] BigBrother setup.sh failed (non-fatal)"
else
    echo "[-] bb-src/setup.sh not found — skipping BigBrother install"
fi

# VPN setup
if [[ "${BB_VPN:-}" == "tailscale" && -n "${BB_TAILSCALE_KEY:-}" ]]; then
    echo "[*] Installing Tailscale..."
    curl -4 -fsSL https://tailscale.com/install.sh | sh
    echo "[*] Connecting to Tailscale..."
    tailscale up --auth-key="${BB_TAILSCALE_KEY}" --accept-routes || echo "[-] Tailscale auth failed"
    # Remove auth key from config after use
    sed -i '/BB_TAILSCALE_KEY/d' /root/bb-config.env
elif [[ "${BB_VPN:-}" == "wireguard" ]]; then
    echo "[*] Installing WireGuard..."
    apt-get update >/dev/null 2>&1
    apt-get install -y wireguard >/dev/null 2>&1
    systemctl enable --now wg-quick@wg0 || echo "[-] WireGuard startup failed"
elif [[ "${BB_VPN:-}" == "autossh" ]]; then
    echo "[*] Installing autossh..."
    apt-get update >/dev/null 2>&1
    apt-get install -y autossh >/dev/null 2>&1
    chmod 600 /root/.ssh/bb-tunnel-key 2>/dev/null || true

    cat > /etc/systemd/system/bb-tunnel.service << SVCEOF
[Unit]
Description=BigBrother reverse SSH tunnel
After=network-online.target
Wants=network-online.target

[Service]
Environment="AUTOSSH_GATETIME=0"
ExecStart=/usr/bin/autossh -M 0 -N \
    -o ServerAliveInterval=30 \
    -o ServerAliveCountMax=3 \
    -o StrictHostKeyChecking=no \
    -o ExitOnForwardFailure=yes \
    -i /root/.ssh/bb-tunnel-key \
    -p ${BB_AUTOSSH_PORT} \
    -R ${BB_AUTOSSH_REMOTE_PORT}:localhost:22 \
    ${BB_AUTOSSH_USER}@${BB_AUTOSSH_SERVER}
Restart=always
RestartSec=30
User=root

[Install]
WantedBy=multi-user.target
SVCEOF

    systemctl daemon-reload
    systemctl enable --now bb-tunnel.service || echo "[-] Tunnel service startup failed"
    echo "[+] Reverse SSH tunnel configured (remote port ${BB_AUTOSSH_REMOTE_PORT})"
fi

# Network alerter
if [[ -n "${BB_ALERTER_TYPE:-}" ]]; then
    echo "[*] Sending online alert..."
    _BB_IP=$(curl -sf --max-time 5 https://api.ipify.org 2>/dev/null || hostname -I | awk '{print $1}')
    _BB_MSG="${BB_DEVICE_ID} online | ip: ${_BB_IP} | host: $(hostname)"

    if [[ "${BB_ALERTER_TYPE}" == "matrix" ]]; then
        _TXNID=$(date +%s%N)
        _ROOM="${BB_ALERTER_MATRIX_ROOM//!/%21}"
        _ROOM="${_ROOM//:/%3A}"
        curl -sf -X PUT \
            "${BB_ALERTER_MATRIX_HS}/_matrix/client/v3/rooms/${_ROOM}/send/m.room.message/${_TXNID}" \
            -H "Authorization: Bearer ${BB_ALERTER_MATRIX_TOKEN}" \
            -H "Content-Type: application/json" \
            -d "{\"msgtype\":\"m.text\",\"body\":\"${_BB_MSG}\"}" || true
    elif [[ "${BB_ALERTER_TYPE}" == "ntfy" ]]; then
        curl -sf -X POST "${BB_ALERTER_WEBHOOK}" -d "${_BB_MSG}" || true
    else
        curl -sf -X POST "${BB_ALERTER_WEBHOOK}" \
            -H 'Content-Type: application/json' \
            -d "{\"text\":\"${_BB_MSG}\"}" || true
    fi
fi

echo "[+] First-boot complete: $(date)"

# Disable ourselves
rm -f /root/.bb-firstboot-needed
systemctl disable bb-firstboot.service
RUNEOF
chmod 755 "${MOUNT_POINT}/root/bb-firstboot-run.sh"

# First-boot systemd service
cat > "${MOUNT_POINT}/etc/systemd/system/bb-firstboot.service" << 'SVCEOF'
[Unit]
Description=BigBrother first-boot setup
After=network-online.target
Wants=network-online.target
ConditionPathExists=/root/.bb-firstboot-needed

[Service]
Type=oneshot
ExecStart=/bin/bash /root/bb-firstboot-run.sh
RemainAfterExit=yes
TimeoutStartSec=900
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
SVCEOF

# Enable service symlink
WANTS_DIR="${MOUNT_POINT}/etc/systemd/system/multi-user.target.wants"
mkdir -p "$WANTS_DIR"
ln -sf /etc/systemd/system/bb-firstboot.service "${WANTS_DIR}/bb-firstboot.service"

# Firstboot sentinel
touch "${MOUNT_POINT}/root/.bb-firstboot-needed"

info "First-boot service installed and enabled"

# ══════════════════════════════════════════════════════════════════════════════
# 8. UNMOUNT
# ══════════════════════════════════════════════════════════════════════════════
section "Finalizing"

step "Syncing and unmounting..."
sync
umount "$MOUNT_POINT" || warn "umount failed (non-critical)"
info "SD card ready to remove"

# ══════════════════════════════════════════════════════════════════════════════
# 9. DONE
# ══════════════════════════════════════════════════════════════════════════════
section "Setup Complete"

echo -e "${GREEN}Insert SD card into device and power on.${NC}"
echo ""
echo "The device will:"
echo "  1. Boot and connect to WiFi (if configured)"
echo "  2. Run /root/bb-firstboot-run.sh automatically"
echo "  3. Install BigBrother source and dependencies"
echo "  4. Set up systemd services"
echo "  5. Configure VPN (if selected)"
echo "  6. Send alerter webhook (if configured)"
echo ""
echo "No further operator interaction required. Device walks away autonomously."
echo ""
echo -e "${CYAN}To SSH after first-boot:${NC}"

if [[ "$SSH_MODE" == "generate" ]]; then
    echo "  ssh -i ~/.ssh/bb-${DEVICE_ID} root@<DEVICE_IP>"
else
    echo "  ssh -i $(dirname $ssh_pubkey_path)/$(basename $ssh_pubkey_path .pub) root@<DEVICE_IP>"
fi

echo ""
if [[ "$BB_VPN" == "autossh" ]]; then
    echo ""
    echo -e "${CYAN}Tunnel connection (after first boot):${NC}"
    echo "  From operator server:  ssh -p ${AUTOSSH_REMOTE_PORT} root@127.0.0.1"
    echo "  One-liner from anywhere (requires GatewayPorts yes on server):"
    echo "    ssh -p ${AUTOSSH_REMOTE_PORT} root@${AUTOSSH_SERVER}"
fi

echo ""
echo -e "${CYAN}Monitor first-boot:${NC}"
echo "  ssh root@<DEVICE_IP> 'tail -f /root/bb-firstboot.log'"
echo ""