#!/bin/bash
# enroll-biometrics.sh — TUI for face biometric setup.
#
# Two subsystems:
#   1. Presence detection — configure/test the webcam used by presence-detect.sh
#   2. Howdy face auth    — enroll/manage/test face models for PAM authentication

BACKTITLE="Biometric Enrollment"
PRESENCE_CFG="${XDG_CONFIG_HOME:-$HOME/.config}/presence-detect.conf"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PYTHON_DETECT="$SCRIPT_DIR/python/presence_detect.py"

# ── PAM 2FA (howdy face + FIDO U2F) ────────────────────────────────────────────
# Both factors required, with the normal password stack kept as a fallback so a
# dead camera or absent key can never lock you out.
PAM_TARGETS=(/etc/pam.d/sudo /etc/pam.d/hyprlock /etc/pam.d/login)
PAM_HOWDY_SO="/usr/lib/security/pam_howdy.so"
PAM_U2F_SO="/usr/lib/security/pam_u2f.so"
U2F_MAP="/etc/u2f_mappings"
PAM_MARK_BEGIN="# enroll-biometrics:begin howdy+u2f 2fa"
PAM_MARK_END="# enroll-biometrics:end howdy+u2f 2fa"

# ── Dialog theme (Cyberqueer) ─────────────────────────────────────────────────
TMP_D=$(mktemp -d)
trap 'rm -rf "$TMP_D"' EXIT INT TERM

export DIALOGRC="$TMP_D/dialogrc"
cat > "$DIALOGRC" << 'RCEOF'
use_shadow = ON
use_colors = ON
screen_color              = (BLACK,BLACK,ON)
shadow_color              = (BLACK,BLACK,ON)
title_color               = (MAGENTA,BLACK,ON)
border_color              = (MAGENTA,BLACK,ON)
button_active_color       = (BLACK,MAGENTA,ON)
button_inactive_color     = (WHITE,BLACK,OFF)
button_key_active_color   = (BLACK,CYAN,ON)
button_key_inactive_color = (CYAN,BLACK,ON)
button_label_active_color   = (BLACK,MAGENTA,ON)
button_label_inactive_color = (WHITE,BLACK,OFF)
inputbox_color        = (WHITE,BLACK,OFF)
inputbox_border_color = (MAGENTA,BLACK,ON)
menubox_color         = (WHITE,BLACK,OFF)
menubox_border_color  = (MAGENTA,BLACK,ON)
item_color            = (WHITE,BLACK,OFF)
item_selected_color   = (BLACK,MAGENTA,ON)
tag_color             = (CYAN,BLACK,ON)
tag_selected_color    = (BLACK,CYAN,ON)
check_color           = (WHITE,BLACK,OFF)
check_selected_color  = (BLACK,MAGENTA,ON)
uarrow_color          = (MAGENTA,BLACK,ON)
darrow_color          = (MAGENTA,BLACK,ON)
RCEOF

# ── Helpers ───────────────────────────────────────────────────────────────────
die() { clear; printf "\nError: %s\n\n" "$1" >&2; exit 1; }

msg() {
    dialog --backtitle "$BACKTITLE" --title " $1 " --msgbox "\n$2\n" "$3" "$4"
}

require_dialog() {
    command -v dialog &>/dev/null && return
    sudo pacman -S --noconfirm dialog || die "dialog not found"
}

# ── Camera helpers ─────────────────────────────────────────────────────────────
list_cameras() {
    for dev in /dev/video*; do
        [[ -c "$dev" ]] || continue
        local id="${dev#/dev/video}"
        local name
        name=$(v4l2-ctl --device="$dev" --info 2>/dev/null \
               | awk -F': ' '/Card type/{print $2}' | head -1)
        [[ -z "$name" ]] && name="Camera $id"
        printf '%s\n%s\n' "$id" "$name"
    done
}

get_camera_id() {
    if [[ -n "$PRESENCE_DETECT_CAMERA" ]]; then
        echo "$PRESENCE_DETECT_CAMERA"
    elif [[ -f "$PRESENCE_CFG" ]]; then
        grep -oP 'CAMERA=\K[0-9]+' "$PRESENCE_CFG" 2>/dev/null || echo 0
    else
        echo 0
    fi
}

set_camera_id() {
    mkdir -p "$(dirname "$PRESENCE_CFG")"
    printf 'CAMERA=%s\n' "$1" > "$PRESENCE_CFG"
}

# ── Presence detection ────────────────────────────────────────────────────────
presence_configure_camera() {
    local -a cam_items
    mapfile -t cam_items < <(list_cameras)

    if [[ ${#cam_items[@]} -eq 0 ]]; then
        msg "No Cameras Found" \
            "No video devices found at /dev/video*.\n\nMake sure your webcam is connected." \
            10 55
        return
    fi

    local current; current=$(get_camera_id)
    local choice
    choice=$(dialog --backtitle "$BACKTITLE" \
        --title " Select Camera " \
        --menu "\nCurrent: /dev/video${current}\n\nSelect camera to use for presence detection:" \
        16 62 6 \
        "${cam_items[@]}" \
        3>&1 1>&2 2>&3) || return

    set_camera_id "$choice"
    msg "Camera Set" \
        "Presence detection will use /dev/video${choice}.\n\nRestart presence-detect.sh for the change to take effect." \
        9 62
}

presence_test_camera() {
    local cam; cam=$(get_camera_id)

    clear
    printf "\n\033[1;35m  Testing presence detection on /dev/video%s...\033[0m\n" "$cam"
    printf "\033[35m  ─────────────────────────────────────────\033[0m\n\n"
    printf "  Look at the camera and stay still.\n\n"

    python3 "$PYTHON_DETECT" "$cam" 2>/dev/null
    local rc=$?

    case $rc in
        0) msg "Test Result" "Face detected!\n\nPresence detection is working correctly." 8 52 ;;
        1) msg "Test Result" \
               "No face detected.\n\nMake sure you are in front of the camera\nand there is adequate lighting." \
               10 56 ;;
        2) msg "Camera Error" \
               "Could not open /dev/video${cam}.\n\nTry configuring a different camera first." \
               9 56 ;;
        *) msg "Error" "Unexpected exit code ($rc) from detection script." 7 52 ;;
    esac
}

# ── Howdy helpers ─────────────────────────────────────────────────────────────
howdy_installed() { command -v howdy &>/dev/null; }

howdy_require() {
    howdy_installed && return 0

    # howdy lives only in the AUR — it is not in the official repos, so a
    # plain `pacman -S howdy` always fails. Require an AUR helper.
    local helper
    for h in yay paru; do
        command -v "$h" &>/dev/null && { helper="$h"; break; }
    done
    if [[ -z "$helper" ]]; then
        msg "No AUR Helper" \
            "howdy is an AUR package and needs an AUR helper (yay or paru) to install.\n\nInstall yay first, then re-run this menu." \
            10 60
        return 1
    fi

    dialog --backtitle "$BACKTITLE" --title " Howdy Not Found " \
        --yesno "\nhowdy is not installed.\n\nInstall it now via ${helper}?" 8 48 || return 1
    clear
    printf "\nInstalling howdy via %s...\n\n" "$helper"
    "$helper" -S --noconfirm --needed howdy
    howdy_installed || { msg "Install Failed" "howdy installation failed.\nInstall it manually: ${helper} -S howdy" 8 52; return 1; }
}

# ── Howdy operations ──────────────────────────────────────────────────────────
howdy_add() {
    howdy_require || return

    local name
    name=$(dialog --backtitle "$BACKTITLE" \
        --title " Add Face Model " \
        --inputbox "\nEnter a label for this face model:" 8 52 "$USER" \
        3>&1 1>&2 2>&3) || return
    [[ -z "$name" ]] && return

    clear
    printf "\n\033[1;35m  Enrolling face model '%s'...\033[0m\n" "$name"
    printf "\033[35m  ─────────────────────────────────────────\033[0m\n\n"
    printf "  Look at the camera and hold still.\n\n"
    # howdy has no label flag; the label is read from stdin during `add`.
    printf '%s\n' "$name" | sudo howdy -U "$USER" add
    local rc=${PIPESTATUS[1]}

    if [[ $rc -eq 0 ]]; then
        msg "Enrolled" \
            "Face model '$name' added.\n\nYou may need to relogin for PAM changes to take effect." \
            9 58
    else
        msg "Failed" "Enrollment failed (code $rc).\n\nCheck camera and lighting." 8 52
    fi
}

howdy_list() {
    howdy_require || return
    local out; out=$(sudo howdy -U "$USER" list 2>&1)
    msg "Enrolled Face Models" "$out" 20 64
}

howdy_remove() {
    howdy_require || return

    local models; models=$(sudo howdy -U "$USER" list 2>&1)
    if [[ -z "$models" || "$models" == *"no models"* ]]; then
        msg "No Models" "No face models are currently enrolled." 7 44
        return
    fi

    local id
    id=$(dialog --backtitle "$BACKTITLE" \
        --title " Remove Face Model " \
        --inputbox "\nCurrent models:\n\n${models}\n\nEnter model ID to remove:" 20 64 \
        3>&1 1>&2 2>&3) || return
    [[ -z "$id" || ! "$id" =~ ^[0-9]+$ ]] && {
        msg "Invalid ID" "Please enter a numeric model ID." 6 36; return
    }

    dialog --backtitle "$BACKTITLE" --title " Confirm " \
        --yesno "\nRemove face model ID $id?" 7 40 || return

    sudo howdy -U "$USER" remove "$id"
    msg "Done" "Face model $id removed." 6 38
}

howdy_test() {
    howdy_require || return
    clear
    printf "\n\033[1;35m  Testing howdy authentication...\033[0m\n"
    printf "\033[35m  ─────────────────────────────────────────\033[0m\n\n"
    printf "  Look at the camera.\n\n"
    sudo howdy -U "$USER" test
    local rc=$?
    if [[ $rc -eq 0 ]]; then
        msg "Result" "Authentication successful." 6 38
    else
        msg "Result" "Authentication failed (code $rc).\n\nCheck enrolled models and camera setup." 9 52
    fi
}

# ── FIDO U2F helpers ────────────────────────────────────────────────────────────
fido_installed() { [[ -e "$PAM_U2F_SO" ]] && command -v pamu2fcfg &>/dev/null; }

fido_require() {
    fido_installed && return 0
    dialog --backtitle "$BACKTITLE" --title " FIDO U2F Not Found " \
        --yesno "\npam-u2f is not installed.\n\nInstall it now via pacman?" 8 50 || return 1
    clear
    printf "\nInstalling pam-u2f...\n\n"
    sudo pacman -S --noconfirm --needed pam-u2f
    fido_installed || { msg "Install Failed" "pam-u2f installation failed.\nInstall it manually: sudo pacman -S pam-u2f" 8 56; return 1; }
}

# Ensure $USER has a registered key in the central mapping; register one if not.
fido_register() {
    if sudo grep -q "^${USER}:" "$U2F_MAP" 2>/dev/null; then
        dialog --backtitle "$BACKTITLE" --title " FIDO Key " \
            --yesno "\nA FIDO key is already registered for $USER.\n\nRegister an additional key?" 9 56 \
            || return 0
        clear
        printf "\n\033[1;35m  Insert your FIDO key and touch it when it blinks...\033[0m\n\n"
        local extra
        extra=$(pamu2fcfg -n -u "$USER") || { msg "Registration Failed" "Could not read FIDO key." 7 44; return 1; }
        [[ -z "$extra" ]] && { msg "No Key" "No FIDO key detected." 7 36; return 1; }
        # pamu2fcfg -n prints just the credential; append it to the user's line.
        sudo sed -i "s|^\(${USER}:.*\)\$|\1:${extra}|" "$U2F_MAP"
        return 0
    fi

    clear
    printf "\n\033[1;35m  Registering FIDO key for %s...\033[0m\n" "$USER"
    printf "\033[35m  ─────────────────────────────────────────\033[0m\n\n"
    printf "  Insert your FIDO key and touch it when it blinks.\n\n"
    local line
    line=$(pamu2fcfg -u "$USER") || { msg "Registration Failed" "Could not read FIDO key.\n\nMake sure it is plugged in." 9 48; return 1; }
    [[ -z "$line" ]] && { msg "No Key" "No FIDO key detected." 7 36; return 1; }
    printf '%s\n' "$line" | sudo tee -a "$U2F_MAP" >/dev/null
    sudo chmod 600 "$U2F_MAP"
    sudo chown root:root "$U2F_MAP"
    return 0
}

# ── PAM file editing ──────────────────────────────────────────────────────────
# The inserted block: try howdy, then U2F. Both must pass to satisfy auth and
# skip the password stack; if either fails, control falls through to the
# existing (password) auth lines below.
pam_emit_block() {
    printf '%s\n' "$PAM_MARK_BEGIN"
    printf 'auth      [success=ignore default=1] pam_howdy.so\n'
    printf 'auth      sufficient                 pam_u2f.so authfile=%s cue\n' "$U2F_MAP"
    printf '%s\n' "$PAM_MARK_END"
}

pam_apply_file() {
    local f="$1"
    [[ -f "$f" ]] || { printf '  skip  %s (not present)\n' "$f"; return 0; }
    if sudo grep -qF "$PAM_MARK_BEGIN" "$f"; then
        printf '  ok    %s (already configured)\n' "$f"; return 0
    fi
    sudo cp -a "$f" "${f}.bak.$(date +%s)"
    local tmp first; tmp=$(mktemp); first=$(sudo head -n1 "$f")
    if [[ "$first" == "#%PAM"* ]]; then
        { printf '%s\n' "$first"; pam_emit_block; sudo tail -n +2 "$f"; } > "$tmp"
    else
        { pam_emit_block; sudo cat "$f"; } > "$tmp"
    fi
    sudo cp "$tmp" "$f"; rm -f "$tmp"
    printf '  added %s\n' "$f"
}

pam_remove_file() {
    local f="$1"
    [[ -f "$f" ]] || return 0
    sudo grep -qF "$PAM_MARK_BEGIN" "$f" || { printf '  skip  %s (not configured)\n' "$f"; return 0; }
    sudo cp -a "$f" "${f}.bak.$(date +%s)"
    local tmp; tmp=$(mktemp)
    sudo awk -v b="$PAM_MARK_BEGIN" -v e="$PAM_MARK_END" \
        '$0==b{s=1} !s{print} $0==e{s=0}' "$f" > "$tmp"
    sudo cp "$tmp" "$f"; rm -f "$tmp"
    printf '  removed from %s\n' "$f"
}

# hyprlock ships as `auth include login`. Since the 2FA block also lands in
# /etc/pam.d/login, a plain include would re-run the block on the failure path.
# Give hyprlock its own clean fallback by pointing it at system-auth instead.
HYPRLOCK_REDIRECT='s/^(auth[[:space:]]+include[[:space:]]+)login[[:space:]]*$/\1system-auth/'
HYPRLOCK_RESTORE='s/^(auth[[:space:]]+include[[:space:]]+)system-auth[[:space:]]*$/\1login/'

pam_apply_hyprlock() {
    local f="/etc/pam.d/hyprlock"
    [[ -f "$f" ]] || { printf '  skip  %s (not present)\n' "$f"; return 0; }
    if sudo grep -qF "$PAM_MARK_BEGIN" "$f"; then
        printf '  ok    %s (already configured)\n' "$f"; return 0
    fi
    sudo cp -a "$f" "${f}.bak.$(date +%s)"
    local tmp first; tmp=$(mktemp); first=$(sudo head -n1 "$f")
    if [[ "$first" == "#%PAM"* ]]; then
        { printf '%s\n' "$first"; pam_emit_block
          sudo tail -n +2 "$f" | sed -E "$HYPRLOCK_REDIRECT"; } > "$tmp"
    else
        { pam_emit_block; sudo cat "$f" | sed -E "$HYPRLOCK_REDIRECT"; } > "$tmp"
    fi
    sudo cp "$tmp" "$f"; rm -f "$tmp"
    printf '  added %s (fallback → system-auth)\n' "$f"
}

pam_remove_hyprlock() {
    local f="/etc/pam.d/hyprlock"
    [[ -f "$f" ]] || return 0
    sudo grep -qF "$PAM_MARK_BEGIN" "$f" || { printf '  skip  %s (not configured)\n' "$f"; return 0; }
    sudo cp -a "$f" "${f}.bak.$(date +%s)"
    local tmp; tmp=$(mktemp)
    sudo awk -v b="$PAM_MARK_BEGIN" -v e="$PAM_MARK_END" \
        '$0==b{s=1} !s{print} $0==e{s=0}' "$f" | sed -E "$HYPRLOCK_RESTORE" > "$tmp"
    sudo cp "$tmp" "$f"; rm -f "$tmp"
    printf '  removed from %s (fallback → login)\n' "$f"
}

pam_setup() {
    howdy_require || return
    fido_require  || return

    if [[ ! -e "$PAM_HOWDY_SO" ]]; then
        msg "Missing PAM Module" "pam_howdy.so not found at:\n$PAM_HOWDY_SO\n\nReinstall howdy and try again." 10 60
        return
    fi

    dialog --backtitle "$BACKTITLE" --title " Set Up PAM 2FA " --yesno \
"\nThis will require BOTH face (howdy) and a FIDO key for:\n\n  ${PAM_TARGETS[*]}\n\nYour password still works if a factor is unavailable.\nBackups of each file are saved as *.bak.<timestamp>.\n\nContinue?" 15 64 || return

    fido_register || return

    clear
    printf "\n\033[1;35m  Configuring PAM...\033[0m\n"
    printf "\033[35m  ─────────────────────────────────────────\033[0m\n\n"
    local f
    for f in "${PAM_TARGETS[@]}"; do
        if [[ "$f" == */hyprlock ]]; then pam_apply_hyprlock; else pam_apply_file "$f"; fi
    done
    printf "\n"
    msg "PAM Configured" \
"Face + FIDO 2FA is active (password fallback kept).\n\n*** KEEP THIS TERMINAL OPEN. ***\nIn a NEW terminal, run:  sudo -k; sudo true\nand confirm you can authenticate.\n\nIf locked out, restore a *.bak.<timestamp> file." 15 64
}

pam_teardown() {
    dialog --backtitle "$BACKTITLE" --title " Remove PAM 2FA " --yesno \
"\nRemove the howdy + FIDO 2FA block from:\n\n  ${PAM_TARGETS[*]}\n\nContinue?" 12 64 || return
    clear
    printf "\n\033[1;35m  Removing PAM 2FA...\033[0m\n\n"
    local f
    for f in "${PAM_TARGETS[@]}"; do
        if [[ "$f" == */hyprlock ]]; then pam_remove_hyprlock; else pam_remove_file "$f"; fi
    done
    printf "\n"
    msg "PAM Restored" "The 2FA block was removed.\nBackups remain as *.bak.<timestamp>." 8 56
}

# ── Main menu ─────────────────────────────────────────────────────────────────
main_menu() {
    local choice
    choice=$(dialog --backtitle "$BACKTITLE" \
        --title " Biometric Enrollment " \
        --menu "\nSelect an option:" 22 70 9 \
        "1" "Presence detection  — configure camera" \
        "2" "Presence detection  — test detection" \
        "3" "Howdy face auth     — add face model" \
        "4" "Howdy face auth     — list enrolled models" \
        "5" "Howdy face auth     — remove face model" \
        "6" "Howdy face auth     — test authentication" \
        "7" "PAM 2FA            — set up howdy + FIDO key" \
        "8" "PAM 2FA            — remove howdy + FIDO key" \
        3>&1 1>&2 2>&3) || { clear; exit 0; }

    case "$choice" in
        1) presence_configure_camera ;;
        2) presence_test_camera       ;;
        3) howdy_add                  ;;
        4) howdy_list                 ;;
        5) howdy_remove               ;;
        6) howdy_test                 ;;
        7) pam_setup                  ;;
        8) pam_teardown               ;;
    esac
    main_menu
}

require_dialog
main_menu