#!/bin/bash

set -e

# Function to log messages
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# Function to check if VPN is connected
check_vpn() {
    if [ "$VPN_TYPE" = "openvpn" ]; then
        # Check if tun interface exists and has an IP
        if ip addr show tun0 >/dev/null 2>&1; then
            TUN_IP=$(ip addr show tun0 | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1)
            if [ -n "$TUN_IP" ]; then
                return 0
            fi
        fi
    elif [ "$VPN_TYPE" = "wireguard" ]; then
        # Check if wg interface exists and has an IP
        if ip addr show wg0 >/dev/null 2>&1; then
            WG_IP=$(ip addr show wg0 | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1)
            if [ -n "$WG_IP" ]; then
                return 0
            fi
        fi
    fi
    return 1
}

# Function to start OpenVPN
start_openvpn() {
    log "Starting OpenVPN with config: $VPN_CONFIG"

    if [ ! -f "$VPN_CONFIG" ]; then
        log "ERROR: VPN config file not found: $VPN_CONFIG"
        exit 1
    fi

    # Start OpenVPN in background
    openvpn --config "$VPN_CONFIG" --daemon --log /var/log/openvpn.log

    # Wait for VPN to connect
    log "Waiting for VPN connection..."
    for i in {1..30}; do
        if check_vpn; then
            log "VPN connected successfully"
            return 0
        fi
        sleep 2
    done

    log "ERROR: VPN failed to connect within 60 seconds"
    exit 1
}

# Function to start WireGuard
start_wireguard() {
    log "Starting WireGuard with config: $VPN_CONFIG"

    if [ ! -f "$VPN_CONFIG" ]; then
        log "ERROR: WireGuard config file not found: $VPN_CONFIG"
        exit 1
    fi

    # Start WireGuard
    wg-quick up "$VPN_CONFIG"

    # Wait for WireGuard to connect
    log "Waiting for WireGuard connection..."
    for i in {1..30}; do
        if check_vpn; then
            log "WireGuard connected successfully"
            return 0
        fi
        sleep 2
    done

    log "ERROR: WireGuard failed to connect within 60 seconds"
    exit 1
}

# Function to setup routing
setup_routing() {
    log "Setting up routing for VPN proxy"

    # Get the VPN interface IP
    if [ "$VPN_TYPE" = "openvpn" ]; then
        VPN_IP=$(ip addr show tun0 | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1)
        VPN_INTERFACE="tun0"
    elif [ "$VPN_TYPE" = "wireguard" ]; then
        VPN_IP=$(ip addr show wg0 | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1)
        VPN_INTERFACE="wg0"
    fi

    log "VPN Interface: $VPN_INTERFACE, IP: $VPN_IP"

    # Set up iptables rules to route traffic through VPN
    iptables -t nat -A OUTPUT -p tcp --dport 80 -j DNAT --to-destination $VPN_IP:80
    iptables -t nat -A OUTPUT -p tcp --dport 443 -j DNAT --to-destination $VPN_IP:443

    # Allow traffic through VPN interface
    iptables -A OUTPUT -o $VPN_INTERFACE -j ACCEPT
    iptables -A INPUT -i $VPN_INTERFACE -j ACCEPT
}

# Function to start Squid proxy
start_squid() {
    log "Starting Squid proxy server"

    # Initialize Squid cache
    squid -z -N -d 1

    # Start Squid
    squid -N -d 1 &
    SQUID_PID=$!

    # Wait for Squid to start
    sleep 5

    if kill -0 $SQUID_PID 2>/dev/null; then
        log "Squid proxy started successfully on port 3128"
    else
        log "ERROR: Failed to start Squid proxy"
        exit 1
    fi
}

# Function to test proxy
test_proxy() {
    log "Testing proxy connection"

    # Test if we can reach the target URL through the proxy
    if [ -n "$PROXY_URL" ]; then
        log "Testing connection to: $PROXY_URL"
        if curl -x localhost:3128 --connect-timeout 10 --max-time 30 -s -o /dev/null "$PROXY_URL"; then
            log "Proxy test successful - can reach $PROXY_URL"
        else
            log "WARNING: Proxy test failed - cannot reach $PROXY_URL"
        fi
    fi
}

# Main execution
main() {
    log "Starting VPN Proxy container"

    # Check required environment variables
    if [ -z "$PROXY_URL" ]; then
        log "ERROR: PROXY_URL environment variable is required"
        exit 1
    fi

    if [ -z "$VPN_CONFIG" ]; then
        log "ERROR: VPN_CONFIG environment variable is required"
        exit 1
    fi

    log "Configuration:"
    log "  PROXY_URL: $PROXY_URL"
    log "  VPN_CONFIG: $VPN_CONFIG"
    log "  VPN_TYPE: ${VPN_TYPE:-openvpn}"

    # Start VPN based on type
    if [ "$VPN_TYPE" = "wireguard" ]; then
        start_wireguard
    else
        start_openvpn
    fi

    # Setup routing
    setup_routing

    # Start Squid proxy
    start_squid

    # Test proxy
    test_proxy

    log "VPN Proxy is ready and listening on port 3128"
    log "Use this container as a proxy: http://localhost:3128"

    # Keep container running and monitor
    while true; do
        if ! check_vpn; then
            log "ERROR: VPN connection lost, restarting..."
            exit 1
        fi

        if ! kill -0 $SQUID_PID 2>/dev/null; then
            log "ERROR: Squid proxy died, restarting..."
            exit 1
        fi

        sleep 30
    done
}

# Handle shutdown
cleanup() {
    log "Shutting down VPN Proxy"
    if [ -n "$SQUID_PID" ]; then
        kill $SQUID_PID 2>/dev/null || true
    fi

    if [ "$VPN_TYPE" = "wireguard" ] && [ -f "$VPN_CONFIG" ]; then
        wg-quick down "$VPN_CONFIG" 2>/dev/null || true
    fi

    exit 0
}

# Set up signal handlers
trap cleanup SIGTERM SIGINT

# Run main function
main "$@"