#!/bin/bash

# Cistech VPN Connection Script with OpenConnect-SSO
# Usage: ./openconnect-vpn [-c|--connect] [-d|--disconnect] [-m|--menu] [-r|--routes] [-s|--status] [--help]
#
# Options:
#   -c, --connect     Connect to VPN and exit
#   -d, --disconnect  Disconnect VPN and exit
#   -m, --menu        Skip auto-connect, show menu directly
#   -r, --routes      Show routing table and exit
#   -s, --status      Show VPN status and exit
#   --help            Show this help message

# Credentials from environment variables (set by runtipi)
VPN_EMAIL="${VPN_EMAIL:-}"
VPN_PASSWORD="${VPN_PASSWORD:-}"
VPN_TOTP_SECRET="${VPN_TOTP_SECRET:-}"
VPN_HOST="${VPN_HOST:-}"
TARGET_IP="${TARGET_IP:-10.3.1.0}"
VPN_INTERFACE="${VPN_INTERFACE:-tun0}"
CONTAINER_NETWORK="172.30.0.0/24"

# Hardcoded test host (IBM i server)
IBMI_HOST="10.3.1.201"

# Log directory
LOG_DIR="/var/log/openconnect-vpn"
LOG_RETENTION_DAYS=7
mkdir -p "$LOG_DIR" 2>/dev/null

get_log_file() {
    echo "$LOG_DIR/$(date '+%Y-%m-%d').log"
}

cleanup_old_logs() {
    find "$LOG_DIR" -name "*.log" -type f -mtime +$LOG_RETENTION_DAYS -delete 2>/dev/null
}

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

print_banner() {
    echo -e "${CYAN}========================================${NC}"
    echo -e "${CYAN}    Cistech VPN Connection Script       ${NC}"
    echo -e "${CYAN}========================================${NC}"
    echo ""
}

# Flags
SKIP_AUTO_CONNECT=false
DO_CONNECT=false
DO_DISCONNECT=false

# Logging
log() {
    local level="$1"
    local msg="$2"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    local timestamp_short=$(date '+%H:%M:%S')
    local log_file=$(get_log_file)

    local msg_plain=$(echo -e "$msg" | sed 's/\x1b\[[0-9;]*m//g')
    echo "[$timestamp] [$level] $msg_plain" >> "$log_file"

    case $level in
        INFO)  echo -e "${GRAY}[$timestamp_short]${NC} ${GREEN}[INFO]${NC}  $msg" ;;
        WARN)  echo -e "${GRAY}[$timestamp_short]${NC} ${YELLOW}[WARN]${NC}  $msg" ;;
        ERROR) echo -e "${GRAY}[$timestamp_short]${NC} ${RED}[ERROR]${NC} $msg" ;;
        DEBUG) echo -e "${GRAY}[$timestamp_short]${NC} ${CYAN}[DEBUG]${NC} $msg" ;;
        CMD)   echo -e "${GRAY}[$timestamp_short]${NC} ${GRAY}[CMD]${NC}   $msg" ;;
        *)     echo -e "${GRAY}[$timestamp_short]${NC} $msg" ;;
    esac
}

run_cmd() {
    local desc="$1"
    shift
    log CMD "$desc: $*"
    output=$("$@" 2>&1)
    local rc=$?
    if [ -n "$output" ]; then
        echo "$output" | while IFS= read -r line; do
            echo -e "           ${GRAY}│${NC} $line"
        done
    fi
    return $rc
}

show_help() {
    echo -e "${CYAN}Cistech VPN Connection Script${NC}"
    echo ""
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -c, --connect     Connect to VPN and exit"
    echo "  -d, --disconnect  Disconnect VPN and exit"
    echo "  -m, --menu        Skip auto-connect, show menu directly"
    echo "  -r, --routes      Show routing table and exit"
    echo "  -s, --status      Show VPN status and exit"
    echo "  --help            Show this help message"
    echo ""
    echo "Menu Options:"
    echo "  1 - Connect VPN"
    echo "  2 - Disconnect VPN"
    echo "  3 - Show VPN status"
    echo "  4 - Setup IP forwarding"
    echo "  5 - Test connection"
    echo "  6 - Show network status"
    echo "  7 - Show routing table"
    echo "  8 - Show live TOTP"
    echo "  q - Quit"
}

# Fetch server certificate fingerprint
get_server_cert() {
    local host="$1"
    local server=$(echo "$host" | sed -E 's|^https?://||; s|/.*$||')
    [[ "$server" != *:* ]] && server="${server}:443"

    log DEBUG "Fetching server certificate from $server..."

    local cert_pin=$(echo | openssl s_client -connect "$server" 2>/dev/null \
        | openssl x509 -pubkey -noout 2>/dev/null \
        | openssl pkey -pubin -outform der 2>/dev/null \
        | openssl dgst -sha256 -binary \
        | openssl enc -base64)

    if [[ -n "$cert_pin" ]]; then
        echo "pin-sha256:${cert_pin}"
    else
        log WARN "Could not fetch server certificate"
        echo ""
    fi
}

# Setup keyring with credentials
setup_keyring() {
    log INFO "Setting up keyring credentials..."

    if [[ -z "$VPN_EMAIL" ]]; then
        log ERROR "VPN_EMAIL is not set"
        return 1
    fi

    python3 << PYTHON
import keyring
from keyrings.alt.file import PlaintextKeyring

keyring.set_keyring(PlaintextKeyring())

email = "$VPN_EMAIL"
password = "$VPN_PASSWORD"
totp_secret = "$VPN_TOTP_SECRET"

if password:
    keyring.set_password('openconnect-sso', email, password)
    print(f"Password stored in keyring for {email}")

if totp_secret:
    keyring.set_password('openconnect-sso', f'totp/{email}', totp_secret.upper())
    print(f"TOTP secret stored in keyring for {email}")

print("Keyring setup complete")
PYTHON

    if [ $? -eq 0 ]; then
        log INFO "Keyring credentials configured"
    else
        log ERROR "Failed to setup keyring"
        return 1
    fi
}

get_totp() {
    oathtool --totp -b "$VPN_TOTP_SECRET"
}

# Test connection to IBMI_HOST
test_connection() {
    if [[ -z "$IBMI_HOST" ]]; then
        log WARN "IBMI_HOST not set"
        return 1
    fi

    log INFO "Testing connection to $IBMI_HOST..."
    if ping -c 3 -W 3 "$IBMI_HOST" &>/dev/null; then
        log INFO "Connection test: ${GREEN}SUCCESS${NC}"
        return 0
    else
        log WARN "Connection test: ${RED}FAILED${NC} (may need manual route on host)"
        return 1
    fi
}

show_totp() {
    log INFO "Starting live TOTP display (Ctrl+C to stop)"
    echo ""
    while true; do
        TOTP=$(get_totp)
        SECONDS_LEFT=$((30 - ($(date +%s) % 30)))
        echo -ne "\r           ${CYAN}Current TOTP:${NC} ${GREEN}$TOTP${NC} (expires in ${YELLOW}${SECONDS_LEFT}s${NC})   "
        sleep 1
    done
}

get_vpn_interface() {
    local iface=$(ip link show | grep -oP 'tun\d+(?=:.*UP)' | head -1)
    if [ -z "$iface" ]; then
        iface=$(ip link show | grep -oP 'tun\d+' | head -1)
    fi
    echo "$iface"
}

get_container_ip() {
    ip addr show eth0 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1
}

get_vpn_ip() {
    local iface=$(get_vpn_interface)
    if [ -n "$iface" ]; then
        ip addr show "$iface" 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1
    fi
}

check_vpn_status() {
    local vpn_iface=$(get_vpn_interface)
    if [ -n "$vpn_iface" ]; then
        local vpn_ip=$(get_vpn_ip)
        log INFO "VPN is ${GREEN}CONNECTED${NC}"
        log DEBUG "  Interface: $vpn_iface"
        log DEBUG "  VPN IP: $vpn_ip"
        return 0
    else
        log WARN "VPN is ${RED}NOT CONNECTED${NC}"
        return 1
    fi
}

show_routes() {
    echo -e "${CYAN}========================================${NC}"
    echo -e "${CYAN}         Current Routing Table          ${NC}"
    echo -e "${CYAN}========================================${NC}"
    echo ""
    echo -e "${GREEN}IPv4 Routes:${NC}"
    ip -4 route show | while IFS= read -r line; do
        if echo "$line" | grep -qE "tun[0-9]"; then
            echo -e "  ${YELLOW}$line${NC}"
        elif echo "$line" | grep -qE "10\.|172\.(1[6-9]|2[0-9]|3[01])\.|192\.168\."; then
            echo -e "  ${CYAN}$line${NC}"
        else
            echo "  $line"
        fi
    done
    echo ""
    echo -e "${GREEN}VPN Interface:${NC}"
    local vpn_iface=$(get_vpn_interface)
    if [ -n "$vpn_iface" ]; then
        ip addr show "$vpn_iface" 2>/dev/null | grep -E "inet|link" | while IFS= read -r line; do
            echo "  $line"
        done
    else
        echo -e "  ${RED}No VPN interface found${NC}"
    fi
}

show_network_status() {
    log INFO "Current network status:"
    echo ""
    log DEBUG "Container Network Interfaces:"
    ip -4 addr show | grep -E "inet |^[0-9]+:" | while IFS= read -r line; do
        echo -e "           ${GRAY}│${NC} $line"
    done

    echo ""
    local vpn_iface=$(get_vpn_interface)
    if [ -n "$vpn_iface" ]; then
        local vpn_ip=$(get_vpn_ip)
        log INFO "VPN Status: ${GREEN}CONNECTED${NC}"
        log DEBUG "  Interface: $vpn_iface"
        log DEBUG "  VPN IP: $vpn_ip"
    else
        log WARN "VPN Status: ${RED}NOT CONNECTED${NC}"
    fi

    local container_ip=$(get_container_ip)
    if [ -n "$container_ip" ]; then
        log DEBUG "Container IP: $container_ip"
    fi

    echo ""
    log DEBUG "Default gateway:"
    ip route show default | while IFS= read -r line; do
        echo -e "           ${GRAY}│${NC} $line"
    done
    echo ""
}

kill_vpn_processes() {
    log INFO "Killing VPN processes..."
    local killed=0

    for pid in $(pgrep -x "openconnect" 2>/dev/null); do
        log DEBUG "Killing openconnect (PID $pid)"
        kill -9 "$pid" 2>/dev/null && ((killed++))
    done

    for pid in $(pgrep -f "openconnect-sso" 2>/dev/null); do
        log DEBUG "Killing openconnect-sso (PID $pid)"
        kill -9 "$pid" 2>/dev/null && ((killed++))
    done

    if [ $killed -eq 0 ]; then
        log INFO "No VPN processes were running"
    else
        log INFO "Killed $killed process(es)"
        sleep 2
    fi
}

disconnect_vpn() {
    log INFO "Disconnecting VPN..."
    kill_vpn_processes

    if check_vpn_status; then
        log WARN "VPN still appears connected"
        return 1
    fi

    log INFO "VPN disconnected"
    return 0
}

setup_forwarding() {
    log INFO "Setting up IP forwarding rules for $TARGET_IP..."

    local vpn_iface=$(get_vpn_interface)
    if [ -z "$vpn_iface" ]; then
        log ERROR "No VPN interface found! Is VPN connected?"
        return 1
    fi

    local vpn_ip=$(get_vpn_ip)
    local container_ip=$(get_container_ip)

    log DEBUG "VPN interface: $vpn_iface"
    log DEBUG "VPN IP: $vpn_ip"
    log DEBUG "Container IP: $container_ip"

    # Enable IP forwarding
    run_cmd "Enabling IP forwarding" sysctl -w net.ipv4.ip_forward=1

    # NAT masquerade for container network going through VPN
    if ! iptables -t nat -C POSTROUTING -s "$CONTAINER_NETWORK" -o "$vpn_iface" -j MASQUERADE 2>/dev/null; then
        run_cmd "Adding NAT masquerade for container network -> VPN" iptables -t nat -A POSTROUTING -s "$CONTAINER_NETWORK" -o "$vpn_iface" -j MASQUERADE
    else
        log DEBUG "NAT masquerade for container network already exists"
    fi

    # Forward rules at position 1
    iptables -D FORWARD -d "$TARGET_IP" -j ACCEPT 2>/dev/null || true
    iptables -D FORWARD -s "$TARGET_IP" -j ACCEPT 2>/dev/null || true
    iptables -D FORWARD -s "$CONTAINER_NETWORK" -j ACCEPT 2>/dev/null || true
    iptables -D FORWARD -d "$CONTAINER_NETWORK" -j ACCEPT 2>/dev/null || true

    run_cmd "Inserting forward rule (to container network)" iptables -I FORWARD 1 -d "$CONTAINER_NETWORK" -j ACCEPT
    run_cmd "Inserting forward rule (from container network)" iptables -I FORWARD 1 -s "$CONTAINER_NETWORK" -j ACCEPT
    run_cmd "Inserting forward rule (from target)" iptables -I FORWARD 1 -s "$TARGET_IP" -j ACCEPT
    run_cmd "Inserting forward rule (to target)" iptables -I FORWARD 1 -d "$TARGET_IP" -j ACCEPT

    log INFO "Forwarding rules configured"
    echo ""

    # Trigger host routing service restart
    log INFO "Triggering host routing service restart..."
    touch /runtime/restart-routing 2>/dev/null || true
    sleep 2
    if [ ! -f /runtime/restart-routing ]; then
        log INFO "Host routing service restarted"
    else
        log WARN "Host watcher may not be running (trigger file still exists)"
    fi

    log INFO "Routing configured for $TARGET_IP through VPN tunnel"
    echo ""
}

connect_vpn() {
    log INFO "=== Starting OpenConnect-SSO VPN ==="
    echo ""

    # Kill any existing VPN processes
    kill_vpn_processes

    # Reset DNS to public servers (VPN script may have overwritten it)
    log DEBUG "Resetting DNS to public servers..."
    echo "nameserver 8.8.8.8" > /etc/resolv.conf
    echo "nameserver 1.1.1.1" >> /etc/resolv.conf

    # Clean up stale tun interface
    ip link delete tun0 2>/dev/null || true

    # Validate required variables
    if [[ -z "$VPN_HOST" ]]; then
        log ERROR "VPN_HOST is not set"
        return 1
    fi

    # Setup keyring credentials
    setup_keyring

    # Fetch server certificate
    log INFO "Fetching server certificate..."
    local servercert=$(get_server_cert "$VPN_HOST")

    log INFO "Connecting to: $VPN_HOST"
    log DEBUG "Interface: $VPN_INTERFACE"

    # Build openconnect-sso command
    local sso_args=()
    sso_args+=("-s" "$VPN_HOST")

    if [[ -n "$VPN_EMAIL" ]]; then
        sso_args+=("-u" "$VPN_EMAIL")
    fi

    # Use hidden browser for automated login
    sso_args+=("--browser-display-mode" "hidden")

    # Build openconnect args
    local oc_args=()
    oc_args+=("--protocol=anyconnect")
    oc_args+=("--interface" "$VPN_INTERFACE")
    oc_args+=("--script" "/usr/share/vpnc-scripts/vpnc-script")

    if [[ -n "$servercert" ]]; then
        oc_args+=("--servercert" "$servercert")
        log DEBUG "Server cert: $servercert"
    fi

    # Launch openconnect-sso
    log INFO "Launching openconnect-sso..."
    openconnect-sso "${sso_args[@]}" -- /usr/sbin/openconnect "${oc_args[@]}" &
    OC_PID=$!
    disown $OC_PID
    log DEBUG "openconnect-sso started with PID $OC_PID"

    # Wait for VPN to connect
    log INFO "Waiting for VPN connection..."
    local wait_count=0
    local max_wait=300

    while [ -z "$(get_vpn_interface)" ]; do
        sleep 2
        ((wait_count+=2))
        if [ $((wait_count % 10)) -eq 0 ]; then
            log DEBUG "Still waiting for VPN... (${wait_count}s)"
        fi
        if [ $wait_count -ge $max_wait ]; then
            log ERROR "Timeout waiting for VPN connection after ${max_wait}s"
            return 1
        fi
    done

    log INFO "VPN connected!"
    local vpn_iface=$(get_vpn_interface)
    local vpn_ip=$(get_vpn_ip)
    log DEBUG "  Interface: $vpn_iface"
    log DEBUG "  VPN IP: $vpn_ip"

    # Wait for routes to stabilize
    log DEBUG "Waiting for routes to stabilize..."
    sleep 3

    # Setup forwarding
    setup_forwarding

    # Test connection to IBMI host
    test_connection

    # Disable screen blanking
    xset s off 2>/dev/null || true
    xset -dpms 2>/dev/null || true
    xset s noblank 2>/dev/null || true
    log DEBUG "Screen blanking disabled"

    log INFO "VPN setup complete"

    # Start watchdog in background
    start_watchdog &
    WATCHDOG_PID=$!
    log DEBUG "Watchdog started with PID $WATCHDOG_PID"

    return 0
}

# Watchdog - monitors VPN and reconnects if dropped
start_watchdog() {
    log INFO "Starting VPN watchdog (check every 60s, keepalive ping every 5min)..."

    local check_interval=60
    local keepalive_interval=300
    local last_keepalive=0
    local reconnect_attempts=0
    local max_reconnect_attempts=3

    while true; do
        sleep $check_interval

        local now=$(date +%s)
        local vpn_iface=$(get_vpn_interface)

        if [ -n "$vpn_iface" ]; then
            # VPN is up
            reconnect_attempts=0

            # Keepalive ping every 5 minutes
            if [ $((now - last_keepalive)) -ge $keepalive_interval ]; then
                if [[ -n "$IBMI_HOST" ]] && ping -c 1 -W 5 "$IBMI_HOST" &>/dev/null; then
                    log DEBUG "Keepalive ping to $IBMI_HOST successful"
                else
                    log WARN "Keepalive ping to $IBMI_HOST failed (VPN may be degraded)"
                fi
                last_keepalive=$now
            fi
        else
            # VPN is down
            ((reconnect_attempts++))
            log WARN "VPN disconnected! Reconnect attempt $reconnect_attempts/$max_reconnect_attempts..."

            if [ $reconnect_attempts -le $max_reconnect_attempts ]; then
                connect_vpn
            else
                log ERROR "Max reconnect attempts reached. Manual intervention required."
                sleep 300
                reconnect_attempts=0
            fi
        fi
    done
}

# Main menu
main_menu() {
    echo -e "${GREEN}Options:${NC}"
    echo -e "  ${CYAN}1${NC} - Connect VPN"
    echo -e "  ${CYAN}2${NC} - Disconnect VPN"
    echo -e "  ${CYAN}3${NC} - Show VPN status"
    echo -e "  ${CYAN}4${NC} - Setup IP forwarding only"
    echo -e "  ${CYAN}5${NC} - Test connection to $IBMI_HOST"
    echo -e "  ${CYAN}6${NC} - Show network status"
    echo -e "  ${CYAN}7${NC} - Show routing table"
    echo -e "  ${CYAN}8${NC} - Show live TOTP"
    echo -e "  ${CYAN}q${NC} - Quit"
    echo ""
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -c|--connect)
                DO_CONNECT=true
                shift
                ;;
            -d|--disconnect)
                DO_DISCONNECT=true
                shift
                ;;
            -m|--menu)
                SKIP_AUTO_CONNECT=true
                shift
                ;;
            -s|--status)
                print_banner
                check_vpn_status
                echo ""
                show_network_status
                exit 0
                ;;
            -r|--routes)
                show_routes
                exit 0
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                echo "Unknown option: $1"
                echo "Use --help for usage information"
                exit 1
                ;;
        esac
    done

    if [ "$DO_CONNECT" = "true" ] && [ "$DO_DISCONNECT" = "true" ]; then
        echo "Error: --connect and --disconnect are mutually exclusive"
        exit 1
    fi
}

# Main
parse_args "$@"
cleanup_old_logs

echo "" >> "$(get_log_file)"
echo "========================================" >> "$(get_log_file)"
log INFO "=== Starting OpenConnect-SSO VPN ==="
echo ""
log DEBUG "VPN_EMAIL=$VPN_EMAIL"
log DEBUG "VPN_HOST=$VPN_HOST"
log DEBUG "TARGET_IP=$TARGET_IP"
log DEBUG "VPN_TOTP_SECRET is $([ -n "$VPN_TOTP_SECRET" ] && echo 'set' || echo 'NOT SET')"

print_banner

if [ "$DO_DISCONNECT" = "true" ]; then
    disconnect_vpn
    exit $?
fi

if [ "$DO_CONNECT" = "true" ]; then
    connect_vpn
    exit $?
fi

# Auto-connect logic (unless -m flag)
if [ "$SKIP_AUTO_CONNECT" = "true" ]; then
    log INFO "Menu mode - skipping auto-connect"
elif check_vpn_status; then
    echo ""
    log INFO "VPN already connected. Setting up forwarding..."
    setup_forwarding
else
    echo ""
    log INFO "Auto-starting VPN connection..."
    echo ""
    connect_vpn
fi

# Interactive menu loop
while true; do
    echo ""
    main_menu
    echo -ne "${CYAN}Choice: ${NC}"
    read -r choice
    echo ""

    [[ -z "${choice// }" ]] && continue

    case $choice in
        1) connect_vpn ;;
        2) disconnect_vpn ;;
        3) check_vpn_status ;;
        4) setup_forwarding ;;
        5) test_connection ;;
        6) show_network_status ;;
        7) show_routes ;;
        8) show_totp ;;
        q|Q) log INFO "Goodbye!"; exit 0 ;;
        *) ;;
    esac
done