#!/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 ""
}

# ══════════════════════════════════════════════════════════════════════════════
# 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 -fL --progress-bar "$IMAGE_URL" -o "$IMAGE" || error "Failed to download image"
        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"
echo -n -e "${CYAN}Configure WiFi? (y/n):${NC} "
read -r enable_wifi
if [[ "$enable_wifi" =~ ^[yY]$ ]]; then
    echo -n -e "${CYAN}WiFi SSID:${NC} "
    read -r wifi_ssid
    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
    [[ "$wifi_pass" == "$wifi_pass2" ]] || error "Passphrases do not match"
    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"
while true; do
    echo -n -e "${CYAN}Root password:${NC} "
    read -s -r root_pass; echo
    echo -n -e "${CYAN}Confirm:${NC} "
    read -s -r root_pass2; echo
    if [[ "$root_pass" == "$root_pass2" && -n "$root_pass" ]]; then
        ROOT_PASS="$root_pass"
        unset root_pass root_pass2
        break
    else
        warn "Passwords do not match or empty. Try again."
    fi
done

# SSH Key
step "SSH Key Configuration"
echo "  [1] Generate new keypair (stored in Infisical)"
echo "  [2] Use existing public key file"
echo -n -e "${CYAN}SSH option (1-2):${NC} "
read -r ssh_opt

if [[ "$ssh_opt" == "1" ]]; then
    SSH_MODE="generate"
elif [[ "$ssh_opt" == "2" ]]; then
    SSH_MODE="existing"
    echo -n -e "${CYAN}Path to public key file:${NC} "
    read -r ssh_pubkey_path
    [[ -f "$ssh_pubkey_path" ]] || error "File not found: $ssh_pubkey_path"
    SSH_PUBKEY=$(cat "$ssh_pubkey_path")
else
    error "Invalid SSH option"
fi

# Device ID
step "Device ID"
echo -n -e "${CYAN}Device label (e.g. bb-01):${NC} "
read -r device_id
[[ -n "$device_id" ]] || error "Device ID cannot be empty"
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)"
echo -n -e "${CYAN}Select mode (1-3):${NC} "
read -r mode_opt

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 -n -e "${CYAN}Select VPN (1-3):${NC} "
read -r vpn_opt

BB_VPN="none"
TAILSCALE_KEY=""
WG_CONFIG=""

if [[ "$vpn_opt" == "1" ]]; then
    BB_VPN="none"
elif [[ "$vpn_opt" == "2" ]]; then
    BB_VPN="tailscale"
    echo -n -e "${CYAN}Tailscale auth key:${NC} "
    read -r tailscale_key
    TAILSCALE_KEY="$tailscale_key"
elif [[ "$vpn_opt" == "3" ]]; then
    BB_VPN="wireguard"
    echo -n -e "${CYAN}Path to WireGuard config (wg0.conf):${NC} "
    read -r wg_config_path
    [[ -f "$wg_config_path" ]] || error "File not found: $wg_config_path"
    WG_CONFIG=$(cat "$wg_config_path")
else
    error "Invalid VPN option"
fi

# Network alerter webhook
step "Network Alerter"
echo -n -e "${CYAN}Enable alerter webhook? (y/n):${NC} "
read -r enable_alerter
if [[ "$enable_alerter" =~ ^[yY]$ ]]; then
    echo -n -e "${CYAN}Webhook URL (Slack/Discord/Matrix):${NC} "
    read -r webhook_url
    BB_ALERTER_WEBHOOK="$webhook_url"
else
    BB_ALERTER_WEBHOOK=""
fi

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

echo "Device ID:          ${BOLD}${DEVICE_ID}${NC}"
echo "Target:             ${BOLD}/dev/${DEVICE}${NC}"
echo "Image:              ${BOLD}$(basename $IMAGE)${NC}"
echo ""
echo "WiFi:               ${BOLD}$([ -n "$WIFI_SSID" ] && echo "$WIFI_SSID" || echo "disabled")${NC}"
echo "Root Pass:          ${BOLD}[set]${NC}"
echo "SSH Mode:           ${BOLD}${SSH_MODE}${NC}"
echo "Auto-start:         ${BOLD}${BB_MODE}${NC}"
echo "VPN:                ${BOLD}${BB_VPN}${NC}"
echo "Alerter:            ${BOLD}$([ -n "$BB_ALERTER_WEBHOOK" ] && echo "enabled" || 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 to ${MOUNT_POINT}..."
mkdir -p "$MOUNT_POINT"
# Try /dev/${DEVICE}1 (most common for single-partition images)
mount "/dev/${DEVICE}1" "$MOUNT_POINT" 2>/dev/null || mount "/dev/${DEVICE}p1" "$MOUNT_POINT" || error "Mount failed. Check device partitioning."

info "Mounted at ${MOUNT_POINT}"

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

# WiFi profile
if [[ -n "$WIFI_SSID" ]]; then
    step "Writing WiFi profile..."
    NM_DIR="${MOUNT_POINT}/etc/NetworkManager/system-connections"
    NM_FILE="${NM_DIR}/wifi-bb.nmconnection"
    mkdir -p "$NM_DIR"

    WIFI_UUID=$(cat /proc/sys/kernel/random/uuid)
    cat > "$NM_FILE" << EOF
[connection]
id=wifi-bb
uuid=${WIFI_UUID}
type=wifi
autoconnect=true
autoconnect-priority=10

[wifi]
mode=infrastructure
ssid=${WIFI_SSID}

[wifi-security]
auth-alg=open
key-mgmt=wpa-psk
psk=${WIFI_PASS}

[ipv4]
method=auto

[ipv6]
addr-gen-mode=stable-privacy
method=auto
EOF

    chmod 600 "$NM_FILE"
    chown root:root "$NM_FILE"
    info "WiFi profile written (SSID: ${WIFI_SSID}, UUID: ${WIFI_UUID})"
fi

# 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_KEY=$(mktemp)
    ssh-keygen -t ed25519 -N "" -C "" -f "$TMP_KEY" -q
    PRIVKEY=$(cat "$TMP_KEY")
    PUBKEY=$(cat "${TMP_KEY}.pub")
    rm -f "$TMP_KEY" "${TMP_KEY}.pub"

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

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

# Hostname
echo "orangepi" > "${MOUNT_POINT}/etc/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_WEBHOOK" ]]; then
    echo "BB_ALERTER_WEBHOOK=\"${BB_ALERTER_WEBHOOK}\"" >> "${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

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
set -euo pipefail
LOG=/root/bb-firstboot.log
exec >> "$LOG" 2>&1

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
cd /root/bb-src
bash setup.sh --enable-service

# VPN setup
if [[ "${BB_VPN:-}" == "tailscale" && -n "${BB_TAILSCALE_KEY:-}" ]]; then
    echo "[*] Installing Tailscale..."
    curl -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"
fi

# Network alerter
if [[ -n "${BB_ALERTER_WEBHOOK:-}" ]]; then
    echo "[*] Sending online alert..."
    curl -sf -X POST "$BB_ALERTER_WEBHOOK" \
        -H 'Content-Type: application/json' \
        -d "{\"text\":\"BigBrother ${BB_DEVICE_ID} online\"}" || true
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 "${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 "${CYAN}To SSH after first-boot:${NC}"

if [[ "$SSH_MODE" == "generate" ]]; then
    INFISICAL_KEY="SSH_PRIVKEY_$(echo "$DEVICE_ID" | tr '[:lower:]-' '[:upper:]_')"
    if [[ -x "$CREDS" ]]; then
        echo "  ${CREDS} get ${INFISICAL_KEY} bigbrother > /tmp/${DEVICE_ID}.key"
        echo "  chmod 600 /tmp/${DEVICE_ID}.key"
        echo "  ssh -i /tmp/${DEVICE_ID}.key root@<DEVICE_IP>"
    fi
else
    echo "  ssh -i $(dirname $ssh_pubkey_path)/$(basename $ssh_pubkey_path .pub) root@<DEVICE_IP>"
fi

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