#!/bin/bash

# Cistech VPN Connection Script with OpenConnect SSO
# Uses config.toml for configuration, openconnect-sso handles auto-connect

# Credentials from environment variables
OC_URL="${OC_URL:-}"
OC_USER="${OC_USER:-}"
OC_PASSWORD="${OC_PASSWORD:-}"
OC_TOTP_SECRET="${OC_TOTP_SECRET:-}"
OC_INTERFACE="${OC_INTERFACE:-tun0}"

CONFIG_DIR="$HOME/.config/openconnect-sso"
CONFIG_FILE="$CONFIG_DIR/config.toml"

# 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 ""
}

log() {
    local level="$1"
    local msg="$2"
    local ts=$(date '+%H:%M:%S')
    case $level in
        INFO)  echo -e "${GRAY}[$ts]${NC} ${GREEN}[INFO]${NC}  $msg" ;;
        WARN)  echo -e "${GRAY}[$ts]${NC} ${YELLOW}[WARN]${NC}  $msg" ;;
        ERROR) echo -e "${GRAY}[$ts]${NC} ${RED}[ERROR]${NC} $msg" ;;
        DEBUG) echo -e "${GRAY}[$ts]${NC} ${CYAN}[DEBUG]${NC} $msg" ;;
        *)     echo -e "${GRAY}[$ts]${NC} $msg" ;;
    esac
}

get_totp() {
    [ -n "$OC_TOTP_SECRET" ] && oathtool --totp -b "$OC_TOTP_SECRET"
}

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)
    [ -n "$iface" ] && ip addr show "$iface" 2>/dev/null | grep -oP 'inet \K[\d.]+' | head -1
}

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

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())
" 2>/dev/null && log INFO "TOTP stored in keyring"
    fi
}

setup_config() {
    mkdir -p "$CONFIG_DIR"

    # Determine browser display mode
    local browser_mode="hidden"
    [ -z "$OC_USER" ] && browser_mode="shown"

    cat > "$CONFIG_FILE" << EOF
[openconnect]
server = "$OC_URL"
interface = "$OC_INTERFACE"

[sso]
browser_display_mode = "$browser_mode"
EOF

    [ -n "$OC_USER" ] && echo "user = \"$OC_USER\"" >> "$CONFIG_FILE"

    log INFO "Config written to $CONFIG_FILE"
}

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!"
        return 1
    fi
    sysctl -w net.ipv4.ip_forward=1 >/dev/null
    iptables -t nat -C POSTROUTING -o "$vpn_iface" -j MASQUERADE 2>/dev/null || \
        iptables -t nat -A POSTROUTING -o "$vpn_iface" -j MASQUERADE
    touch /runtime/restart-routing 2>/dev/null
    log INFO "Forwarding configured"
}

kill_vpn() {
    log INFO "Stopping VPN..."
    pkill -f "openconnect" 2>/dev/null
    sleep 1
}

start_vpn() {
    log INFO "Starting OpenConnect SSO..."

    setup_keyring
    setup_config

    log INFO "Credentials:"
    echo -e "  ${CYAN}URL:  $OC_URL${NC}"
    echo -e "  ${CYAN}User: $OC_USER${NC}"
    [ -n "$OC_TOTP_SECRET" ] && echo -e "  ${CYAN}TOTP: $(get_totp)${NC}"
    echo ""

    # Run openconnect-sso (it handles reconnection)
    if [[ -n "$OC_PASSWORD" ]]; then
        echo "$OC_PASSWORD" | openconnect-sso
    else
        echo "" | openconnect-sso
    fi
}

show_totp() {
    [ -z "$OC_TOTP_SECRET" ] && { log ERROR "No TOTP secret"; return 1; }
    log INFO "Live TOTP (Ctrl+C to stop)"
    while true; do
        echo -ne "\r  TOTP: ${GREEN}$(get_totp)${NC} (${YELLOW}$((30 - $(date +%s) % 30))s${NC})  "
        sleep 1
    done
}

main_menu() {
    echo -e "${GREEN}Options:${NC}"
    echo -e "  ${CYAN}1${NC} - Start VPN"
    echo -e "  ${CYAN}2${NC} - Stop VPN"
    echo -e "  ${CYAN}3${NC} - Show TOTP"
    echo -e "  ${CYAN}4${NC} - Setup forwarding"
    echo -e "  ${CYAN}5${NC} - Check status"
    echo -e "  ${CYAN}6${NC} - Show routes"
    echo -e "  ${CYAN}q${NC} - Quit"
    echo ""
}

# Parse args
for arg in "$@"; do
    case $arg in
        -m|--menu) SKIP_AUTO=true ;;
        -s|--status) check_vpn_status; exit 0 ;;
        --help) echo "Usage: $0 [-m|--menu] [-s|--status]"; exit 0 ;;
    esac
done

# Main
print_banner
log INFO "openconnect-vpn started"

if [ "${SKIP_AUTO:-false}" = "true" ]; then
    log INFO "Menu mode - skipping auto-connect"
elif check_vpn_status; then
    setup_forwarding
else
    log INFO "Auto-starting VPN..."
    start_vpn &
    VPN_PID=$!

    # Wait for connection
    for i in {1..60}; do
        sleep 2
        if [ -n "$(get_vpn_interface)" ]; then
            log INFO "VPN connected!"
            setup_forwarding
            break
        fi
    done
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 ;;
        2) kill_vpn ;;
        3) show_totp ;;
        4) setup_forwarding ;;
        5) check_vpn_status ;;
        6) ip -4 route show ;;
        q|Q) log INFO "Goodbye!"; exit 0 ;;
    esac
done