#!/bin/bash

set -e

# Configuration =======================================================================================================
#
# Connection settings (SSID, passphrase, subnet, DHCP range, interface) can
# be supplied three ways, highest priority first:
#   1. Environment variables (FALLBACK_AP_*) - used for headless installs.
#   2. The persisted config file below, written by `install`.
#   3. The hard-coded defaults in this block.
#
# Everything else (timers, paths) is not meant to be reconfigured per-install
# and stays as plain constants.

CONFIG_FILE="/etc/fallback-ap.conf"
if [ -f "$CONFIG_FILE" ]; then source "$CONFIG_FILE"; fi

WLAN0_IFACE="${FALLBACK_AP_IFACE:-${SAVED_IFACE:-wlan0}}"

AP_NAME="${FALLBACK_AP_SSID:-${SAVED_SSID:-$HOSTNAME}}"
AP_PASSWORD="${FALLBACK_AP_PASSWORD:-${SAVED_PASSWORD:-abcd1234}}"
AP_ENCRYPTION="${FALLBACK_AP_ENCRYPTION:-${SAVED_ENCRYPTION:-wpa-psk}}" # none | wpa-psk | sae

FALLBACK_AP_CIDR="${FALLBACK_AP_CIDR:-${SAVED_CIDR:-10.10.10.1/24}}"
FALLBACK_AP_DHCP_START_OCTET="${FALLBACK_AP_DHCP_START_OCTET:-${SAVED_DHCP_START_OCTET:-100}}"
FALLBACK_AP_DHCP_END_OCTET="${FALLBACK_AP_DHCP_END_OCTET:-${SAVED_DHCP_END_OCTET:-250}}"

AP_CIDR="$FALLBACK_AP_CIDR"
AP_NETWORK_BASE="${AP_CIDR%%/*}"   # e.g. "10.10.10.1" out of "10.10.10.1/24"
AP_NETWORK_BASE="${AP_NETWORK_BASE%.*}" # e.g. "10.10.10"
DHCP_START="${AP_NETWORK_BASE}.${FALLBACK_AP_DHCP_START_OCTET}"
DHCP_END="${AP_NETWORK_BASE}.${FALLBACK_AP_DHCP_END_OCTET}"

# How often, while the AP has no clients, to stop the AP and give
# NetworkManager a chance to reconnect to a known Wi-Fi network.
RECOVERY_CHECK_INTERVAL=60

# How long to leave the AP down during a recovery attempt before giving
# up and bringing the AP back.
RECOVERY_WINDOW=20

# Small delay after changing wlan0 state, to let things settle before
# we query it again.
STATE_SETTLE_DELAY=1

CONNECTION_NAME="Fallback AP"
SERVICE_NAME="fallback-ap"
SERVICE_FILE="/etc/systemd/system/$SERVICE_NAME.service"

STATE_DIR="/var/lib/fallback-ap"
LAST_WIFI_FILE="$STATE_DIR/last-wifi"
LOCK_FILE="$STATE_DIR/lock"
STATE_FILE="$STATE_DIR/state"

# End of Configuration ================================================================================================


# Resolve the actual location of this script.
SCRIPT_PATH="$(readlink -f "$0")"

if [ -z "$SCRIPT_PATH" ] || [ ! -f "$SCRIPT_PATH" ]; then
    echo "Unable to determine script path."
    exit 1
fi


log() {
    echo "[$(date '+%H:%M:%S')] $*"
}


set_state() {
    printf '%s\n' "$1" > "$STATE_FILE"
}


get_state() {
    cat "$STATE_FILE" 2>/dev/null || echo "unknown"
}


is_ap_on() {
    nmcli -t -f NAME connection show --active 2>/dev/null |
        grep -Fxq "$CONNECTION_NAME"
}


is_wifi_connected() {
    nmcli -t -f DEVICE,TYPE,STATE,CONNECTION device status 2>/dev/null |
        awk -F: -v iface="$WLAN0_IFACE" -v ap="$CONNECTION_NAME" \
        '$1 == iface && $2 == "wifi" && $3 == "connected" && $4 != ap {found=1} END {exit !found}'
}


get_wifi_connection() {
    nmcli -t -f DEVICE,TYPE,STATE,CONNECTION device status 2>/dev/null |
        awk -F: -v iface="$WLAN0_IFACE" \
        '$1 == iface && $2 == "wifi" && $3 == "connected" {print $4; exit}'
}


remember_wifi() {
    local connection

    connection="$(get_wifi_connection)"

    [ -z "$connection" ] && return
    [ "$connection" = "$CONNECTION_NAME" ] && return

    mkdir -p "$STATE_DIR"

    if [ "$(cat "$LAST_WIFI_FILE" 2>/dev/null || true)" != "$connection" ]; then
        log "Remembering Wi-Fi connection: $connection"
        printf '%s\n' "$connection" > "$LAST_WIFI_FILE"
    fi
}


get_last_wifi() {
    [ -f "$LAST_WIFI_FILE" ] && cat "$LAST_WIFI_FILE"
}


get_client_count() {
    iw dev "$WLAN0_IFACE" station dump 2>/dev/null |
        grep -c '^Station ' || true
}


wait_for_ap_state() {
    local wanted="$1"
    local timeout="${2:-5}"

    for ((i=0; i<timeout*10; i++)); do
        if [ "$wanted" = "on" ] && is_ap_on; then
            return 0
        fi

        if [ "$wanted" = "off" ] && ! is_ap_on; then
            return 0
        fi

        sleep 0.1
    done

    return 1
}


wait_for_wifi() {
    local timeout="$1"

    for ((i=0; i<timeout*10; i++)); do
        if is_wifi_connected; then
            return 0
        fi

        sleep 0.1
    done

    return 1
}


start_ap() {
    if is_ap_on; then
        return 0
    fi

    log "Enabling AP: $AP_NAME"

    nmcli connection up "$CONNECTION_NAME" ifname "$WLAN0_IFACE" >/dev/null 2>&1 || {
        log "Failed to start AP"
        return 1
    }

    if ! wait_for_ap_state on 5; then
        log "AP failed to become active"
        return 1
    fi

    sleep "$STATE_SETTLE_DELAY"

    log "AP enabled"

    return 0
}


stop_ap() {
    if ! is_ap_on; then
        return 0
    fi

    log "Disabling AP"

    nmcli connection down "$CONNECTION_NAME" >/dev/null 2>&1 || true

    if ! wait_for_ap_state off 5; then
        log "AP failed to stop cleanly"
        return 1
    fi

    sleep "$STATE_SETTLE_DELAY"

    return 0
}

tick() {
    local now="$1"

    if is_wifi_connected; then
        remember_wifi

        if is_ap_on; then
            stop_ap
        fi

        set_state "wifi"
        last_attempt="$now"
        return
    fi

    if ! is_ap_on; then
        start_ap
        set_state "fallback"
        last_attempt="$now"
        return
    fi

    set_state "fallback"

    local clients
    clients="$(get_client_count)"

    if [ "$clients" -gt 0 ]; then
        last_attempt="$now"
        return
    fi

    if (( now - last_attempt >= RECOVERY_CHECK_INTERVAL )); then
        last_attempt="$now"
        attempt_recovery
    fi
}

attempt_recovery() {
    log "No AP clients, checking for known Wi-Fi (AP down for up to ${RECOVERY_WINDOW}s)"

    set_state "checking"

    stop_ap || true

    if is_ap_on; then
        log "AP did not go down, aborting recovery attempt"
        start_ap
        set_state "fallback"
        return 1
    fi

    if wait_for_wifi "$RECOVERY_WINDOW"; then
        sleep "$STATE_SETTLE_DELAY"

        if is_wifi_connected; then
            log "Wi-Fi reconnected"
            remember_wifi
            set_state "wifi"
            return 0
        fi
    fi

    log "No Wi-Fi available, restoring AP"

    start_ap
    set_state "fallback"

    return 1
}


save_config() {
    mkdir -p "$(dirname "$CONFIG_FILE")"
    cat > "$CONFIG_FILE" <<EOF
SAVED_IFACE="$WLAN0_IFACE"
SAVED_SSID="$AP_NAME"
SAVED_PASSWORD="$AP_PASSWORD"
SAVED_ENCRYPTION="$AP_ENCRYPTION"
SAVED_CIDR="$FALLBACK_AP_CIDR"
SAVED_DHCP_START_OCTET="$FALLBACK_AP_DHCP_START_OCTET"
SAVED_DHCP_END_OCTET="$FALLBACK_AP_DHCP_END_OCTET"
EOF
    chmod 600 "$CONFIG_FILE"
}


install() {
    if [ "$EUID" -ne 0 ]; then
        echo "Run this script as root."
        exit 1
    fi

    log "Installing NetworkManager fallback AP"

    apt-get update
    apt-get install -y network-manager iw

    systemctl enable NetworkManager
    systemctl start NetworkManager

    mkdir -p "$STATE_DIR"

    if is_wifi_connected; then
        remember_wifi
    fi

    create_connection
    save_config

    log "Creating recovery service"

    cat > "$SERVICE_FILE" <<EOF
[Unit]
Description=NetworkManager Fallback AP
After=NetworkManager.service
Wants=NetworkManager.service

[Service]
Type=simple
ExecStart=$SCRIPT_PATH run
Restart=always
RestartSec=2

[Install]
WantedBy=multi-user.target
EOF

    chmod 644 "$SERVICE_FILE"

    systemctl daemon-reload
    systemctl enable "$SERVICE_NAME"
    systemctl restart "$SERVICE_NAME"

    log "Installation complete"
    log "Script: $SCRIPT_PATH"
    log "Config file: $CONFIG_FILE"
    log "SSID: $AP_NAME"
    log "Interface: $WLAN0_IFACE"
    log "Network: $AP_CIDR"
    log "DHCP: $DHCP_START - $DHCP_END"
    log "Encryption: $AP_ENCRYPTION"
    log "Recovery check interval: ${RECOVERY_CHECK_INTERVAL}s"
    log "Recovery window: ${RECOVERY_WINDOW}s"
    log "Last Wi-Fi: $(get_last_wifi || echo none)"

    echo
    nmcli device status
}


create_connection() {
    log "Creating NetworkManager AP connection..."

    nmcli connection delete "$CONNECTION_NAME" >/dev/null 2>&1 || true

    nmcli connection add \
        type wifi \
        ifname "$WLAN0_IFACE" \
        con-name "$CONNECTION_NAME" \
        autoconnect no \
        ssid "$AP_NAME"

    nmcli connection modify "$CONNECTION_NAME" \
        connection.autoconnect no \
        connection.autoconnect-priority -999 \
        connection.autoconnect-retries 0 \
        802-11-wireless.mode ap \
        802-11-wireless.powersave 2 \
        ipv4.method shared \
        ipv4.addresses "$AP_CIDR" \
        ipv4.shared-dhcp-range "$DHCP_START,$DHCP_END" \
        ipv6.method disabled

    case "$AP_ENCRYPTION" in
        none)
            nmcli connection modify "$CONNECTION_NAME" \
                wifi-sec.key-mgmt none
            ;;

        wpa-psk)
            if [ "${#AP_PASSWORD}" -lt 8 ] || [ "${#AP_PASSWORD}" -gt 63 ]; then
                echo "WPA-PSK password must be between 8 and 63 characters."
                exit 1
            fi

            nmcli connection modify "$CONNECTION_NAME" \
                wifi-sec.key-mgmt wpa-psk \
                wifi-sec.psk "$AP_PASSWORD"
            ;;

        sae)
            if [ "${#AP_PASSWORD}" -lt 8 ]; then
                echo "WPA3-SAE password must be at least 8 characters."
                exit 1
            fi

            nmcli connection modify "$CONNECTION_NAME" \
                wifi-sec.key-mgmt sae \
                wifi-sec.psk "$AP_PASSWORD"
            ;;

        *)
            echo "Unsupported encryption: $AP_ENCRYPTION"
            exit 1
            ;;
    esac
}


uninstall() {
    if [ "$EUID" -ne 0 ]; then
        echo "Run this script as root."
        exit 1
    fi

    log "Uninstalling fallback AP"

    systemctl disable --now "$SERVICE_NAME" >/dev/null 2>&1 || true

    nmcli connection down "$CONNECTION_NAME" >/dev/null 2>&1 || true
    nmcli connection delete "$CONNECTION_NAME" >/dev/null 2>&1 || true

    rm -f "$SERVICE_FILE"
    rm -f "$CONFIG_FILE"
    rm -rf "$STATE_DIR"

    systemctl daemon-reload

    log "Fallback AP removed"
}


start() {
    mkdir -p "$STATE_DIR"

    (
        flock -n 9 || exit 1
        start_ap
    ) 9>"$LOCK_FILE"
}


stop() {
    mkdir -p "$STATE_DIR"

    (
        flock -n 9 || exit 1
        stop_ap
    ) 9>"$LOCK_FILE"
}


status() {
    echo
    echo "=== Service ==="
    systemctl status "$SERVICE_NAME" --no-pager || true

    echo
    echo "=== Devices ==="
    nmcli device status

    echo
    echo "=== Active Connections ==="
    nmcli connection show --active

    echo
    echo "=== Saved Wi-Fi Profiles ==="
    nmcli -f NAME,TYPE,AUTOCONNECT,AUTOCONNECT-PRIORITY connection show

    echo
    echo "=== Fallback AP ==="
    nmcli connection show "$CONNECTION_NAME" 2>/dev/null || echo "Not installed"

    echo
    echo "=== Config ($CONFIG_FILE) ==="
    echo "SSID: $AP_NAME"
    echo "Interface: $WLAN0_IFACE"
    echo "Network: $AP_CIDR"
    echo "DHCP: $DHCP_START - $DHCP_END"

    echo
    echo "=== Last Known Wi-Fi ==="
    get_last_wifi || echo "None"

    echo
    echo "=== State ==="
    get_state

    echo
    echo "=== AP Clients ==="
    get_client_count

    echo
    echo "=== Script ==="
    echo "$SCRIPT_PATH"
}


run() {
    local last_attempt=0
    local now

    mkdir -p "$STATE_DIR"

    log "Fallback AP service started"
    log "Script: $SCRIPT_PATH"
    exec 9>"$LOCK_FILE"

    flock 9

    if is_wifi_connected; then
        remember_wifi
        set_state "wifi"
        log "Wi-Fi already connected"
    else
        log "No Wi-Fi connection, starting fallback AP"
        start_ap
        set_state "fallback"
    fi

    flock -u 9

    while true; do
        now="$(date +%s)"

        if flock -n 9; then
            tick "$now"
            flock -u 9
        fi

        sleep 1
    done
}


if [ "$EUID" -ne 0 ]; then
    echo "Run this script as root."
    exit 1
fi


case "${1:-}" in
    install)
        install
        ;;

    uninstall)
        uninstall
        ;;

    start)
        start
        ;;

    stop)
        stop
        ;;

    status)
        status
        ;;

    run)
        run
        ;;

    *)
        echo "Usage: $SCRIPT_PATH {install|uninstall|start|stop|status|run}"
        exit 1
        ;;
esac