#!/bin/sh
# Parahub Mesh — Local Mullvad WireGuard Manager
#
# Allows the node owner to run Mullvad directly on the router.
# Guest traffic routes through the local WireGuard tunnel instead of
# the default VPS gateway path, giving lower latency and using the
# nearest Mullvad server.
#
# Usage:
#   parahub-mullvad setup <account_key> [country_code]
#   parahub-mullvad status
#   parahub-mullvad remove

set -e

PARAHUB_DIR="/etc/parahub"
ACCOUNT_FILE="$PARAHUB_DIR/mullvad_account"

# ============================================================================
# Helpers
# ============================================================================

# Find the firewall forwarding index where src=guest
find_guest_forwarding() {
    local idx=0
    while uci -q get "firewall.@forwarding[$idx]" >/dev/null 2>&1; do
        local src=$(uci -q get "firewall.@forwarding[$idx].src")
        if [ "$src" = "guest" ]; then
            echo "$idx"
            return 0
        fi
        idx=$((idx + 1))
    done
    return 1
}

# Find firewall zone index by name
find_zone_index() {
    local name="$1"
    local idx=0
    while uci -q get "firewall.@zone[$idx]" >/dev/null 2>&1; do
        local zname=$(uci -q get "firewall.@zone[$idx].name")
        if [ "$zname" = "$name" ]; then
            echo "$idx"
            return 0
        fi
        idx=$((idx + 1))
    done
    return 1
}

# ============================================================================
# setup <account_key> [country_code]
# ============================================================================

cmd_setup() {
    local ACCOUNT="$1"
    local COUNTRY="$2"

    if [ -z "$ACCOUNT" ]; then
        echo "Usage: parahub-mullvad setup <account_key> [country_code]"
        echo ""
        echo "Examples:"
        echo "  parahub-mullvad setup 1234567890123456 pt   # Portugal"
        echo "  parahub-mullvad setup 1234567890123456 de   # Germany"
        echo "  parahub-mullvad setup 1234567890123456       # Auto-detect"
        echo ""
        echo "Countries: us gb de nl se pt fr es ch at it jp sg br ..."
        exit 1
    fi

    # Auto-detect country from WAN IP
    if [ -z "$COUNTRY" ]; then
        echo "Auto-detecting country from WAN IP..."
        COUNTRY=$(curl -s --max-time 5 https://ipinfo.io/country 2>/dev/null | tr 'A-Z' 'a-z')
        if [ -z "$COUNTRY" ] || [ ${#COUNTRY} -ne 2 ]; then
            echo "Error: Could not detect country."
            echo "Specify manually: parahub-mullvad setup $ACCOUNT <country>"
            exit 1
        fi
        echo "Detected: $COUNTRY"
    fi

    COUNTRY=$(echo "$COUNTRY" | tr 'A-Z' 'a-z')

    # --- Step 1: Generate WireGuard keys ---
    echo "Generating WireGuard keys..."
    umask 077
    wg genkey > "$PARAHUB_DIR/wg_private.key"
    wg pubkey < "$PARAHUB_DIR/wg_private.key" > "$PARAHUB_DIR/wg_public.key"

    PRIVKEY=$(cat "$PARAHUB_DIR/wg_private.key")
    PUBKEY=$(cat "$PARAHUB_DIR/wg_public.key")

    # --- Step 2: Register with Mullvad API ---
    echo "Registering key with Mullvad..."
    RESULT=$(curl -s --max-time 15 -X POST https://api.mullvad.net/wg/ \
        --data-urlencode "account=$ACCOUNT" \
        --data-urlencode "pubkey=$PUBKEY")

    if echo "$RESULT" | grep -q "^[0-9]"; then
        MULLVAD_IPV4=$(echo "$RESULT" | cut -d',' -f1)
        echo "Mullvad IP: $MULLVAD_IPV4"
    else
        echo "Error from Mullvad: $RESULT"
        exit 1
    fi

    # --- Step 3: Find server in target country ---
    echo "Finding Mullvad server in '$COUNTRY'..."
    curl -s --max-time 15 https://api.mullvad.net/www/relays/wireguard/ | \
        tr '{' '\n' | \
        grep "\"country_code\":\"${COUNTRY}\"" | \
        grep '"active":true' | \
        head -1 > /tmp/mullvad_relay.tmp

    if [ ! -s /tmp/mullvad_relay.tmp ]; then
        echo "Error: No active WireGuard server for country '$COUNTRY'"
        rm -f /tmp/mullvad_relay.tmp
        exit 1
    fi

    SERVER_IP=$(sed 's/.*"ipv4_addr_in":"\([^"]*\)".*/\1/' /tmp/mullvad_relay.tmp)
    SERVER_PUBKEY=$(sed 's/.*"pubkey":"\([^"]*\)".*/\1/' /tmp/mullvad_relay.tmp)
    SERVER_HOST=$(sed 's/.*"hostname":"\([^"]*\)".*/\1/' /tmp/mullvad_relay.tmp)
    rm -f /tmp/mullvad_relay.tmp

    echo "Server: $SERVER_HOST ($SERVER_IP)"

    # --- Step 4: Clean previous config ---
    uci -q delete network.mullvad_local 2>/dev/null || true
    while uci -q delete network.@wireguard_mullvad_local[0] 2>/dev/null; do :; done

    # --- Step 5: Create WireGuard interface (routes in table 100) ---
    echo "Configuring WireGuard..."
    uci batch <<-WG_EOF
set network.mullvad_local=interface
set network.mullvad_local.proto='wireguard'
set network.mullvad_local.private_key='${PRIVKEY}'
add_list network.mullvad_local.addresses='${MULLVAD_IPV4}'
set network.mullvad_local.mtu='1420'
set network.mullvad_local.ip4table='100'

add network wireguard_mullvad_local
set network.@wireguard_mullvad_local[-1].public_key='${SERVER_PUBKEY}'
set network.@wireguard_mullvad_local[-1].endpoint_host='${SERVER_IP}'
set network.@wireguard_mullvad_local[-1].endpoint_port='51820'
add_list network.@wireguard_mullvad_local[-1].allowed_ips='0.0.0.0/0'
set network.@wireguard_mullvad_local[-1].route_allowed_ips='1'
set network.@wireguard_mullvad_local[-1].persistent_keepalive='25'
WG_EOF

    # Disable VPS gateway (Mullvad replaces it in table 100)
    uci -q set network.vps_gateway.auto='0' 2>/dev/null || true
    uci commit network

    # --- Step 6: Firewall zone for mullvad_local ---
    local zone_idx
    if zone_idx=$(find_zone_index "mullvad_local"); then
        uci delete "firewall.@zone[$zone_idx]"
    fi

    uci batch <<-FW_EOF
add firewall zone
set firewall.@zone[-1].name='mullvad_local'
set firewall.@zone[-1].input='REJECT'
set firewall.@zone[-1].output='ACCEPT'
set firewall.@zone[-1].forward='REJECT'
set firewall.@zone[-1].masq='1'
set firewall.@zone[-1].mtu_fix='1'
add_list firewall.@zone[-1].network='mullvad_local'
FW_EOF

    # Switch guest forwarding to mullvad_local
    local fwd_idx
    if fwd_idx=$(find_guest_forwarding); then
        uci set "firewall.@forwarding[$fwd_idx].dest=mullvad_local"
    fi
    uci commit firewall

    # --- Step 7: Save config ---
    cat > "$ACCOUNT_FILE" <<-ACCT_EOF
MULLVAD_ACCOUNT=${ACCOUNT}
MULLVAD_COUNTRY=${COUNTRY}
MULLVAD_SERVER=${SERVER_HOST}
MULLVAD_SERVER_IP=${SERVER_IP}
MULLVAD_LOCAL_IP=${MULLVAD_IPV4}
ACCT_EOF
    chmod 600 "$ACCOUNT_FILE"

    # --- Step 8: Apply ---
    echo "Restarting network..."
    /etc/init.d/network restart
    /etc/init.d/firewall restart

    echo ""
    echo "Done! Guest traffic now routes directly through Mullvad."
    echo "Server: $SERVER_HOST ($COUNTRY)"
    echo "Test: connect to parahub.io/free, visit https://am.i.mullvad.net"
}

# ============================================================================
# status
# ============================================================================

cmd_status() {
    echo "=== Parahub Mesh VPN Status ==="
    echo ""

    if [ -f "$ACCOUNT_FILE" ]; then
        echo "Mode: LOCAL MULLVAD (direct)"
        cat "$ACCOUNT_FILE"
        echo ""
        wg show mullvad_local 2>/dev/null || echo "WireGuard interface: not up"
    else
        echo "Mode: VPS GATEWAY (WireGuard)"
        echo "VPS: 185.47.131.84 -> Mullvad"
        echo ""
        wg show vps_gateway 2>/dev/null || echo "VPS gateway: not active"
    fi

    echo ""
    echo "Guest forwarding:"
    local fwd_idx
    if fwd_idx=$(find_guest_forwarding); then
        echo "  guest -> $(uci -q get "firewall.@forwarding[$fwd_idx].dest")"
    else
        echo "  (none)"
    fi
}

# ============================================================================
# remove
# ============================================================================

cmd_remove() {
    if [ ! -f "$ACCOUNT_FILE" ]; then
        echo "No local Mullvad configured. Nothing to remove."
        exit 0
    fi

    echo "Removing local Mullvad, reverting to VPS gateway..."

    # Remove WireGuard interface
    uci -q delete network.mullvad_local 2>/dev/null || true
    while uci -q delete network.@wireguard_mullvad_local[0] 2>/dev/null; do :; done

    # Re-enable VPS gateway
    uci -q delete network.vps_gateway.auto 2>/dev/null || true
    uci commit network

    # Remove firewall zone
    local zone_idx
    if zone_idx=$(find_zone_index "mullvad_local"); then
        uci delete "firewall.@zone[$zone_idx]"
    fi

    # Switch guest forwarding back to vps_gateway
    local fwd_idx
    if fwd_idx=$(find_guest_forwarding); then
        uci set "firewall.@forwarding[$fwd_idx].dest=vps_gateway"
    fi
    uci commit firewall

    # Clean up files
    rm -f "$PARAHUB_DIR/wg_private.key" "$PARAHUB_DIR/wg_public.key" "$ACCOUNT_FILE"

    echo "Restarting network..."
    /etc/init.d/network restart
    /etc/init.d/firewall restart

    echo "Done! Guest traffic now routes through VPS gateway."
}

# ============================================================================
# Main
# ============================================================================

case "${1:-}" in
    setup)
        cmd_setup "$2" "${3:-}"
        ;;
    status)
        cmd_status
        ;;
    remove)
        cmd_remove
        ;;
    *)
        echo "Parahub Mesh — Mullvad WireGuard Manager"
        echo ""
        echo "Commands:"
        echo "  parahub-mullvad setup <account_key> [country]"
        echo "  parahub-mullvad status"
        echo "  parahub-mullvad remove"
        echo ""
        echo "By default, guest traffic goes through the VPS gateway."
        echo "With 'setup', it routes directly through your own Mullvad"
        echo "account — faster, and uses the nearest server."
        ;;
esac