#!/bin/bash
set -e

VERSION_FILE="/etc/wireguard/VERSION"
if [ -f "$VERSION_FILE" ]; then
    VERSION=$(cat "$VERSION_FILE")
else
    VERSION="unknown"
fi
echo "[init] VPN Container Entrypoint ${VERSION}"

INTERFACE="wg0"
MOUNT_CONFIG="/etc/wireguard/${INTERFACE}.conf"
CONFIG_DIR="/run/wireguard"
CONFIG_FILE="${CONFIG_DIR}/${INTERFACE}.conf"
CHECK_INTERVAL="${HEALTH_CHECK_INTERVAL:-10}"
CHECK_HOST="${HEALTH_CHECK_HOST:-1.1.1.1}"

# ──────────────────────────────────────────────
# Validate config exists, copy to writable location
# ──────────────────────────────────────────────
if [ ! -f "$MOUNT_CONFIG" ]; then
    echo "[error] WireGuard config not found at ${MOUNT_CONFIG}"
    echo "[error] Mount your config file: -v /path/to/your.conf:/etc/wireguard/wg0.conf:ro"
    exit 1
fi

mkdir -p "$CONFIG_DIR"
cp "$MOUNT_CONFIG" "$CONFIG_FILE"
chmod 600 "$CONFIG_FILE"

# Extract endpoint IP and port from the config
VPN_ENDPOINT=$(grep -i '^Endpoint' "$CONFIG_FILE" | head -1 | sed 's/.*= *//;s/:.*//;s/ //g')
VPN_PORT=$(grep -i '^Endpoint' "$CONFIG_FILE" | head -1 | sed 's/.*://;s/ //g')
# Extract address
VPN_ADDRESS=$(grep -i '^Address' "$CONFIG_FILE" | head -1 | sed 's/.*= *//;s/ //g')

if [ -z "$VPN_ENDPOINT" ] || [ -z "$VPN_PORT" ]; then
    echo "[error] Could not parse Endpoint from ${CONFIG_FILE}"
    exit 1
fi

echo "[init] Config: ${CONFIG_FILE}"
echo "[init] Endpoint: ${VPN_ENDPOINT}:${VPN_PORT}"
echo "[init] Address: ${VPN_ADDRESS}"

# ──────────────────────────────────────────────
# Kill switch: only allow traffic through wg0
# ──────────────────────────────────────────────
setup_killswitch() {
    echo "[killswitch] Setting up iptables kill switch..."

    # Flush existing rules
    iptables -F
    iptables -X
    iptables -t nat -F

    # Default policy: DROP everything
    iptables -P INPUT DROP
    iptables -P FORWARD DROP
    iptables -P OUTPUT DROP

    # Allow loopback
    iptables -A INPUT  -i lo -j ACCEPT
    iptables -A OUTPUT -o lo -j ACCEPT

    # Allow traffic to/from VPN endpoint (needed to establish tunnel)
    iptables -A OUTPUT -d "$VPN_ENDPOINT" -p udp --dport "$VPN_PORT" -j ACCEPT
    iptables -A INPUT  -s "$VPN_ENDPOINT" -p udp --sport "$VPN_PORT" -j ACCEPT

    # Allow all traffic through the WireGuard interface
    iptables -A INPUT  -i "$INTERFACE" -j ACCEPT
    iptables -A OUTPUT -o "$INTERFACE" -j ACCEPT

    # Allow DNS (VPN DNS servers are routed through wg0; allow before routing decision)
    iptables -A OUTPUT -p udp --dport 53 -j ACCEPT
    iptables -A OUTPUT -p tcp --dport 53 -j ACCEPT
    iptables -A INPUT  -p udp --sport 53 -j ACCEPT
    iptables -A INPUT  -p tcp --sport 53 -j ACCEPT

    # Allow DHCP (for container networking)
    iptables -A OUTPUT -p udp --dport 67:68 -j ACCEPT
    iptables -A INPUT  -p udp --sport 67:68 -j ACCEPT

    # Allow established/related connections
    iptables -A INPUT  -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    iptables -A OUTPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

    # ── Allow incoming connections to exposed service ports (e.g. app on 8000) ──
    # LOCAL_PORTS can be set as env var, e.g. "8000,8080,3000"
    if [ -n "${LOCAL_PORTS:-}" ]; then
        for port in $(echo "$LOCAL_PORTS" | tr ',' ' '); do
            echo "[killswitch] Allowing incoming traffic on port ${port}"
            iptables -A INPUT  -p tcp --dport "$port" -j ACCEPT
            iptables -A OUTPUT -p tcp --sport "$port" -j ACCEPT
        done
    fi

    # ── FORWARDING (so other containers can use this VPN) ──
    iptables -A FORWARD -i eth0 -o "$INTERFACE" -j ACCEPT
    iptables -A FORWARD -i "$INTERFACE" -o eth0 -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

    # NAT: masquerade traffic from other containers going out through wg0
    iptables -t nat -A POSTROUTING -o "$INTERFACE" -j MASQUERADE

    echo "[killswitch] Kill switch active. Traffic blocked if VPN drops."
}

# ──────────────────────────────────────────────
# Enable IP forwarding so other containers can route through us
# ──────────────────────────────────────────────
enable_forwarding() {
    echo "[init] Enabling IP forwarding..."
    if cat /proc/sys/net/ipv4/ip_forward 2>/dev/null | grep -q 1; then
        echo "[init] IP forwarding already enabled."
    elif echo 1 > /proc/sys/net/ipv4/ip_forward 2>/dev/null; then
        echo "[init] IP forwarding enabled via /proc."
    else
        echo "[init] /proc read-only — relying on --sysctl net.ipv4.ip_forward=1"
    fi
}

# ──────────────────────────────────────────────
# Start WireGuard manually (no wg-quick, avoids sysctl issues)
# ──────────────────────────────────────────────
start_vpn() {
    echo "[vpn] Starting WireGuard interface ${INTERFACE}..."

    # Create the interface
    ip link add "$INTERFACE" type wireguard

    # Apply the WireGuard config (keys, peer, endpoint)
    # Filter out wg-quick directives that wg setconf doesn't understand
    wg setconf "$INTERFACE" <(grep -v -i '^\(Address\|DNS\|MTU\|Table\|PreUp\|PostUp\|PreDown\|PostDown\|SaveConfig\)' "$CONFIG_FILE")

    # Log public key so it can be verified against the server's peer list
    local PUBKEY
    PUBKEY=$(wg show "$INTERFACE" public-key 2>/dev/null || echo "unknown")
    echo "[vpn] Public key: ${PUBKEY}"

    # Assign the address
    ip -4 address add "$VPN_ADDRESS" dev "$INTERFACE"

    # Set MTU and bring up
    ip link set mtu 1420 up dev "$INTERFACE"

    # ── fwmark-based routing (mirrors wg-quick behavior) ──
    # WireGuard marks its own encapsulated UDP packets with this fwmark.
    # Policy rules then ensure:
    #   - Normal packets (no mark) → VPN routing table → wg0
    #   - WireGuard-encapsulated packets (marked) → main table → eth0
    local FW_MARK=51820
    local FW_TABLE=51820
    wg set "$INTERFACE" fwmark "$FW_MARK"

    # Remove any auto-created default route on wg0
    ip route del default dev "$INTERFACE" 2>/dev/null || true

    # VPN routing table: send everything through the tunnel
    ip -4 route add default dev "$INTERFACE" table "$FW_TABLE"

    # Policy rules:
    # 1. Packets NOT marked by WireGuard use the VPN table (→ wg0)
    # 2. suppress_prefixlength 0: ignore bare default routes in main table,
    #    but keep more-specific routes (e.g. LAN, endpoint) working
    ip -4 rule add not fwmark "$FW_MARK" table "$FW_TABLE"
    ip -4 rule add table main suppress_prefixlength 0

    # Find default gateway/interface
    DEFAULT_GW=$(ip route | grep '^default' | head -1 | awk '{print $3}')
    DEFAULT_IF=$(ip route | grep '^default' | head -1 | awk '{print $5}')

    # ── Policy routing: ensure responses to incoming LAN traffic go back via eth0 ──
    if [ -n "$DEFAULT_GW" ] && [ -n "$DEFAULT_IF" ]; then
        # Get the container's eth0 IP address (BusyBox-compatible, no grep -P)
        ETH0_IP=$(ip -4 addr show "$DEFAULT_IF" | awk '/inet / {split($2, a, "/"); print a[1]}' | head -1)
        ETH0_SUBNET=$(ip -4 route show dev "$DEFAULT_IF" | grep -v default | head -1 | awk '{print $1}')
        if [ -n "$ETH0_IP" ] && [ -n "$ETH0_SUBNET" ]; then
            echo "[vpn] Setting up policy routing for incoming traffic (${ETH0_IP} on ${DEFAULT_IF})"
            ip route add default via "$DEFAULT_GW" dev "$DEFAULT_IF" table 100 2>/dev/null || true
            ip route add "$ETH0_SUBNET" dev "$DEFAULT_IF" table 100 2>/dev/null || true
            ip rule add from "$ETH0_IP" table 100 priority 100 2>/dev/null || true
            echo "[vpn] Policy routing active — incoming connections will be routed back via ${DEFAULT_IF}"
        fi
    fi

    # Set up DNS (handle comma-separated DNS servers)
    VPN_DNS=$(grep -i '^DNS' "$CONFIG_FILE" | head -1 | sed 's/.*= *//;s/ //g')
    if [ -n "$VPN_DNS" ]; then
        # Clear resolv.conf and add each DNS server on its own line
        > /etc/resolv.conf
        for dns in $(echo "$VPN_DNS" | tr ',' ' '); do
            echo "nameserver $dns" >> /etc/resolv.conf
            # Add explicit route to DNS server through wg0 so it's found in main table
            # (suppress_prefixlength 0 ignores default routes but allows host routes)
            ip -4 route add "$dns" dev "$INTERFACE" 2>/dev/null || true
        done
        echo "[vpn] DNS set to: ${VPN_DNS}"
    fi

    # Add explicit host route for the health-check target so it is picked up by
    # the 'lookup main suppress_prefixlength 0' rule (same as DNS servers above).
    # Without this, CHECK_HOST falls through to the VPN table default route whose
    # source-address selection can be defeated by the priority-100 'from ETH0_IP'
    # policy rule, causing pings to bypass wg0 and be dropped by the kill switch.
    ip -4 route add "${CHECK_HOST}" dev "$INTERFACE" 2>/dev/null || true
    echo "[vpn] Health-check route: ${CHECK_HOST} → ${INTERFACE}"

    echo "[vpn] WireGuard interface ${INTERFACE} is up."
    echo "[vpn] Main routes:"
    ip route show | sed 's/^/[vpn]   /'
    echo "[vpn] VPN table ($FW_TABLE):"
    ip route show table "$FW_TABLE" 2>/dev/null | sed 's/^/[vpn]   /'
    echo "[vpn] Policy rules:"
    ip rule show | sed 's/^/[vpn]   /'
    echo "[vpn] WireGuard status:"
    wg show "$INTERFACE" 2>/dev/null | sed 's/^/[vpn]   /'
}

# ──────────────────────────────────────────────
# Stop WireGuard manually
# ──────────────────────────────────────────────
stop_vpn() {
    echo "[vpn] Stopping WireGuard interface ${INTERFACE}..."

    local FW_MARK=51820
    local FW_TABLE=51820

    # Remove fwmark-based policy rules
    ip -4 rule del not fwmark "$FW_MARK" table "$FW_TABLE" 2>/dev/null || true
    ip -4 rule del table main suppress_prefixlength 0 2>/dev/null || true

    # Flush VPN routing table
    ip -4 route flush table "$FW_TABLE" 2>/dev/null || true

    # Remove LAN policy routing
    ip -4 rule del table 100 2>/dev/null || true
    ip -4 route flush table 100 2>/dev/null || true

    ip link del "$INTERFACE" 2>/dev/null || true
}

# ──────────────────────────────────────────────
# Health check loop — restarts VPN if tunnel dies
# ──────────────────────────────────────────────
health_loop() {
    local failures=0
    local max_failures=3

    echo "[health] Starting health check (every ${CHECK_INTERVAL}s, target ${CHECK_HOST})..."

    while true; do
        sleep "$CHECK_INTERVAL"

        if ping -c 1 -W 5 "$CHECK_HOST" > /dev/null 2>&1; then
            if [ "$failures" -gt 0 ]; then
                echo "[health] VPN recovered."
                failures=0
            fi
            # Secondary DNS check
            if ping -c 1 -W 5 "google.com" > /dev/null 2>&1; then
                : # DNS OK — silent
            else
                echo "[health] WARN google.com unreachable — possible DNS issue"
            fi
        else
            failures=$((failures + 1))
            echo "[health] Check failed ($failures/$max_failures) — ping ${CHECK_HOST} failed"
            # Secondary check: distinguish IP failure from DNS failure
            if ping -c 1 -W 5 "google.com" > /dev/null 2>&1; then
                echo "[health] INFO google.com reachable — DNS works, ${CHECK_HOST} may be filtered"
            else
                echo "[health] INFO google.com also unreachable — DNS or general routing failure"
            fi
            # Dump WireGuard stats to show if handshake is stale and how much data flows
            echo "[health] wg stats:"
            wg show "$INTERFACE" 2>/dev/null | grep -E 'latest handshake|transfer|endpoint' | sed 's/^/[health]   /' || echo "[health]   wg0 not found"
            echo "[health] routes:"
            ip route show | grep -E 'wg0|default' | sed 's/^/[health]   /'

            if [ "$failures" -ge "$max_failures" ]; then
                echo "[health] VPN appears down. Restarting WireGuard..."
                stop_vpn
                sleep 2
                start_vpn
                failures=0
                echo "[health] WireGuard restarted."
            fi
        fi
    done
}

# ──────────────────────────────────────────────
# Graceful shutdown
# ──────────────────────────────────────────────
cleanup() {
    echo "[shutdown] Stopping WireGuard..."
    stop_vpn
    echo "[shutdown] Flushing iptables..."
    iptables -F
    iptables -t nat -F
    echo "[shutdown] Done."
    exit 0
}

trap cleanup SIGTERM SIGINT

# ──────────────────────────────────────────────
# Startup connectivity checks — diagnose issues early
# ──────────────────────────────────────────────
check_vpn_connectivity() {
    echo "[check] ── Startup connectivity checks ──"

    # 1. Wait for WireGuard handshake (up to 15s)
    local elapsed=0
    local handshake_ts=0
    echo "[check] Waiting for WireGuard handshake (up to 15s)..."
    while [ "$elapsed" -lt 15 ]; do
        handshake_ts=$(wg show "$INTERFACE" latest-handshakes 2>/dev/null | awk '{print $2}' | head -1)
        if [ -n "$handshake_ts" ] && [ "$handshake_ts" != "0" ]; then
            local age=$(( $(date +%s) - handshake_ts ))
            echo "[check] OK   Handshake established ${age}s ago"
            break
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done
    if [ "$elapsed" -ge 15 ]; then
        echo "[check] FAIL No WireGuard handshake after 15s — tunnel is not up"
        echo "[check]      This container's public key (must be on the server):"
        echo "[check]        PublicKey  = $(wg show "$INTERFACE" public-key 2>/dev/null || echo 'unknown')"
        echo "[check]        AllowedIPs = ${VPN_ADDRESS}"
        echo "[check]      Verify on server: wg show"
    fi

    # 2. Check whether traffic actually flows through the tunnel
    echo "[check] Testing traffic through tunnel (ping ${CHECK_HOST})..."
    local rx_before
    rx_before=$(wg show "$INTERFACE" transfer 2>/dev/null | awk '{print $2}' | head -1)

    if ping -c 1 -W 8 "${CHECK_HOST}" > /dev/null 2>&1; then
        echo "[check] OK   Traffic flows — tunnel is fully working"
    else
        local rx_after
        rx_after=$(wg show "$INTERFACE" transfer 2>/dev/null | awk '{print $2}' | head -1)
        echo "[check] FAIL ping ${CHECK_HOST} unreachable through tunnel"

        if [ -n "$rx_before" ] && [ -n "$rx_after" ]; then
            if [ "$rx_after" -le "$rx_before" ]; then
                echo "[check]      RX bytes unchanged (${rx_before} → ${rx_after})"
                echo "[check]      Server receives packets but does NOT route them back"
                echo "[check]      Fix on VPN server (${VPN_ENDPOINT}):"
                echo "[check]        sysctl net.ipv4.ip_forward              # must output 1"
                echo "[check]        iptables -t nat -L POSTROUTING -v -n    # must have MASQUERADE"
                echo "[check]        wg show                                  # check peer + AllowedIPs"
            else
                echo "[check]      RX increased (${rx_before} → ${rx_after}) — tunnel passes data"
                echo "[check]      Issue may be specific to ${CHECK_HOST} or DNS"
            fi
        fi

        local transfer
        transfer=$(wg show "$INTERFACE" transfer 2>/dev/null | awk '{printf "rx=%s tx=%s", $2, $3}')
        echo "[check]      wg transfer: ${transfer}"
    fi

    # 3. DNS check
    echo "[check] Testing DNS resolution..."
    if nslookup 1.1.1.1 > /dev/null 2>&1 || nslookup google.com > /dev/null 2>&1; then
        echo "[check] OK   DNS resolves"
    else
        echo "[check] FAIL DNS resolution failed"
        echo "[check]      resolv.conf: $(tr '\n' ' ' < /etc/resolv.conf)"
        echo "[check]      Check that DNS servers are reachable through wg0"
        echo "[check] ── End of checks ──"
        exit 1
    fi

    echo "[check] ── End of checks ──"
}

# ── Main ──
enable_forwarding
setup_killswitch
start_vpn
check_vpn_connectivity
health_loop
