#!/bin/bash
# WireGuard setup for Isekai VPN on kiosks
# Usage: sudo bash setup_kiosk.sh --test   (for test environment)
#        sudo bash setup_kiosk.sh --prod   (for production)
#
# NOTE: This script ALWAYS exits 0 (never fails) because it's called from fgupdate.sh
# Errors are logged as warnings but don't stop execution

# Error handler - convert any uncaught errors to warnings
trap 'echo "⚠ WARNING: Script encountered error at line $LINENO - $BASH_COMMAND"; exit 0' ERR

# Parse arguments (default to prod)
ENV_MODE="prod"
for arg in "$@"; do
    case $arg in
        --test) ENV_MODE="test" ;;
        --prod) ENV_MODE="prod" ;;
        -h|--help)
            echo "Usage: sudo bash setup_kiosk.sh [--test|--prod]"
            echo "  --test  Connect to test environment (port 5001/51821)"
            echo "  --prod  Connect to production environment (port 5000/51820, default)"
            exit 0
            ;;
    esac
done

# Configuration based on environment
# Register keys are low-risk (can only join VPN, not list/delete peers)
if [ "$ENV_MODE" = "test" ]; then
    ISEKAI_API="https://50.151.230.237:5001"
    ISEKAI_SECRET="GmXyWdCSV5hJAFepaSw6YhGegTza1ALTYbDMkxT3Tnn6jLoJFMZ5yKHGZvtfg8uO"
    ENV_LABEL="TEST"
else
    ISEKAI_API="https://50.151.230.237:5000"
    ISEKAI_SECRET="gptsDBgdefva7REQ98mpXz03XiPYVIfLYnn0tHgYYZUlIGqwFx9vXFkKABgSNq7d"
    ENV_LABEL="PRODUCTION"
fi

echo "========================================="
echo "  Isekai WireGuard Setup ($ENV_LABEL)"
echo "========================================="
echo

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo "⚠ WARNING: This script must be run as root (use sudo)"
    echo "Skipping WireGuard setup (not root)"
    exit 0
fi

# Get machine ID (always ATM_MACHINE_ID from /etc/environment on production kiosks)
if [ -f /etc/environment ]; then
    MACHINE_ID=$(grep '^ATM_MACHINE_ID=' /etc/environment | cut -d'=' -f2 | tr -d '"' | tr -d "'")
fi
if [ -z "$MACHINE_ID" ] && [ -f /etc/machine-id-fg ]; then
    MACHINE_ID=$(cat /etc/machine-id-fg)
fi
if [ -z "$MACHINE_ID" ] && [ -f /etc/hostname ]; then
    MACHINE_ID=$(cat /etc/hostname)
    echo "Note: Using hostname (not a production FG kiosk)"
fi
if [ -z "$MACHINE_ID" ]; then
    echo "⚠ WARNING: No machine ID found in /etc/environment, /etc/machine-id-fg, or /etc/hostname"
    echo "Skipping WireGuard setup (no machine ID)"
    exit 0
fi
echo "Machine ID: $MACHINE_ID"

# Install WireGuard if not present (check both binary and package)
if ! command -v wg &> /dev/null; then
    # Check if pacman package exists
    if pacman -Q wireguard-tools &> /dev/null; then
        echo "✓ WireGuard package installed (via pacman)"
    else
        echo
        echo "Installing WireGuard (static binaries for old systems)..."
        SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

        # Check for local copy first (bundled with script)
        if [ -f "$SCRIPT_DIR/wireguard-static.tar.gz" ]; then
            echo "Using bundled WireGuard binaries..."
            tar xzf "$SCRIPT_DIR/wireguard-static.tar.gz" -C /tmp
        elif [ -f "./wireguard-static.tar.gz" ]; then
            echo "Using local WireGuard binaries..."
            tar xzf ./wireguard-static.tar.gz -C /tmp
        else
            # Download from fgtest
            echo "Downloading WireGuard binaries from fgtest..."
            cd /tmp
            if ! curl -fsSL https://fgtest.moneroatm.xyz/wireguard-static/wireguard-static.tar.gz -o wireguard-static.tar.gz 2>/dev/null; then
                echo "⚠ WARNING: Failed to download WireGuard static binaries from fgtest"
                echo "Skipping WireGuard setup (download failed)"
                exit 0
            fi
            if ! curl -fsSL https://fgtest.moneroatm.xyz/wireguard-static/wireguard-static.tar.gz.md5 -o wireguard-static.tar.gz.md5 2>/dev/null; then
                echo "⚠ WARNING: Failed to download checksum file"
                echo "Proceeding without checksum verification (risky)"
            else
                # Verify checksum
                if md5sum -c wireguard-static.tar.gz.md5 2>/dev/null | grep -q "OK"; then
                    echo "✓ Checksum verified"
                else
                    echo "⚠ WARNING: Checksum verification failed - download may be corrupted"
                    echo "Skipping WireGuard setup (corrupted download)"
                    rm -f wireguard-static.tar.gz wireguard-static.tar.gz.md5
                    exit 0
                fi
            fi
            tar xzf wireguard-static.tar.gz
        fi

        # Install from /tmp
        install -m 755 /tmp/wireguard-static/wg /usr/bin/wg
        install -m 755 /tmp/wireguard-static/wg-quick /usr/bin/wg-quick
        rm -rf /tmp/wireguard-static /tmp/wireguard-static.tar.gz
        echo "✓ WireGuard installed to /usr/bin"
    fi
else
    echo "✓ WireGuard already installed"
fi

# Check if tunnel is healthy before reconfiguring
if [ -f /etc/wireguard/wg-isekai.conf ]; then
    echo
    echo "Checking existing WireGuard tunnel health..."

    # Check if interface exists and has recent handshake
    if wg show wg-isekai 2>/dev/null | grep -q "latest handshake:"; then
        HANDSHAKE_AGE=$(wg show wg-isekai latest-handshakes 2>/dev/null | head -1 | awk '{print $2}')
        CURRENT_TIME=$(date +%s)

        if [ -n "$HANDSHAKE_AGE" ] && [ "$((CURRENT_TIME - HANDSHAKE_AGE))" -lt 180 ]; then
            echo "✓ Tunnel is healthy (recent handshake), skipping reconfiguration"
            exit 0
        else
            echo "⚠ Tunnel unhealthy (stale/no handshake), will reuse keys and re-register..."
        fi
    else
        echo "⚠ Tunnel down, will reuse keys and re-register..."
    fi

    # Stop interface but preserve config (keys will be reused)
    systemctl stop wg-isekai.service 2>/dev/null || true
    ip link delete wg-isekai 2>/dev/null || true
fi

# Get or generate WireGuard keys
echo
if [ -f /etc/wireguard/wg-isekai.conf ]; then
    echo "Reusing existing WireGuard keys..."
    PRIVATE_KEY=$(grep 'PrivateKey' /etc/wireguard/wg-isekai.conf | sed 's/.*PrivateKey *= *//' | tr -d ' \t\r\n')
    if [ -n "$PRIVATE_KEY" ]; then
        PUBLIC_KEY=$(echo "$PRIVATE_KEY" | wg pubkey)
        echo "✓ Keys loaded from existing config"
    else
        echo "⚠ Could not extract key from config, generating new..."
        PRIVATE_KEY=$(wg genkey)
        PUBLIC_KEY=$(echo "$PRIVATE_KEY" | wg pubkey)
        echo "✓ Keys generated"
    fi
else
    echo "Generating new WireGuard keys..."
    PRIVATE_KEY=$(wg genkey)
    PUBLIC_KEY=$(echo "$PRIVATE_KEY" | wg pubkey)
    echo "✓ Keys generated"
fi
echo "  Public key: $PUBLIC_KEY"

# Register with Isekai API
echo
echo "Registering with Isekai control plane..."
RESPONSE=$(curl -s -k -X POST "$ISEKAI_API/api/v1/peers/register" \
    -H "Content-Type: application/json" \
    -d "{\"isekai_secret\": \"$ISEKAI_SECRET\", \"machine_id\": \"$MACHINE_ID\", \"public_key\": \"$PUBLIC_KEY\"}")

echo "API Response:"
echo "$RESPONSE" | python3 -m json.tool || echo "$RESPONSE"

# Extract config from API response
ALLOCATED_IP=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('allocated_wg_ip', ''))" 2>/dev/null || echo "")
API_HUB_ENDPOINT=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('hub_endpoint', ''))" 2>/dev/null || echo "")
API_HUB_PUBKEY=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('hub_public_key', ''))" 2>/dev/null || echo "")
API_ALLOWED_IPS=$(echo "$RESPONSE" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('allowed_ips', ''))" 2>/dev/null || echo "")

if [ -z "$ALLOCATED_IP" ] || [ -z "$API_HUB_ENDPOINT" ]; then
    echo
    echo "⚠ WARNING: Failed to get config from Isekai API!"
    echo "Response: $RESPONSE"
    echo "Skipping WireGuard setup (API registration failed)"
    exit 0
fi

echo
echo "✓ Registered successfully"
echo "  Allocated VPN IP: $ALLOCATED_IP"
echo "  Hub Endpoint: $API_HUB_ENDPOINT"

# Exclude wg-isekai from systemd-networkd (let wg-quick manage it)
echo
echo "Configuring systemd-networkd to ignore wg-isekai..."
cat > /etc/systemd/network/10-wg-isekai.network <<EOF
[Match]
Name=wg-isekai

[Link]
Unmanaged=yes
EOF
systemctl restart systemd-networkd
echo "✓ systemd-networkd configured"

# Write WireGuard config
echo
echo "Writing /etc/wireguard/wg-isekai.conf..."
mkdir -p /etc/wireguard
cat > /etc/wireguard/wg-isekai.conf <<EOF
[Interface]
PrivateKey = $PRIVATE_KEY
Address = $ALLOCATED_IP/32

[Peer]
# Isekai Hub Server
PublicKey = $API_HUB_PUBKEY
Endpoint = $API_HUB_ENDPOINT
AllowedIPs = $API_ALLOWED_IPS
PersistentKeepalive = 25
EOF

chmod 600 /etc/wireguard/wg-isekai.conf
echo "✓ Config written"

# Start WireGuard
echo
echo "Starting WireGuard tunnel..."
wg-quick up wg-isekai

# Allow incoming traffic from WireGuard subnet
echo
echo "Configuring firewall to allow traffic from VPN..."
if command -v ufw &> /dev/null; then
    ufw allow in on wg-isekai
    echo "✓ UFW rule added"
elif command -v iptables &> /dev/null; then
    if [ "$ENV_MODE" = "test" ]; then
        iptables -I INPUT 1 -i wg-isekai -s 10.201.0.0/16 -j ACCEPT
    else
        iptables -I INPUT 1 -i wg-isekai -s 10.200.0.0/16 -j ACCEPT
    fi
    echo "✓ iptables rule added"
fi

# Enable on boot via systemd service
echo
echo "Enabling WireGuard service to start on boot..."
cat > /etc/systemd/system/wg-isekai.service <<'SVCEOF'
[Unit]
Description=WireGuard VPN Tunnel - Isekai
After=network-online.target nss-lookup.target
Wants=network-online.target nss-lookup.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/bin/wg-quick up wg-isekai
ExecStop=/usr/bin/wg-quick down wg-isekai

[Install]
WantedBy=multi-user.target
SVCEOF
systemctl daemon-reload
systemctl enable wg-isekai.service

# Show status
echo
echo "========================================="
echo "  ✓ Setup Complete!"
echo "========================================="
echo
echo "Machine ID:    $MACHINE_ID"
echo "VPN IP:        $ALLOCATED_IP"
echo "Hub Endpoint:  $API_HUB_ENDPOINT"
echo
echo "WireGuard Status:"
wg show wg-isekai
echo
echo "You can now SSH to this kiosk from the office using:"
echo "  ssh freedomgateway@$ALLOCATED_IP"
echo "  OR use: isekai ssh $MACHINE_ID"
echo
echo "To check status: wg show wg-isekai"
echo "To restart:      sudo systemctl restart wg-quick@wg-isekai"
echo "To stop:         sudo systemctl stop wg-quick@wg-isekai"
echo
