#!/bin/bash

# Cistech VPN Connection Script with OpenConnect SSO
# Usage: ./openconnect-vpn [-c|--connect] [-d|--disconnect] [-m|--menu] [-r|--routes] [-s|--status] [--help]

# Credentials from environment variables (set by runtipi)
OC_URL="${OC_URL:-}"
OC_USER="${OC_USER:-}"
OC_PASSWORD="${OC_PASSWORD:-}"
OC_TOTP_SECRET="${OC_TOTP_SECRET:-}"
OC_SERVERCERT="${OC_SERVERCERT:-}"
OC_INTERFACE="${OC_INTERFACE:-tun0}"
OC_AUTHGROUP="${OC_AUTHGROUP:-}"
OC_USERAGENT="${OC_USERAGENT:-}"
OC_EXTRA_ARGS="${OC_EXTRA_ARGS:-}"

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

# Get current log file (changes daily)
get_log_file() {
    echo "$LOG_DIR/$(date '+%Y-%m-%d').log"
}

# Cleanup old logs
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'

# Flags
SKIP_AUTO_LOGIN=false
DO_CONNECT=false
DO_DISCONNECT=false

# Logging function
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 command with logging
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
}

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

# Show help
show_help() {
    echo -e "${CYAN}Cistech VPN Connection Script${NC}"
    echo ""
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  -c, --connect     Start VPN connection and exit"
    echo "  -d, --disconnect  Disconnect VPN and exit"
    echo "  -m, --menu        Skip auto-connect, show menu directly"
    echo "  -r, --routes      Show current routing table and exit"
    echo "  -s, --status      Show VPN and network status and exit"
    echo "  --help            Show this help message"
}

# Auto-fetch server certificate pin
get_server_cert_pin() {
    local url="$1"
    local host=$(echo "$url" | sed -E 's|https?://([^/:]+).*|\1|')
    local port=443

    log DEBUG "Fetching certificate pin from $host:$port..."

    local pin=$(echo | openssl s_client -connect "$host:$port" -servername "$host" 2>/dev/null | \
        openssl x509 -pubkey -noout 2>/dev/null | \
        openssl pkey -pubin -outform DER 2>/dev/null | \
        openssl dgst -sha256 -binary | \
        base64)

    if [[ -n "$pin" ]]; then
        echo "pin-sha256:$pin"
    else
        log ERROR "Failed to fetch certificate from $host"
        return 1
    fi
}

# Get current TOTP
get_totp() {
    if [ -n "$OC_TOTP_SECRET" ]; then
        oathtool --totp -b "$OC_TOTP_SECRET"
    fi
}

# Check if VPN is connected
get_vpn_interface() {
    ip link show "$OC_INTERFACE" 2>/dev/null | grep -q "UP" && echo "$OC_INTERFACE"
}

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
}

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

# Check VPN status
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
}

# Setup keyring with TOTP secret
setup_keyring() {
    if [[ -n "$OC_TOTP_SECRET" && -n "$OC_USER" ]]; then
        python3 -c "
import keyring
keyring.set_password('openconnect-sso', 'totp/$OC_USER', '$OC_TOTP_SECRET'.upper())
print('TOTP secret stored in keyring for $OC_USER')
"
    fi
}

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

    for pid in $(pgrep -f "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 1
    fi
}

# Disconnect VPN
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 IP forwarding
setup_forwarding() {
    log INFO "Setting up IP forwarding..."

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

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

    # NAT masquerade for traffic going through VPN
    if ! iptables -t nat -C POSTROUTING -o "$vpn_iface" -j MASQUERADE 2>/dev/null; then
        run_cmd "Adding NAT masquerade" iptables -t nat -A POSTROUTING -o "$vpn_iface" -j MASQUERADE
    fi

    # Trigger host routing service
    log INFO "Triggering host routing service..."
    touch /runtime/restart-routing
    sleep 2
    if [ ! -f /runtime/restart-routing ]; then
        log INFO "Host routing service restarted"
    else
        log WARN "Host watcher may not be running"
    fi

    log INFO "Forwarding configured"
}

# Show routing table
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}"
        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
show_network_status() {
    log INFO "Current network status:"

    echo ""
    log DEBUG "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 ""
}

# Print TOTP
show_totp() {
    if [ -z "$OC_TOTP_SECRET" ]; then
        log ERROR "TOTP secret not configured"
        return 1
    fi

    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
}

# Start VPN connection
start_vpn() {
    local do_auto_login="$1"
    log INFO "=== Starting OpenConnect VPN ==="
    echo ""

    # Kill existing
    kill_vpn_processes

    # Get server cert if not set
    if [[ -z "$OC_SERVERCERT" ]]; then
        log INFO "Fetching server certificate..."
        OC_SERVERCERT=$(get_server_cert_pin "$OC_URL")
        log DEBUG "Server cert: $OC_SERVERCERT"
    fi

    # Setup keyring for TOTP
    setup_keyring

    # Build openconnect command
    OPENCONNECT_CMD="/usr/sbin/openconnect --protocol=anyconnect --servercert $OC_SERVERCERT --interface $OC_INTERFACE --script /usr/share/vpnc-scripts/vpnc-script"
    [[ -n "$OC_AUTHGROUP" ]] && OPENCONNECT_CMD+=" --authgroup $OC_AUTHGROUP"
    [[ -n "$OC_USERAGENT" ]] && OPENCONNECT_CMD+=" --useragent $OC_USERAGENT"
    [[ -n "$OC_EXTRA_ARGS" ]] && OPENCONNECT_CMD+=" $OC_EXTRA_ARGS"

    # Determine SSO args
    if [[ -n "$OC_USER" ]]; then
        OC_SSO_ARGS="--browser-display-mode hidden -u $OC_USER"
    else
        OC_SSO_ARGS="--browser-display-mode shown"
    fi

    # Show credentials
    log INFO "Credentials:"
    echo -e "           ${CYAN}URL:      $OC_URL${NC}"
    echo -e "           ${CYAN}User:     $OC_USER${NC}"
    if [ -n "$OC_PASSWORD" ]; then
        echo -e "           ${CYAN}Password: (set)${NC}"
    fi
    if [ -n "$OC_TOTP_SECRET" ]; then
        TOTP=$(get_totp)
        echo -e "           ${CYAN}TOTP:     $TOTP${NC}"
    fi
    echo ""

    log INFO "Launching openconnect-sso..."

    # Run openconnect-sso
    if [[ -n "$OC_USER" && -n "$OC_PASSWORD" ]]; then
        echo "$OC_PASSWORD" | openconnect-sso -s "$OC_URL" $OC_SSO_ARGS -- $OPENCONNECT_CMD &
    elif [[ -n "$OC_USER" ]]; then
        echo "" | openconnect-sso -s "$OC_URL" $OC_SSO_ARGS -- $OPENCONNECT_CMD &
    else
        openconnect-sso -s "$OC_URL" $OC_SSO_ARGS -- $OPENCONNECT_CMD &
    fi
    VPN_PID=$!

    # Wait for connection
    log INFO "Waiting for VPN connection..."
    local wait_count=0
    local max_wait=120
    while [ -z "$(get_vpn_interface)" ]; do
        sleep 2
        ((wait_count+=2))
        if [ $((wait_count % 10)) -eq 0 ]; then
            log DEBUG "Still waiting... (${wait_count}s)"
        fi
        if [ $wait_count -ge $max_wait ]; then
            log ERROR "Timeout waiting for VPN after ${max_wait}s"
            return 1
        fi
        # Check if process died
        if ! kill -0 $VPN_PID 2>/dev/null; then
            log ERROR "openconnect-sso process died"
            return 1
        fi
    done

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

    sleep 3
    setup_forwarding

    log INFO "VPN setup complete"
    return 0
}

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

# Watchdog
start_watchdog() {
    log INFO "Starting VPN watchdog (check every 60s)..."

    local check_interval=60
    local reconnect_attempts=0
    local max_reconnect_attempts=3

    while true; do
        sleep $check_interval

        local vpn_iface=$(get_vpn_interface)

        if [ -n "$vpn_iface" ]; then
            reconnect_attempts=0
        else
            ((reconnect_attempts++))
            log WARN "VPN disconnected! Attempt $reconnect_attempts/$max_reconnect_attempts..."

            if [ $reconnect_attempts -le $max_reconnect_attempts ]; then
                kill_vpn_processes
                sleep 2
                start_vpn "true" &

                local wait_count=0
                while [ -z "$(get_vpn_interface)" ] && [ $wait_count -lt 120 ]; do
                    sleep 2
                    ((wait_count+=2))
                done

                if [ -n "$(get_vpn_interface)" ]; then
                    log INFO "VPN reconnected!"
                    reconnect_attempts=0
                else
                    log ERROR "Reconnect failed"
                fi
            else
                log ERROR "Max reconnect attempts reached"
                sleep 300
                reconnect_attempts=0
            fi
        fi
    done
}

# Parse args
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_LOGIN=true
                shift
                ;;
            -r|--routes)
                show_routes
                exit 0
                ;;
            -s|--status)
                print_banner
                check_vpn_status
                echo ""
                show_network_status
                exit 0
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                echo "Unknown option: $1"
                exit 1
                ;;
        esac
    done
}

# Main
parse_args "$@"
cleanup_old_logs

echo "" >> "$(get_log_file)"
log INFO "openconnect-vpn script started"
log DEBUG "OC_URL=$OC_URL"
log DEBUG "OC_USER=$OC_USER"

print_banner

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

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

# Check current status or auto-connect
if [ "$SKIP_AUTO_LOGIN" = "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 ""
    if start_vpn "true"; then
        # Start watchdog in background
        start_watchdog &
        WATCHDOG_PID=$!
        log DEBUG "Watchdog started (PID $WATCHDOG_PID)"
    fi
fi

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

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

    case $choice in
        1) start_vpn "true" && { start_watchdog & } ;;
        2) disconnect_vpn ;;
        3) show_totp ;;
        4) setup_forwarding ;;
        5) show_network_status ;;
        6) show_routes ;;
        7) kill_vpn_processes ;;
        q|Q) log INFO "Goodbye!"; exit 0 ;;
        *) ;;
    esac
done