#!/bin/bash
# Save to: ~/.ssh/ssh-failover.sh

TARGET_HOST="$1"
CONFIG_DIR="$HOME/.ssh/hosts.d"

# Force lower-case comparison to prevent case mismatches from SSH %h token
TARGET_LOWER=$(echo "$TARGET_HOST" | tr '[:upper:]' '[:lower:]')
CONFIG_FILE="$CONFIG_DIR/proxyrule.${TARGET_LOWER}.conf"

TIMEOUT=0.5
DEBUG=${DEBUG:-false}

if [ ! -f "$CONFIG_FILE" ]; then
    echo "FAILOVER ERROR: Config file missing at $CONFIG_FILE" >&2
    exit 1
fi

source "$CONFIG_FILE"

if [ ${#ENDPOINTS[@]} -eq 0 ]; then
    echo "FAILOVER ERROR: No ENDPOINTS array found inside $CONFIG_FILE" >&2
    exit 1
fi

# Helper function to resolve hostnames to IP addresses before the timer starts
resolve_to_ip() {
    local target=$1
    # Check if it's already an IPv4 or IPv6 address to save process execution
    if [[ "$target" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ || "$target" =~ : ]]; then
        echo "$target"
    else
        # Resolve using standard system hosts databases (handles DNS, /etc/hosts, etc.)
        local ip
        ip=$(getent hosts "$target" | awk '{print $1; exit}')
        # Fall back to dig if getent isn't providing a quick single target ip
        [ -z "$ip" ] && ip=$(dig +short "$target" | tail -n1)
        # If both fail, return original target so nc can try anyway
        echo "${ip:-$target}"
    fi
}

get_tcp_ms() {
    local host=$1
    local port=$2
    local start=$EPOCHREALTIME
    
    if timeout "$TIMEOUT" nc -z "$host" "$port" 2>/dev/null; then
        local end=$EPOCHREALTIME
        awk "BEGIN {print int(($end - $start) * 1000)}"
    else
        echo "9999"
    fi
}

declare -a results
declare -a resolved_hosts
fastest_ms=9999
chosen_index=""

for ((i=0; i<${#ENDPOINTS[@]}; i++)); do
    IFS="|" read -r label host port <<< "${ENDPOINTS[$i]}"
    
    # CRITICAL BUGFIX: Strip any hidden trailing newlines (\n) or carriage returns (\r)
    port=$(echo "$port" | tr -d '\r\n[:space:]')
    
    # Pre-resolve hostnames to eliminate DNS lookup latency skew
    resolved_host=$(resolve_to_ip "$host")
    resolved_hosts[$i]=$resolved_host
    
    ms=$(get_tcp_ms "$resolved_host" "$port")
    results[$i]=$ms
    
    # if [ "$DEBUG" = true ]; then
    #     echo "Comparing $ms to $fastest_ms"
    # fi

    if [ "$ms" -lt "$fastest_ms" ]; then
        fastest_ms=$ms
        chosen_index=$i
    fi
done

# --- DYNAMIC DEBUG DISPLAY LOOP ---
if [ "$DEBUG" = true ]; then
    echo "--------------------------------------------------------" >&2
    printf "Testing routes for: \e[1;34m%s\e[0m\n" "$TARGET_HOST" >&2
    echo "--------------------------------------------------------" >&2
    printf "%-25s | %-12s | %-10s\n" "Endpoint / Host" "Status" "Latency" >&2
    echo "--------------------------------------------------------" >&2
    
    for ((i=0; i<${#ENDPOINTS[@]}; i++)); do
        IFS="|" read -r label host port <<< "${ENDPOINTS[$i]}"
        ms=${results[$i]}
        
        if [ "$ms" -eq 9999 ]; then
            printf "%-25s | \e[31m%-12s\e[0m | %-10s\n" "$label" "OFFLINE" "---" >&2
        else
            printf "%-25s | \e[32m%-12s\e[0m | %-10s\n" "$label" "ONLINE" "${ms}ms" >&2
        fi
    done
    echo "--------------------------------------------------------" >&2
fi

# 2. Route traffic to the fastest endpoint found
if [ -n "$chosen_index" ] && [ "$fastest_ms" -lt 9999 ]; then
    IFS="|" read -r label host port <<< "${ENDPOINTS[$chosen_index]}"
    port=$(echo "$port" | tr -d '\r\n[:space:]')
    
    [ "$DEBUG" = true ] && echo "-> Routing via $label ($host:$port)..." >&2
    exec nc "$host" "$port"
else
    echo "FAILOVER ERROR: All paths to $TARGET_HOST are dead." >&2
    exit 1
fi