#!/bin/bash
# =============================================================================
# mGuard VPN - Multi-Server OpenVPN Entrypoint
# =============================================================================
# This script manages multiple OpenVPN server instances dynamically.
# It polls the API for active servers and starts/stops them as needed.
# =============================================================================

set -e

# Configuration
API_URL="${API_URL:-http://127.0.0.1:8000/api/internal}"
API_TIMEOUT="${API_TIMEOUT:-120}"
API_RETRY_INTERVAL="${API_RETRY_INTERVAL:-5}"
POLL_INTERVAL="${POLL_INTERVAL:-30}"

# Directories
SERVERS_DIR="/etc/openvpn/servers"
SUPERVISOR_DIR="/etc/openvpn/supervisor.d"
LOG_DIR="/var/log/openvpn"
RUN_DIR="/var/run/openvpn"

# Ensure directories exist (volumes may override Dockerfile-created dirs)
mkdir -p "$SERVERS_DIR" "$SUPERVISOR_DIR" "$LOG_DIR" "$RUN_DIR"

# State tracking
RUNNING_SERVERS=""

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

log_error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1" >&2
}

# =============================================================================
# Wait for API to be ready
# =============================================================================
wait_for_api() {
    log "Waiting for API at $API_URL..."
    local attempts=0
    local max_attempts=$((API_TIMEOUT / API_RETRY_INTERVAL))

    while [ $attempts -lt $max_attempts ]; do
        if curl -sf "$API_URL/health" > /dev/null 2>&1; then
            log "API is ready"
            return 0
        fi
        attempts=$((attempts + 1))
        log "API not ready, retrying in ${API_RETRY_INTERVAL}s... ($attempts/$max_attempts)"
        sleep $API_RETRY_INTERVAL
    done

    log_error "API not available after ${API_TIMEOUT}s"
    return 1
}

# =============================================================================
# Fetch active servers from API
# =============================================================================
fetch_active_servers() {
    curl -sf "$API_URL/vpn-servers/active" 2>/dev/null || echo "[]"
}

# =============================================================================
# Fetch CCD files for a server
# =============================================================================
fetch_ccd_files() {
    local server_id="$1"
    local ccd_dir="$2"

    log "  Fetching CCD files for server $server_id..."

    # Get all CCD files from API
    local ccd_json
    ccd_json=$(curl -sf "$API_URL/vpn-servers/$server_id/ccd" 2>/dev/null || echo '{"files":{}}')

    # Parse and create CCD files
    echo "$ccd_json" | jq -r '.files | to_entries[] | @base64' | while read -r entry; do
        local cn=$(echo "$entry" | base64 -d | jq -r '.key')
        local content=$(echo "$entry" | base64 -d | jq -r '.value')

        if [ -n "$cn" ] && [ "$cn" != "null" ]; then
            echo "$content" > "$ccd_dir/$cn"
            log "    Created CCD file: $cn"
        fi
    done

    local count=$(echo "$ccd_json" | jq -r '.count // 0')
    log "  Fetched $count CCD files"
}

# =============================================================================
# Setup a single VPN server
# =============================================================================
setup_server() {
    local server_id="$1"
    local server_name="$2"
    local port="$3"
    local protocol="$4"
    local mgmt_port="$5"

    local server_dir="$SERVERS_DIR/$server_id"

    log "Setting up server $server_id ($server_name) on port $port/$protocol"

    # Create server directory
    mkdir -p "$server_dir"

    # Fetch all required files
    log "  Fetching configuration files..."

    if ! curl -sf "$API_URL/vpn-servers/$server_id/config" > "$server_dir/server.conf"; then
        log_error "Failed to fetch server config for $server_id"
        return 1
    fi

    if ! curl -sf "$API_URL/vpn-servers/$server_id/ca" > "$server_dir/ca.crt"; then
        log_error "Failed to fetch CA for $server_id"
        return 1
    fi

    if ! curl -sf "$API_URL/vpn-servers/$server_id/cert" > "$server_dir/server.crt"; then
        log_error "Failed to fetch server cert for $server_id"
        return 1
    fi

    if ! curl -sf "$API_URL/vpn-servers/$server_id/key" > "$server_dir/server.key"; then
        log_error "Failed to fetch server key for $server_id"
        return 1
    fi
    chmod 600 "$server_dir/server.key"

    if ! curl -sf "$API_URL/vpn-servers/$server_id/dh" > "$server_dir/dh.pem"; then
        log_error "Failed to fetch DH params for $server_id"
        return 1
    fi

    # TA key is optional
    curl -sf "$API_URL/vpn-servers/$server_id/ta" > "$server_dir/ta.key" 2>/dev/null || true
    if [ -f "$server_dir/ta.key" ] && [ -s "$server_dir/ta.key" ]; then
        chmod 600 "$server_dir/ta.key"
    else
        rm -f "$server_dir/ta.key"
    fi

    # CRL
    curl -sf "$API_URL/vpn-servers/$server_id/crl" > "$server_dir/crl.pem" 2>/dev/null || true

    # Create client-config directory and fetch CCD files
    mkdir -p "$server_dir/ccd"
    fetch_ccd_files "$server_id" "$server_dir/ccd"

    # Create status file location
    touch "$RUN_DIR/openvpn-$server_id.status"

    # Create supervisor config for this server
    cat > "$SUPERVISOR_DIR/openvpn-$server_id.conf" << EOF
[program:openvpn-$server_id]
command=/usr/sbin/openvpn --config $server_dir/server.conf
directory=$server_dir
autostart=true
autorestart=true
startsecs=5
startretries=3
exitcodes=0
stopsignal=TERM
stopwaitsecs=10
stdout_logfile=$LOG_DIR/server-$server_id.log
stdout_logfile_maxbytes=10MB
stdout_logfile_backups=3
stderr_logfile=$LOG_DIR/server-$server_id-error.log
stderr_logfile_maxbytes=5MB
stderr_logfile_backups=2
EOF

    log "  Server $server_id configured successfully"
    return 0
}

# =============================================================================
# Remove a VPN server
# =============================================================================
remove_server() {
    local server_id="$1"

    log "Removing server $server_id..."

    # Stop the process via supervisorctl if supervisor is running
    if [ -S /var/run/supervisor.sock ]; then
        supervisorctl stop "openvpn-$server_id" 2>/dev/null || true
        supervisorctl remove "openvpn-$server_id" 2>/dev/null || true
    fi

    # Remove supervisor config
    rm -f "$SUPERVISOR_DIR/openvpn-$server_id.conf"

    # Keep server directory for logs, but mark as inactive
    # rm -rf "$SERVERS_DIR/$server_id"

    log "  Server $server_id removed"
}

# =============================================================================
# Notify API that server started
# =============================================================================
notify_started() {
    local server_id="$1"
    curl -sf -X POST "$API_URL/vpn-servers/$server_id/started" > /dev/null 2>&1 || true
}

# =============================================================================
# Notify API that server stopped
# =============================================================================
notify_stopped() {
    local server_id="$1"
    curl -sf -X POST "$API_URL/vpn-servers/$server_id/stopped" > /dev/null 2>&1 || true
}

# =============================================================================
# Initial setup - configure all active servers
# =============================================================================
initial_setup() {
    log "Performing initial server setup..."

    local servers_json
    servers_json=$(fetch_active_servers)

    if [ "$servers_json" = "[]" ]; then
        log "No active VPN servers found. Waiting for configuration via web UI..."
        return 0
    fi

    # Parse JSON and setup each ready server
    echo "$servers_json" | jq -c '.[] | select(.is_ready == true and .has_ca == true and .has_cert == true)' | while read -r server; do
        local id=$(echo "$server" | jq -r '.id')
        local name=$(echo "$server" | jq -r '.name')
        local port=$(echo "$server" | jq -r '.port')
        local protocol=$(echo "$server" | jq -r '.protocol')
        local mgmt_port=$(echo "$server" | jq -r '.management_port')

        if setup_server "$id" "$name" "$port" "$protocol" "$mgmt_port"; then
            RUNNING_SERVERS="$RUNNING_SERVERS $id"
        fi
    done

    log "Initial setup complete. Running servers:$RUNNING_SERVERS"
}

# =============================================================================
# Poll for changes and update servers
# =============================================================================
poll_for_changes() {
    local servers_json
    servers_json=$(fetch_active_servers)

    # Get list of ready server IDs from API
    local api_server_ids
    api_server_ids=$(echo "$servers_json" | jq -r '.[] | select(.is_ready == true and .has_ca == true and .has_cert == true) | .id' | sort)

    # Get list of currently configured servers
    local current_server_ids
    current_server_ids=$(ls -1 "$SUPERVISOR_DIR" 2>/dev/null | grep "^openvpn-" | sed 's/openvpn-\([0-9]*\)\.conf/\1/' | sort)

    # Find new servers to add
    for id in $api_server_ids; do
        if ! echo "$current_server_ids" | grep -q "^${id}$"; then
            log "New server detected: $id"
            local server
            server=$(echo "$servers_json" | jq -c ".[] | select(.id == $id)")
            local name=$(echo "$server" | jq -r '.name')
            local port=$(echo "$server" | jq -r '.port')
            local protocol=$(echo "$server" | jq -r '.protocol')
            local mgmt_port=$(echo "$server" | jq -r '.management_port')

            if setup_server "$id" "$name" "$port" "$protocol" "$mgmt_port"; then
                # Add to supervisor
                if [ -S /var/run/supervisor.sock ]; then
                    supervisorctl reread > /dev/null 2>&1 || true
                    supervisorctl add "openvpn-$id" > /dev/null 2>&1 || true
                fi
            fi
        fi
    done

    # Find servers to remove (no longer active)
    for id in $current_server_ids; do
        if ! echo "$api_server_ids" | grep -q "^${id}$"; then
            log "Server $id no longer active"
            remove_server "$id"
        fi
    done

    # Update CRL and CCD files for all running servers
    for id in $api_server_ids; do
        if [ -d "$SERVERS_DIR/$id" ]; then
            # Update CRL
            curl -sf "$API_URL/vpn-servers/$id/crl" > "$SERVERS_DIR/$id/crl.pem.new" 2>/dev/null && \
                mv "$SERVERS_DIR/$id/crl.pem.new" "$SERVERS_DIR/$id/crl.pem" || \
                rm -f "$SERVERS_DIR/$id/crl.pem.new"

            # Update CCD files (gateway routes)
            fetch_ccd_files "$id" "$SERVERS_DIR/$id/ccd"
        fi
    done
}

# =============================================================================
# Setup iptables for NAT
# =============================================================================
setup_iptables() {
    log "Setting up iptables NAT rules..."

    # Enable IP forwarding (may fail in container, that's OK if host has it enabled)
    if [ -w /proc/sys/net/ipv4/ip_forward ]; then
        echo 1 > /proc/sys/net/ipv4/ip_forward
        log "  IP forwarding enabled"
    else
        log "  IP forwarding: /proc/sys not writable (ensure host has net.ipv4.ip_forward=1)"
    fi

    # Basic NAT masquerade (adjust interface as needed)
    if iptables -t nat -C POSTROUTING -s 10.0.0.0/8 -j MASQUERADE 2>/dev/null; then
        log "  NAT masquerade rule already exists"
    elif iptables -t nat -A POSTROUTING -s 10.0.0.0/8 -j MASQUERADE 2>/dev/null; then
        log "  NAT masquerade rule added"
    else
        log "  WARNING: Could not add NAT masquerade rule"
    fi

    log "iptables setup complete"
}

# =============================================================================
# Cleanup on exit
# =============================================================================
cleanup() {
    log "Shutting down..."

    # Notify API about all servers stopping
    for conf in "$SUPERVISOR_DIR"/openvpn-*.conf; do
        if [ -f "$conf" ]; then
            local id=$(basename "$conf" | sed 's/openvpn-\([0-9]*\)\.conf/\1/')
            notify_stopped "$id"
        fi
    done

    # Stop supervisor
    if [ -S /var/run/supervisor.sock ]; then
        supervisorctl shutdown 2>/dev/null || true
    fi

    exit 0
}

# =============================================================================
# Main
# =============================================================================
main() {
    log "=== mGuard VPN Multi-Server Container Starting ==="
    log "API URL: $API_URL"
    log "Poll Interval: ${POLL_INTERVAL}s"

    # Setup signal handlers
    trap cleanup SIGTERM SIGINT

    # Wait for API
    if ! wait_for_api; then
        log_error "Cannot start without API"
        exit 1
    fi

    # Setup iptables
    setup_iptables

    # Initial setup of all servers
    initial_setup

    # Check if any servers were configured
    if [ -z "$(ls -A $SUPERVISOR_DIR 2>/dev/null)" ]; then
        log "No VPN servers configured yet."
        log "Please create a CA and VPN server via the web UI."
        log "Container will poll every ${POLL_INTERVAL}s for new servers..."

        # Start supervisor anyway (will have no programs)
        /usr/bin/supervisord -c /etc/supervisord.conf &
        SUPERVISOR_PID=$!

        # Polling loop waiting for first server
        while true; do
            sleep $POLL_INTERVAL
            poll_for_changes

            # Check if supervisor is still running
            if ! kill -0 $SUPERVISOR_PID 2>/dev/null; then
                log_error "Supervisor died unexpectedly"
                exit 1
            fi
        done
    else
        log "Starting supervisord with configured servers..."

        # Start supervisor in background
        /usr/bin/supervisord -c /etc/supervisord.conf &
        SUPERVISOR_PID=$!

        # Wait a moment for supervisor to start
        sleep 3

        # Notify API about started servers
        for conf in "$SUPERVISOR_DIR"/openvpn-*.conf; do
            if [ -f "$conf" ]; then
                local id=$(basename "$conf" | sed 's/openvpn-\([0-9]*\)\.conf/\1/')
                # Check if actually running
                if supervisorctl status "openvpn-$id" 2>/dev/null | grep -q RUNNING; then
                    notify_started "$id"
                fi
            fi
        done

        log "All servers started. Entering polling mode..."

        # Polling loop for changes
        while true; do
            sleep $POLL_INTERVAL
            poll_for_changes

            # Check if supervisor is still running
            if ! kill -0 $SUPERVISOR_PID 2>/dev/null; then
                log_error "Supervisor died unexpectedly"
                exit 1
            fi
        done
    fi
}

# Run main
main