#!/bin/bash
#
# Lilith Platform - VPN Access Setup
#
# Sets up VPN/SOCKS5 proxy access for development environment.
# Enables access to status.atlilith.com (IP whitelisted, 403 without VPN).
#
# Usage:
#   ./setup-vpn-access.sh              # Interactive setup
#   ./setup-vpn-access.sh --check      # Check current status only
#   ./setup-vpn-access.sh --socks5     # Setup SOCKS5 tunnel only
#   ./setup-vpn-access.sh --systemd    # Install systemd services
#
# Requirements:
#   - WireGuard (for primary VPN) OR SSH access (for SOCKS5 fallback)
#   - autossh (installed automatically if missing)
#

# Don't exit on individual check failures
# set -e

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

# Configuration
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INFRA_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
PROJECT_ROOT="$(cd "$INFRA_DIR/.." && pwd)"
VPN_HOST="${VPN_HOST:-vpn.1984.nasty.sh}"
STATUS_URL="https://status.atlilith.com"
SOCKS_PORT="${SOCKS_PORT:-1080}"

# Logging
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_header() { echo -e "\n${CYAN}═══ $1 ═══${NC}\n"; }

# Header
show_banner() {
    echo -e "${CYAN}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║        Lilith Platform - VPN Access Setup                   ║"
    echo "║        Access status.atlilith.com from whitelisted IP       ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

# Check if command exists
command_exists() {
    command -v "$1" &>/dev/null
}

# Check WireGuard status
check_wireguard() {
    log_header "WireGuard Status"

    if ! command_exists wg; then
        log_warn "WireGuard not installed"
        echo "  Install with: sudo dnf install wireguard-tools  # Fedora"
        echo "                sudo apt install wireguard         # Debian/Ubuntu"
        return 1
    fi

    log_success "WireGuard installed: $(wg --version 2>&1 | head -1)"

    if ip link show wg0 &>/dev/null; then
        local wg_ip=$(ip addr show wg0 2>/dev/null | grep -oP 'inet \K[\d.]+')
        log_success "WireGuard interface wg0 is UP"
        log_info "  VPN IP: $wg_ip"

        # Show peer info
        if sudo wg show wg0 2>/dev/null | grep -q "peer"; then
            local last_handshake=$(sudo wg show wg0 2>/dev/null | grep "latest handshake" | awk '{print $3, $4, $5}')
            log_info "  Last handshake: ${last_handshake:-unknown}"
        fi
        return 0
    else
        log_warn "WireGuard interface wg0 not configured"
        return 1
    fi
}

# Check SOCKS5 tunnel status
check_socks5() {
    log_header "SOCKS5 Tunnel Status"

    if ! command_exists autossh; then
        log_warn "autossh not installed (required for persistent SOCKS5 tunnel)"
        echo "  Install with: sudo dnf install autossh  # Fedora"
        echo "                sudo apt install autossh   # Debian/Ubuntu"
    else
        log_success "autossh installed"
    fi

    # Check if tunnel is running
    if pgrep -f "ssh.*-D.*$SOCKS_PORT" &>/dev/null; then
        log_success "SOCKS5 tunnel running on port $SOCKS_PORT"
        return 0
    else
        log_info "SOCKS5 tunnel not running"
        return 1
    fi
}

# Check status.atlilith.com access
check_status_access() {
    log_header "status.atlilith.com Access Test"

    # Direct access (via WireGuard VPN)
    log_info "Testing direct access..."
    local direct_code=$(curl -s -o /dev/null -w "%{http_code}" --connect-timeout 5 "$STATUS_URL" 2>/dev/null || echo "000")

    if [ "$direct_code" = "200" ]; then
        log_success "Direct access: HTTP $direct_code (VPN working)"
        return 0
    elif [ "$direct_code" = "403" ]; then
        log_warn "Direct access: HTTP 403 (IP not whitelisted, need VPN)"
    else
        log_error "Direct access: HTTP $direct_code (connection issue)"
    fi

    # SOCKS5 access (if tunnel running)
    if pgrep -f "ssh.*-D.*$SOCKS_PORT" &>/dev/null; then
        log_info "Testing SOCKS5 proxy access..."
        local socks_code=$(curl -s -o /dev/null -w "%{http_code}" --socks5 "localhost:$SOCKS_PORT" --connect-timeout 5 "$STATUS_URL" 2>/dev/null || echo "000")

        if [ "$socks_code" = "200" ]; then
            log_success "SOCKS5 proxy access: HTTP $socks_code (tunnel working)"
            return 0
        else
            log_error "SOCKS5 proxy access: HTTP $socks_code"
        fi
    fi

    return 1
}

# Start SOCKS5 tunnel
start_socks5_tunnel() {
    log_header "Starting SOCKS5 Tunnel"

    # Check if already running
    if pgrep -f "ssh.*-D.*$SOCKS_PORT" &>/dev/null; then
        log_warn "SOCKS5 tunnel already running on port $SOCKS_PORT"
        return 0
    fi

    # Check SSH access
    log_info "Testing SSH access to $VPN_HOST..."
    if ! ssh -o ConnectTimeout=5 -o BatchMode=yes "$VPN_HOST" "echo connected" &>/dev/null; then
        log_error "Cannot connect to $VPN_HOST via SSH"
        echo ""
        echo "Ensure you have SSH key access configured:"
        echo "  1. Add your SSH key to the VPN server"
        echo "  2. Add entry to ~/.ssh/config:"
        echo ""
        echo "  Host vpn.1984.nasty.sh"
        echo "      User root"
        echo "      IdentityFile ~/.ssh/id_ed25519"
        echo ""
        return 1
    fi

    log_success "SSH access to $VPN_HOST verified"

    # Check for autossh
    if command_exists autossh; then
        log_info "Starting persistent SOCKS5 tunnel with autossh..."
        autossh -M 0 -f -N -D "$SOCKS_PORT" \
            -o "ServerAliveInterval=30" \
            -o "ServerAliveCountMax=3" \
            -o "ExitOnForwardFailure=yes" \
            "$VPN_HOST"
    else
        log_info "Starting SOCKS5 tunnel with ssh (install autossh for persistence)..."
        ssh -f -N -D "$SOCKS_PORT" \
            -o "ServerAliveInterval=30" \
            -o "ServerAliveCountMax=3" \
            "$VPN_HOST"
    fi

    sleep 2

    if pgrep -f "ssh.*-D.*$SOCKS_PORT" &>/dev/null; then
        log_success "SOCKS5 tunnel started on localhost:$SOCKS_PORT"
        echo ""
        echo "Configure your browser/applications to use:"
        echo "  Proxy Type: SOCKS5"
        echo "  Host: localhost (127.0.0.1)"
        echo "  Port: $SOCKS_PORT"
        echo ""
        echo "Test with: curl --socks5 localhost:$SOCKS_PORT $STATUS_URL"
        return 0
    else
        log_error "Failed to start SOCKS5 tunnel"
        return 1
    fi
}

# Stop SOCKS5 tunnel
stop_socks5_tunnel() {
    log_header "Stopping SOCKS5 Tunnel"

    if pkill -f "ssh.*-D.*$SOCKS_PORT" 2>/dev/null; then
        log_success "SOCKS5 tunnel stopped"
    else
        log_info "No SOCKS5 tunnel running"
    fi
}

# Install systemd services
install_systemd_services() {
    log_header "Installing Systemd Services"

    if [ "$EUID" -ne 0 ]; then
        log_error "Must run as root to install systemd services"
        echo "  Run: sudo $0 --systemd"
        return 1
    fi

    local user="${SUDO_USER:-$USER}"
    local systemd_dir="$INFRA_DIR/systemd"

    # Check source files exist
    if [ ! -f "$systemd_dir/vpn-socks5-tunnel.service" ]; then
        log_error "Systemd service files not found in $systemd_dir"
        return 1
    fi

    # Update service files with correct paths
    log_info "Creating service files with correct paths..."

    # SOCKS5 tunnel service
    cat > /etc/systemd/system/vpn-socks5-tunnel.service << EOF
[Unit]
Description=VPN SOCKS5 Tunnel to nasty.sh
After=network-online.target
Wants=network-online.target
ConditionPathExists=!/sys/class/net/wg0
Documentation=file://$INFRA_DIR/VPN_AUTO_CONNECTION.md

[Service]
Type=simple
User=$user
Group=$user
Environment="AUTOSSH_GATETIME=0"
Environment="AUTOSSH_PORT=0"
ExecStart=/usr/bin/autossh -M 0 -N -D $SOCKS_PORT -o "ServerAliveInterval=30" -o "ServerAliveCountMax=3" -o "ExitOnForwardFailure=yes" -o "StrictHostKeyChecking=no" $VPN_HOST
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal

# Security hardening
NoNewPrivileges=true
PrivateTmp=true

[Install]
WantedBy=default.target
EOF
    log_success "Created vpn-socks5-tunnel.service"

    # Health monitor service
    cat > /etc/systemd/system/vpn-health-monitor.service << EOF
[Unit]
Description=VPN Connection Health Monitor
Documentation=file://$INFRA_DIR/VPN_AUTO_CONNECTION.md

[Service]
Type=oneshot
ExecStart=$SCRIPT_DIR/vpn-health-check.sh
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF
    log_success "Created vpn-health-monitor.service"

    # Health monitor timer
    cat > /etc/systemd/system/vpn-health-monitor.timer << EOF
[Unit]
Description=VPN Health Monitor Timer (runs every 60 seconds)
Documentation=file://$INFRA_DIR/VPN_AUTO_CONNECTION.md

[Timer]
OnBootSec=30s
OnUnitActiveSec=60s
AccuracySec=5s

[Install]
WantedBy=timers.target
EOF
    log_success "Created vpn-health-monitor.timer"

    # Reload systemd
    systemctl daemon-reload
    log_success "Systemd daemon reloaded"

    echo ""
    log_info "Enable services with:"
    echo "  sudo systemctl enable --now vpn-socks5-tunnel.service"
    echo "  sudo systemctl enable --now vpn-health-monitor.timer"
}

# Full status check
check_all() {
    show_banner
    check_wireguard
    check_socks5
    check_status_access

    echo ""
    log_header "Summary"

    if check_status_access &>/dev/null; then
        log_success "status.atlilith.com is accessible"
    else
        log_error "status.atlilith.com is NOT accessible"
        echo ""
        echo "Options to fix:"
        echo "  1. Configure WireGuard VPN (see $INFRA_DIR/VPN_SETUP.md)"
        echo "  2. Start SOCKS5 tunnel: $0 --socks5"
    fi
}

# Show usage
show_usage() {
    show_banner
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options:"
    echo "  --check       Check current VPN/SOCKS5 status"
    echo "  --socks5      Start SOCKS5 tunnel"
    echo "  --stop        Stop SOCKS5 tunnel"
    echo "  --systemd     Install systemd services (requires sudo)"
    echo "  --help        Show this help message"
    echo ""
    echo "Environment Variables:"
    echo "  VPN_HOST      VPN server hostname (default: vpn.1984.nasty.sh)"
    echo "  SOCKS_PORT    SOCKS5 proxy port (default: 1080)"
    echo ""
    echo "Examples:"
    echo "  $0 --check              # Check if VPN is working"
    echo "  $0 --socks5             # Start SOCKS5 proxy tunnel"
    echo "  sudo $0 --systemd       # Install auto-start services"
}

# Main
case "${1:-}" in
    --check)
        check_all
        ;;
    --socks5)
        start_socks5_tunnel
        ;;
    --stop)
        stop_socks5_tunnel
        ;;
    --systemd)
        install_systemd_services
        ;;
    --help|-h)
        show_usage
        ;;
    "")
        check_all
        ;;
    *)
        log_error "Unknown option: $1"
        show_usage
        exit 1
        ;;
esac