#!/usr/bin/env bash
set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

info()   { echo -e "${GREEN}[+]${NC} $*"; }
warn()   { echo -e "${YELLOW}[!]${NC} $*"; }
die()    { echo -e "${RED}[ERROR]${NC} $*" >&2; exit 1; }
header() { echo -e "\n${CYAN}══════════════════════════════════════════${NC}"; echo -e "${CYAN}  $*${NC}"; echo -e "${CYAN}══════════════════════════════════════════${NC}"; }

check_deps() {
    local missing=()
    for cmd in "$@"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done
    if [ ${#missing[@]} -gt 0 ]; then
        die "Missing required dependencies: ${missing[*]}"
    fi
}

[ "$(id -u)" -eq 0 ] || die "Run as root"

check_deps ip netplan systemctl ping hostnamectl

header "TinyBoard Network Setup"
echo ""
echo "  0) Change hostname"
echo "  1) Configure static IP"
echo "  2) Prefer IPv4 over IPv6"
echo "  3) Prefer IPv6 over IPv4"
echo "  4) Change Wireless Network"
echo "  q) Quit"
echo ""
read -rp "Choose: " NET_OPT
echo ""

case "$NET_OPT" in
0)
    header "Change Hostname"
    CURRENT_HOSTNAME=$(hostname)
    echo -e "Current hostname: ${YELLOW}$CURRENT_HOSTNAME${NC}"
    read -rp "Enter new hostname (e.g. rocky): " NEW_HOSTNAME
    [ -n "$NEW_HOSTNAME" ] || die "Hostname cannot be empty."
    [[ "$NEW_HOSTNAME" =~ ^[a-zA-Z0-9._-]+$ ]] || die "Invalid hostname — use only letters, numbers, dots, underscores, hyphens."
    hostnamectl set-hostname "$NEW_HOSTNAME"
    echo "$NEW_HOSTNAME" > /etc/hostname
    sed -i "s/${CURRENT_HOSTNAME}/${NEW_HOSTNAME}/g" /etc/hosts
    info "Hostname changed to: $NEW_HOSTNAME"
    exit 0
    ;;
1)
    ;;
2)
    header "Prefer IPv4 over IPv6"
    if grep -q "precedence ::ffff:0:0/96" /etc/gai.conf 2>/dev/null; then
        warn "IPv4 preference already set."
    else
        echo "precedence ::ffff:0:0/96  100" >> /etc/gai.conf
        info "IPv4 preference set. Outgoing connections will prefer IPv4."
    fi
    exit 0
    ;;
3)
    header "Prefer IPv6 over IPv4"
    sed -i '/precedence ::ffff:0:0\/96/d' /etc/gai.conf 2>/dev/null || true
    info "IPv4 preference removed. System will use default IPv6-first behavior."
    exit 0
    ;;
4)
    header "Change Wireless Network"
    check_deps iw wpa_cli wpa_passphrase

    WIFI_IFACE=$(iw dev 2>/dev/null | awk '/Interface/{print $2}' | head -1)
    [ -n "$WIFI_IFACE" ] || die "No wireless interface found."

    CURRENT_SSID=$(iw dev "$WIFI_IFACE" link 2>/dev/null | awk '/SSID:/{print $2}')

    info "Scanning for networks on ${WIFI_IFACE}..."
    ip link set "$WIFI_IFACE" up 2>/dev/null || true
    iw dev "$WIFI_IFACE" scan >/dev/null 2>&1 || true
    sleep 2

    mapfile -t SCAN_LINES < <(iw dev "$WIFI_IFACE" scan 2>/dev/null \
        | awk '
            /^BSS /         { signal=""; ssid="" }
            /signal:/       { signal=$2 }
            /SSID:/         { ssid=substr($0, index($0,$2)); gsub(/^[[:space:]]+|[[:space:]]+$/, "", ssid) }
            ssid!="" && signal!="" { print signal "\t" ssid; signal=""; ssid="" }
        ' \
        | sort -rn \
        | awk -F'\t' '!seen[$2]++ && $2!="" {print $2}')

    [ ${#SCAN_LINES[@]} -gt 0 ] || die "No wireless networks found. Ensure the interface is up and try again."

    echo ""
    for i in "${!SCAN_LINES[@]}"; do
        SSID="${SCAN_LINES[$i]}"
        if [ "$SSID" = "$CURRENT_SSID" ]; then
            echo -e "  $((i+1))) ${GREEN}${SSID}${NC} ${CYAN}(connected)${NC}"
        else
            echo -e "  $((i+1))) ${SSID}"
        fi
    done
    echo ""

    read -rp "Enter network number to join: " WIFI_CHOICE
    [[ "$WIFI_CHOICE" =~ ^[0-9]+$ ]] || die "Invalid selection."
    WIFI_IDX=$((WIFI_CHOICE - 1))
    [ "$WIFI_IDX" -ge 0 ] && [ "$WIFI_IDX" -lt "${#SCAN_LINES[@]}" ] || die "Selection out of range."

    NEW_SSID="${SCAN_LINES[$WIFI_IDX]}"
    echo ""
    read -rsp "Password for '${NEW_SSID}': " NEW_PASS
    echo ""
    [ -n "$NEW_PASS" ] || die "Password cannot be empty."

    WPA_CONF=$(wpa_passphrase "$NEW_SSID" "$NEW_PASS") \
        || die "Failed to generate WPA config — check SSID and password."

    info "Currently connected to: ${CURRENT_SSID:-none}"
    info "Switching to:           ${NEW_SSID}"
    warn "Your SSH session will drop. Reconnect once the device joins '${NEW_SSID}'."
    echo ""
    read -rp "Proceed? [Y/n]: " CONFIRM
    CONFIRM="${CONFIRM:-y}"
    [[ "${CONFIRM,,}" == "y" ]] || { info "Aborted."; exit 0; }

    EXISTING_IDS=$(wpa_cli -i "$WIFI_IFACE" list_networks 2>/dev/null | awk 'NR>1 {print $1}')
    if [[ -z "$EXISTING_IDS" ]] && ! wpa_cli -i "$WIFI_IFACE" status >/dev/null 2>&1; then
        die "wpa_supplicant is not running on ${WIFI_IFACE}. Start it first."
    fi

    for EID in $EXISTING_IDS; do
        wpa_cli -i "$WIFI_IFACE" disable_network "$EID" >/dev/null 2>&1 || true
    done

    NETWORK_ID=$(wpa_cli -i "$WIFI_IFACE" add_network 2>/dev/null | tr -d '[:space:]')
    [[ "$NETWORK_ID" =~ ^[0-9]+$ ]] || die "Failed to add network — wpa_supplicant may not be running."

    PSK=$(echo "$WPA_CONF" | awk -F= '/^\s*psk=/{print $2}' | grep -v '"' | tr -d '[:space:]')

    wpa_cli -i "$WIFI_IFACE" set_network "$NETWORK_ID" ssid "\"${NEW_SSID}\"" >/dev/null
    wpa_cli -i "$WIFI_IFACE" set_network "$NETWORK_ID" psk "${PSK}" >/dev/null
    wpa_cli -i "$WIFI_IFACE" set_network "$NETWORK_ID" freq_list "2412 2417 2422 2427 2432 2437 2442 2447 2452 2457 2462" >/dev/null
    wpa_cli -i "$WIFI_IFACE" select_network "$NETWORK_ID" >/dev/null

    info "Waiting for association..."
    ASSOCIATED=false
    for i in $(seq 1 10); do
        sleep 2
        STATUS=$(wpa_cli -i "$WIFI_IFACE" status 2>/dev/null | awk -F= '/^wpa_state=/{print $2}')
        CONN_SSID=$(wpa_cli -i "$WIFI_IFACE" status 2>/dev/null | awk -F= '/^ssid=/{print $2}')
        if [ "$STATUS" = "COMPLETED" ] && [ "$CONN_SSID" = "$NEW_SSID" ]; then
            ASSOCIATED=true
            break
        fi
        warn "Attempt $i/10 — state: ${STATUS:-unknown}, ssid: ${CONN_SSID:-none}"
    done

    if [ "$ASSOCIATED" = "false" ]; then
        wpa_cli -i "$WIFI_IFACE" remove_network "$NETWORK_ID" >/dev/null 2>&1 || true
        die "Failed to associate with '${NEW_SSID}'. Check the password and try again."
    fi

    wpa_cli -i "$WIFI_IFACE" save_config >/dev/null 2>&1 || true

    info "Associated — renewing DHCP lease..."
    ip link set "$WIFI_IFACE" down 2>/dev/null || true
    sleep 1
    ip link set "$WIFI_IFACE" up 2>/dev/null || true
    sleep 1

    if systemctl is-active --quiet "systemd-networkd"; then
        networkctl reconfigure "$WIFI_IFACE" 2>/dev/null || true
    fi

    if command -v dhclient >/dev/null 2>&1; then
        dhclient -r "$WIFI_IFACE" 2>/dev/null || true
        dhclient "$WIFI_IFACE" 2>/dev/null || true
    elif command -v udhcpc >/dev/null 2>&1; then
        udhcpc -i "$WIFI_IFACE" -q 2>/dev/null || true
    fi

    sleep 3
    NEW_IP=$(ip -o -4 addr show "$WIFI_IFACE" 2>/dev/null | awk '{print $4}' | head -1)
    if [ -n "$NEW_IP" ]; then
        info "IP address: ${NEW_IP}"
    else
        warn "No IP assigned yet — DHCP may still be in progress."
    fi

    NETPLAN_BACKUP_DIR="/root/.config/tinyboard/netplan-backups"
    mkdir -p "$NETPLAN_BACKUP_DIR"

    mapfile -t NETPLAN_FILES < <(find /etc/netplan -maxdepth 1 -name '*.yaml' 2>/dev/null | sort)

    if [ ${#NETPLAN_FILES[@]} -eq 0 ]; then
        warn "No netplan config files found — skipping persistent config update."
        NETPLAN_FILE=""
    elif [ ${#NETPLAN_FILES[@]} -eq 1 ]; then
        NETPLAN_FILE="${NETPLAN_FILES[0]}"
    else
        warn "Multiple netplan config files found:"
        for i in "${!NETPLAN_FILES[@]}"; do
            echo -e "  $((i+1))) ${NETPLAN_FILES[$i]}"
        done
        echo ""
        read -rp "Which file should be updated with the new WiFi credentials? [1]: " NP_CHOICE
        NP_CHOICE="${NP_CHOICE:-1}"
        [[ "$NP_CHOICE" =~ ^[0-9]+$ ]] || die "Invalid selection."
        NP_IDX=$((NP_CHOICE - 1))
        [ "$NP_IDX" -ge 0 ] && [ "$NP_IDX" -lt "${#NETPLAN_FILES[@]}" ] || die "Selection out of range."
        NETPLAN_FILE="${NETPLAN_FILES[$NP_IDX]}"
    fi

    if [ -n "$NETPLAN_FILE" ] && grep -q "access-points" "$NETPLAN_FILE" 2>/dev/null; then
        BACKUP_FILE="$NETPLAN_BACKUP_DIR/$(basename "${NETPLAN_FILE}").$(date +%Y%m%d%H%M%S)"
        cp "$NETPLAN_FILE" "$BACKUP_FILE"
        info "Backed up: $NETPLAN_FILE → $BACKUP_FILE"

        python3 - "$NETPLAN_FILE" "$NEW_SSID" "$NEW_PASS" <<'PYEOF'
import sys, re

path, ssid, pw = sys.argv[1], sys.argv[2], sys.argv[3]
txt = open(path).read()
txt = re.sub(
    r'(access-points:\s*\n\s+)["\']?[^"\':\n]+["\']?:\s*\n(\s+password:)[^\n]*',
    lambda m: f'{m.group(1)}"{ssid}":\n{m.group(2)} "{pw}"',
    txt
)
open(path, "w").write(txt)
PYEOF
        info "Updated: $NETPLAN_FILE"
        info "Changes will persist on next boot."
    elif [ -n "$NETPLAN_FILE" ]; then
        warn "$NETPLAN_FILE has no access-points section — skipping."
    fi

    info "Connected to '${NEW_SSID}' successfully."
    exit 0
    ;;
q|Q)
    exit 0
    ;;
*)
    die "Invalid choice."
    ;;
esac

info "Available interfaces:"
ip -o link show | awk -F': ' 'NR>1 {print "  " $2}'
echo ""

read -rp "Enter interface name to configure (e.g. wlan0, eth0, end0): " IFACE
[ -n "$IFACE" ] || die "Interface name cannot be empty"
ip link show "$IFACE" >/dev/null 2>&1 || die "Interface $IFACE not found"

IS_WIFI=false
if [[ "$IFACE" == wl* ]]; then
    IS_WIFI=true
    info "Wireless interface detected."
else
    info "Wired interface detected — skipping WiFi credential setup."
fi

CURRENT_IP=$(ip -o -4 addr show "$IFACE" 2>/dev/null | awk '{print $4}' | head -1)
CURRENT_GW=$(ip route show default 2>/dev/null | awk '/default/ {print $3}' | head -1)

echo ""
info "Current IP:      ${CURRENT_IP:-none}"
info "Current gateway: ${CURRENT_GW:-none}"
echo ""

read -rp "Set a static IP for this spoke? [Y/n]: " SET_STATIC
SET_STATIC="${SET_STATIC:-y}"

if [[ "${SET_STATIC,,}" != "y" ]]; then
    info "Keeping DHCP. No changes made."
    exit 0
fi

header "Static IP Configuration"

read -rp "Enter static IP with prefix (e.g. 192.168.1.69/24): " STATIC_IP
[ -n "$STATIC_IP" ] || die "IP address cannot be empty"

DEFAULT_GW="${CURRENT_GW:-192.168.1.1}"
read -rp "Gateway [${DEFAULT_GW}]: " GATEWAY
GATEWAY="${GATEWAY:-$DEFAULT_GW}"

read -rp "DNS servers (comma-separated) [${GATEWAY},8.8.8.8]: " DNS_INPUT
DNS_INPUT="${DNS_INPUT:-${GATEWAY},8.8.8.8}"

DNS_YAML=""
IFS=',' read -ra DNS_LIST <<< "$DNS_INPUT"
for DNS in "${DNS_LIST[@]}"; do
    DNS=$(echo "$DNS" | tr -d ' ')
    if [ -n "$DNS_YAML" ]; then
        DNS_YAML="${DNS_YAML}"$'\n'
    fi
    DNS_YAML="${DNS_YAML}          - ${DNS}"
done

info "Current netplan configs:"
ls /etc/netplan/ | sed 's/^/  /'
echo ""

NETPLAN_FILE=$(ls /etc/netplan/*.yaml 2>/dev/null | head -1)
read -rp "Netplan file to update [${NETPLAN_FILE}]: " INPUT_FILE
NETPLAN_FILE="${INPUT_FILE:-$NETPLAN_FILE}"
NETPLAN_FILE="${NETPLAN_FILE:-$(ls /etc/netplan/*.yaml 2>/dev/null | head -1)}"
[ -n "$NETPLAN_FILE" ] || die "No netplan file specified"

if $IS_WIFI; then
    header "WiFi Credentials"
    CURRENT_SSID=""
    if [ -f "$NETPLAN_FILE" ]; then
        CURRENT_SSID=$(python3 - "$NETPLAN_FILE" <<'PYEOF'
import sys, re
txt = open(sys.argv[1]).read()
m = re.search(r'access-points:\s*\n\s+["\']{0,1}([^"\':\n]+)["\']{0,1}:', txt)
print(m.group(1).strip() if m else "")
PYEOF
)
    fi

    KEEP_WIFI="n"
    if [ -n "$CURRENT_SSID" ]; then
        warn "Existing WiFi config found for: $CURRENT_SSID"
        read -rp "Keep existing WiFi credentials? [Y/n]: " KEEP_WIFI
        KEEP_WIFI="${KEEP_WIFI:-y}"
    fi

    if [[ "${KEEP_WIFI,,}" != "y" ]]; then
        read -rp "WiFi SSID: " WIFI_SSID
        [ -n "$WIFI_SSID" ] || die "SSID cannot be empty"
        read -rsp "WiFi password: " WIFI_PASS
        echo ""
        [ -n "$WIFI_PASS" ] || die "Password cannot be empty"
    else
        WIFI_SSID="$CURRENT_SSID"
        WIFI_PASS=$(grep -FA2 "\"${WIFI_SSID}\"" "$NETPLAN_FILE" 2>/dev/null | grep -F "password" | sed 's/^[^:]*: *//' | tr -d '"' || true)
        [ -n "$WIFI_PASS" ] || die "Could not extract WiFi password from existing config — please re-enter credentials."
    fi
fi

header "Writing Netplan Config"
NETPLAN_BACKUP_DIR="/root/.config/tinyboard/netplan-backups"
mkdir -p "$NETPLAN_BACKUP_DIR"
BACKUP_FILE=""

for OTHER_FILE in /etc/netplan/*.yaml; do
    [ "$OTHER_FILE" = "$NETPLAN_FILE" ] && continue
    BACKUP_OTHER="$NETPLAN_BACKUP_DIR/$(basename "${OTHER_FILE}").$(date +%Y%m%d%H%M%S)"
    cp "$OTHER_FILE" "$BACKUP_OTHER"
    rm "$OTHER_FILE"
    warn "Removed conflicting netplan file: $OTHER_FILE (backed up to $BACKUP_OTHER)"
done

if [ -f "$NETPLAN_FILE" ]; then
    BACKUP_FILE="$NETPLAN_BACKUP_DIR/$(basename "${NETPLAN_FILE}").$(date +%Y%m%d%H%M%S)"
    cp "$NETPLAN_FILE" "$BACKUP_FILE"
    info "Netplan config backed up to $BACKUP_FILE"
    info "To restore: cp $BACKUP_FILE $NETPLAN_FILE && netplan apply"
fi

if $IS_WIFI; then
    cat > "$NETPLAN_FILE" <<NETEOF
network:
  version: 2
  wifis:
    ${IFACE}:
      dhcp4: no
      addresses:
        - ${STATIC_IP}
      routes:
        - to: default
          via: ${GATEWAY}
      nameservers:
        addresses:
${DNS_YAML}
      access-points:
        "${WIFI_SSID}":
          password: "${WIFI_PASS}"
NETEOF
else
    cat > "$NETPLAN_FILE" <<NETEOF
network:
  version: 2
  ethernets:
    ${IFACE}:
      dhcp4: no
      addresses:
        - ${STATIC_IP}
      routes:
        - to: default
          via: ${GATEWAY}
      nameservers:
        addresses:
${DNS_YAML}
NETEOF
fi

info "Netplan config written to $NETPLAN_FILE"

header "Applying Configuration"
warn "Applying netplan config — will revert automatically if network is lost..."
netplan apply

CONNECTED=false
for i in $(seq 1 6); do
    sleep 5
    if ping -c 1 -W 2 "$GATEWAY" >/dev/null 2>&1; then
        CONNECTED=true
        break
    fi
    warn "Network check $i/6 failed, retrying..."
done

if [ "$CONNECTED" = "true" ]; then
    info "Network connectivity confirmed — config applied permanently."
else
    warn "No network connectivity detected after 30 seconds — reverting to backup config."
    if [ -n "$BACKUP_FILE" ] && [ -f "$BACKUP_FILE" ]; then
        cp "$BACKUP_FILE" "$NETPLAN_FILE"
        netplan apply
        die "Config reverted to backup. Check your settings and try again."
    else
        die "No backup found to revert to. Restore $NETPLAN_FILE manually."
    fi
fi

STATIC_ADDR="${STATIC_IP%%/*}"
echo ""
echo -e "${YELLOW}══════════════════════════════════════════${NC}"
echo -e "${YELLOW}  Network reconfigured.${NC}"
echo -e "${YELLOW}  If you are connected via SSH, your session${NC}"
echo -e "${YELLOW}  may drop. Reconnect to: ${STATIC_ADDR}${NC}"
echo -e "${YELLOW}  Then run: cd .. && ./setup.sh${NC}"
echo -e "${YELLOW}══════════════════════════════════════════${NC}"
echo ""