#!/bin/sh
# nordvpn-select.sh
# Picks the lowest-latency NordVPN WireGuard server in Osaka.
# Falls back to Tokyo only if Tokyo is significantly better (TOKYO_THRESHOLD ms).
#
# Requirements: curl, jq
#   apk add jq curl          (OpenWrt 25.12+)
#   opkg update && opkg install jq curl   (OpenWrt 24.10 and earlier)
#
# Usage:
#   chmod +x /etc/nordvpn-select.sh
#   /etc/nordvpn-select.sh             # run normally
#   /etc/nordvpn-select.sh --dry-run   # test: select server but don't apply
#
# Cron example (Services > Scheduled Tasks in LuCI):
#   0 */6 * * * /etc/nordvpn-select.sh
# ─────────────────────────────────────────────────────────────────────────────

# ── Configuration ────────────────────────────────────────────────────────────
WG_INTERFACE="wireguard"   # UCI interface name
PING_COUNT=4               # pings per server
PING_TIMEOUT=3             # seconds per ping attempt
CANDIDATE_LIMIT=5          # how many servers to ping per city (lowest load first)
TOKYO_THRESHOLD=20         # only pick Tokyo if it's this many ms faster than Osaka
STATE_FILE="/tmp/nordvpn-select.state"
SERVER_CACHE="/tmp/nordvpn-servers.json"
CACHE_TTL=3600             # reuse server list if fresher than this many seconds
# ─────────────────────────────────────────────────────────────────────────────

log() { logger -t nordvpn-select "$*"; echo "[$(date '+%H:%M:%S')] $*" >&2; }

# ── Argument parsing ────────────────────────────────────────────────────────────
DRY_RUN=0
for arg in "$@"; do
    case "$arg" in
        --dry-run|-n) DRY_RUN=1 ;;
        *) log "ERROR: Unknown argument '$arg'"; exit 1 ;;
    esac
done
[ "$DRY_RUN" -eq 1 ] && log "*** DRY RUN — no changes will be applied ***"

# ── Dependency check ─────────────────────────────────────────────────────────
for cmd in curl jq ping uci wg; do
   command -v "$cmd" >/dev/null 2>&1 || {
       log "ERROR: '$cmd' not found. Run: apk add $cmd"
       exit 1
   }
done

# ── Fetch full server list (with caching) ────────────────────────────────────
fetch_servers() {
    # Reuse cache if recent enough
    if [ -f "$SERVER_CACHE" ]; then
        age=$(( $(date +%s) - $(date -r "$SERVER_CACHE" +%s 2>/dev/null || echo 0) ))
        if [ "$age" -lt "$CACHE_TTL" ]; then
            log "Using cached server list (${age}s old)"
            return 0
        fi
    fi

    log "Fetching NordVPN server list..."
    curl -sf --max-time 120 'https://api.nordvpn.com/v1/servers?limit=50000' \
        > "$SERVER_CACHE.tmp" || {
        log "ERROR: Failed to fetch server list"
        # Fall back to stale cache if available
        [ -f "$SERVER_CACHE" ] && { log "Using stale cache"; mv "$SERVER_CACHE" "$SERVER_CACHE"; return 0; }
        return 1
    }

    # Sanity check — must be a non-empty JSON array
    count=$(jq 'length' "$SERVER_CACHE.tmp" 2>/dev/null)
    if [ -z "$count" ] || [ "$count" -eq 0 ]; then
        log "ERROR: Server list is empty or invalid JSON"
        rm -f "$SERVER_CACHE.tmp"
        return 1
    fi

    mv "$SERVER_CACHE.tmp" "$SERVER_CACHE"
    log "Fetched $count servers"
}

# ── Get candidates for a city, sorted by load, limited to CANDIDATE_LIMIT ───
# Outputs lines of: ip|pubkey|hostname|load
get_city_candidates() {
    city="$1"
    jq -r --arg city "$city" '
        [
            .[]
            | select(any(.technologies[]; .identifier == "wireguard_udp"))
            | select(any(.locations[]; .country.city.name == $city))
            | select(.status == "online")
        ]
        | sort_by(.load)
        | .[]
        | . as $s
        | .technologies[]
        | select(.identifier == "wireguard_udp")
        | .metadata[]
        | select(.name == "public_key")
        | [$s.station, .value, $s.hostname, ($s.load | tostring)]
        | join("|")
    ' "$SERVER_CACHE" | head -n "$CANDIDATE_LIMIT"
}

# ── Ping test a set of candidates, return best as ip|pubkey|hostname|load|lat ─
# Reads candidate lines from a file to avoid pipe subshell scoping issues
find_best_in_file() {
    city_name="$1"
    candidate_file="$2"

    if [ ! -s "$candidate_file" ]; then
        log "  No online $city_name servers found"
        return 1
    fi

    best_lat=9999
    best_line=""

    while IFS="|" read -r ip pubkey hostname load; do
        [ -z "$ip" ] && continue
        lat=$(ping -c "$PING_COUNT" -W "$PING_TIMEOUT" -q "$ip" 2>/dev/null \
              | awk -F'/' '/^rtt|^round-trip/ { printf "%.0f", $5 }')
        lat=${lat:-9999}
        log "  $hostname ($ip) load=${load}% latency=${lat}ms"
        if [ "$lat" -lt "$best_lat" ]; then
            best_lat=$lat
            best_line="${ip}|${pubkey}|${hostname}|${load}|${lat}"
        fi
    done < "$candidate_file"

    if [ -z "$best_line" ] || [ "$best_lat" -eq 9999 ]; then
        log "  No reachable $city_name servers"
        return 1
    fi

    echo "$best_line"
    return 0
}

# ── Main ─────────────────────────────────────────────────────────────────────
log "=== NordVPN server selection started ==="

fetch_servers || exit 1

# Build candidate files (files, not pipes, to avoid subshell scoping)
get_city_candidates "Osaka" > /tmp/nord_osaka.txt
get_city_candidates "Tokyo" > /tmp/nord_tokyo.txt

osaka_count=$(wc -l < /tmp/nord_osaka.txt | tr -d ' ')
tokyo_count=$(wc -l < /tmp/nord_tokyo.txt | tr -d ' ')
log "Candidates — Osaka: $osaka_count  Tokyo: $tokyo_count"

# Ping test each city
log "Testing Osaka..."
osaka_result=$(find_best_in_file "Osaka" /tmp/nord_osaka.txt)
osaka_ok=$?

log "Testing Tokyo..."
tokyo_result=$(find_best_in_file "Tokyo" /tmp/nord_tokyo.txt)
tokyo_ok=$?

rm -f /tmp/nord_osaka.txt /tmp/nord_tokyo.txt

# Parse results
if [ $osaka_ok -eq 0 ]; then
    IFS="|" read -r osaka_ip osaka_pubkey osaka_host osaka_load osaka_lat <<EOF
$osaka_result
EOF
fi
if [ $tokyo_ok -eq 0 ]; then
    IFS="|" read -r tokyo_ip tokyo_pubkey tokyo_host tokyo_load tokyo_lat <<EOF
$tokyo_result
EOF
fi

# Choose city
if [ $osaka_ok -ne 0 ] && [ $tokyo_ok -ne 0 ]; then
    log "ERROR: No reachable servers in Osaka or Tokyo. Keeping current config."
    exit 1
elif [ $osaka_ok -ne 0 ]; then
    chosen_ip=$tokyo_ip; chosen_pubkey=$tokyo_pubkey; chosen_host=$tokyo_host
    chosen_load=$tokyo_load; chosen_lat=$tokyo_lat; chosen_city="Tokyo"
    log "Osaka unreachable — using Tokyo"
elif [ $tokyo_ok -ne 0 ]; then
    chosen_ip=$osaka_ip; chosen_pubkey=$osaka_pubkey; chosen_host=$osaka_host
    chosen_load=$osaka_load; chosen_lat=$osaka_lat; chosen_city="Osaka"
    log "Tokyo unreachable — using Osaka"
else
    diff=$((osaka_lat - tokyo_lat))
    if [ "$diff" -gt "$TOKYO_THRESHOLD" ]; then
        chosen_ip=$tokyo_ip; chosen_pubkey=$tokyo_pubkey; chosen_host=$tokyo_host
        chosen_load=$tokyo_load; chosen_lat=$tokyo_lat; chosen_city="Tokyo"
        log "Tokyo is ${diff}ms faster (threshold ${TOKYO_THRESHOLD}ms) — switching to Tokyo"
    else
        chosen_ip=$osaka_ip; chosen_pubkey=$osaka_pubkey; chosen_host=$osaka_host
        chosen_load=$osaka_load; chosen_lat=$osaka_lat; chosen_city="Osaka"
        if [ "$diff" -le 0 ]; then
            log "Osaka preferred ($(( -diff ))ms faster than Tokyo)"
        else
            log "Osaka preferred (Tokyo only ${diff}ms faster, threshold ${TOKYO_THRESHOLD}ms)"
        fi
    fi
fi

log "Selected: $chosen_host [$chosen_city] IP=$chosen_ip load=${chosen_load}% latency=${chosen_lat}ms"

# Skip restart if server unchanged
current_ip=$(uci -q get "network.@wireguard_${WG_INTERFACE}[0].endpoint_host" 2>/dev/null)
if [ "$current_ip" = "$chosen_ip" ]; then
    log "Server unchanged — no restart needed."
    echo "$chosen_host|$chosen_ip|$chosen_city|${chosen_lat}ms" > "$STATE_FILE"
    exit 0
fi

# Apply (skipped in dry-run mode)
if [ "$DRY_RUN" -eq 1 ]; then
    log "=== Dry run complete. No changes made."
    echo ""
    echo "# ── UCI commands to apply ────────────────────────────────"
    echo "uci set network.@wireguard_${WG_INTERFACE}[0].endpoint_host='$chosen_ip'"
    echo "uci set network.@wireguard_${WG_INTERFACE}[0].public_key='$chosen_pubkey'"
    echo "uci commit network"
    echo "ifdown $WG_INTERFACE && sleep 2 && ifup $WG_INTERFACE"
    echo ""
    echo "# ── LuCI / manual WireGuard peer params ─────────────────"
    echo "  Server:     $chosen_host ($chosen_city)"
    echo "  Endpoint:   $chosen_ip:51820"
    echo "  Public key: $chosen_pubkey"
    echo "  AllowedIPs: 0.0.0.0/0"
    echo "  Keep-alive: 25"
    echo ""
    exit 0
fi

log "Updating peer: ${current_ip:-<none>} → $chosen_ip"
uci set "network.@wireguard_${WG_INTERFACE}[0].endpoint_host=$chosen_ip"
uci set "network.@wireguard_${WG_INTERFACE}[0].public_key=$chosen_pubkey"
uci commit network

log "Reconnecting interface '$WG_INTERFACE'..."
ifdown "$WG_INTERFACE" 2>/dev/null
sleep 2
ifup "$WG_INTERFACE"

# Verify handshake
sleep 5
handshake=$(wg show "$WG_INTERFACE" latest-handshakes 2>/dev/null | awk '{print $2}')
now=$(date +%s)
if [ -n "$handshake" ] && [ "$handshake" != "0" ]; then
    age=$((now - handshake))
    if [ "$age" -lt 30 ]; then
        log "✓ Handshake confirmed (${age}s ago)"
    else
        log "WARNING: Last handshake was ${age}s ago — tunnel may not be up yet"
    fi
else
    log "WARNING: No handshake detected — check: wg show $WG_INTERFACE"
fi

echo "$chosen_host|$chosen_ip|$chosen_city|${chosen_lat}ms" > "$STATE_FILE"
log "=== Done: $chosen_host ($chosen_city, ${chosen_lat}ms) ==="