#!/bin/bash
#
# wg1-setup.sh - Set up WireGuard wg1 tunnel between earth and hyperstack VM
#
# USAGE:
#   ./wg1-setup.sh <VM_PUBLIC_IP>
#   Example: ./wg1-setup.sh 185.216.20.163
#
# NETWORK DESIGN:
#   Subnet: 192.168.3.0/24 (separate from wg0's 192.168.2.0/24)
#   Port: 56710/udp
#
#   +----------------+                     +------------------+
#   | earth (client) |                     | hyperstack (VM)  |
#   | 192.168.3.2    |<--- WireGuard --->  | 192.168.3.1      |
#   +----------------+     tunnel          +------------------+
#                                          | Ollama :11434    |
#                                          +------------------+
#
# WHAT THIS SCRIPT DOES:
#   On hyperstack VM (via SSH):
#     - Installs WireGuard if not present
#     - Creates /etc/wireguard/wg1.conf
#     - Opens UFW ports: 56710/udp (WireGuard), 11434/tcp from 192.168.3.0/24 (Ollama)
#     - Configures Ollama to listen on 0.0.0.0:11434
#     - Starts wg-quick@wg1
#
#   On earth (locally):
#     - Installs WireGuard if not present (dnf)
#     - Creates /etc/wireguard/wg1.conf
#     - Starts wg-quick@wg1
#
# PREREQUISITES:
#   - SSH access to ubuntu@<VM_IP> with key-based auth
#   - UDP port 56710 open in cloud provider's firewall/security group
#
# RE-RUNNING:
#   When the VM IP changes, simply re-run this script with the new IP.
#   It will regenerate keys and update configs on both sides.
#
# USING OLLAMA REMOTELY:
#   export OLLAMA_HOST=http://192.168.3.1:11434
#   ollama run qwen2.5-coder:14b-instruct
#   # Or with aider:
#   aider --model ollama/qwen2.5-coder:14b-instruct
#

set -euo pipefail

# Configuration constants
WG_INTERFACE="wg1"
WG_PORT="56710"
SERVER_WG_IP="192.168.3.1"
CLIENT_WG_IP="192.168.3.2"
SUBNET_MASK="24"
SSH_USER="ubuntu"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

print_warning() {
    echo -e "${YELLOW}$1${NC}"
}

print_success() {
    echo -e "${GREEN}$1${NC}"
}

print_error() {
    echo -e "${RED}$1${NC}"
}

# Validate arguments
if [[ $# -ne 1 ]]; then
    echo "Usage: $0 <VM_PUBLIC_IP>"
    echo "Example: $0 185.216.20.163"
    exit 1
fi

VM_IP="$1"

# Validate IP format (basic check)
if ! [[ "$VM_IP" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
    print_error "Error: Invalid IP address format: $VM_IP"
    exit 1
fi

echo "=============================================="
print_warning "IMPORTANT: Ensure UDP port ${WG_PORT} is open on the VM!"
print_warning "This must be configured in your cloud provider's"
print_warning "firewall/security group settings."
echo "=============================================="
echo ""
read -p "Press Enter to continue (or Ctrl+C to abort)..."
echo ""

# Create temporary directory for key generation
TMPDIR=$(mktemp -d)
trap "rm -rf $TMPDIR" EXIT

echo "=== Generating WireGuard keys locally ==="

# Generate server (hyperstack) keys
wg genkey > "$TMPDIR/server-privatekey"
wg pubkey < "$TMPDIR/server-privatekey" > "$TMPDIR/server-publickey"
SERVER_PRIVATE_KEY=$(cat "$TMPDIR/server-privatekey")
SERVER_PUBLIC_KEY=$(cat "$TMPDIR/server-publickey")

# Generate client (earth) keys
wg genkey > "$TMPDIR/client-privatekey"
wg pubkey < "$TMPDIR/client-privatekey" > "$TMPDIR/client-publickey"
CLIENT_PRIVATE_KEY=$(cat "$TMPDIR/client-privatekey")
CLIENT_PUBLIC_KEY=$(cat "$TMPDIR/client-publickey")

print_success "Keys generated successfully"

echo ""
echo "=== Creating server (hyperstack) configuration ==="

# Create server wg1.conf
cat > "$TMPDIR/server-wg1.conf" << EOF
# WireGuard wg1 configuration for hyperstack VM
# Server side of earth <-> hyperstack tunnel
# Generated by wg1-setup.sh on $(date)

[Interface]
Address = ${SERVER_WG_IP}/${SUBNET_MASK}
ListenPort = ${WG_PORT}
PrivateKey = ${SERVER_PRIVATE_KEY}

[Peer]
# earth (client)
PublicKey = ${CLIENT_PUBLIC_KEY}
AllowedIPs = ${CLIENT_WG_IP}/32
EOF

print_success "Server config created"

echo ""
echo "=== Creating client (earth) configuration ==="

# Create client wg1.conf
cat > "$TMPDIR/client-wg1.conf" << EOF
# WireGuard wg1 configuration for earth
# Client side of earth <-> hyperstack tunnel
# Generated by wg1-setup.sh on $(date)

[Interface]
Address = ${CLIENT_WG_IP}/${SUBNET_MASK}
PrivateKey = ${CLIENT_PRIVATE_KEY}

[Peer]
# hyperstack VM (server)
PublicKey = ${SERVER_PUBLIC_KEY}
Endpoint = ${VM_IP}:${WG_PORT}
AllowedIPs = ${SERVER_WG_IP}/32
PersistentKeepalive = 25
EOF

print_success "Client config created"

echo ""
echo "=== Setting up hyperstack VM (${VM_IP}) ==="

# Check SSH connectivity
echo "Testing SSH connection..."
if ! ssh -o ConnectTimeout=10 -o BatchMode=yes "${SSH_USER}@${VM_IP}" "echo 'SSH OK'" 2>/dev/null; then
    print_error "Error: Cannot connect to ${SSH_USER}@${VM_IP}"
    print_error "Please ensure SSH access is configured."
    exit 1
fi
print_success "SSH connection OK"

# Install WireGuard on server if not present
echo "Installing WireGuard on hyperstack..."
ssh "${SSH_USER}@${VM_IP}" "which wg >/dev/null 2>&1 || (sudo apt update && sudo apt install -y wireguard)"
print_success "WireGuard installed"

# Copy server config to hyperstack
echo "Copying wg1.conf to hyperstack..."
scp "$TMPDIR/server-wg1.conf" "${SSH_USER}@${VM_IP}:/tmp/wg1.conf"
ssh "${SSH_USER}@${VM_IP}" "sudo mv /tmp/wg1.conf /etc/wireguard/wg1.conf && sudo chmod 600 /etc/wireguard/wg1.conf"
print_success "Server config installed"

# Configure firewall on hyperstack
echo "Configuring firewall (ufw) on hyperstack..."
ssh "${SSH_USER}@${VM_IP}" << 'REMOTE_SCRIPT'
# Ensure ufw is enabled
sudo ufw --force enable >/dev/null 2>&1 || true

# Allow WireGuard port
sudo ufw allow 56710/udp comment 'WireGuard wg1' 2>/dev/null || true

# Allow Ollama access from wg1 subnet
sudo ufw allow from 192.168.3.0/24 to any port 11434 proto tcp comment 'Ollama via wg1' 2>/dev/null || true

echo "Firewall rules added"
REMOTE_SCRIPT
print_success "Firewall configured"

# Configure Ollama to listen on all interfaces
echo "Configuring Ollama to listen on 0.0.0.0..."
ssh "${SSH_USER}@${VM_IP}" << 'REMOTE_SCRIPT'
# Create override directory if it doesn't exist
sudo mkdir -p /etc/systemd/system/ollama.service.d

# Create or update override.conf to bind Ollama to all interfaces
cat << 'OVERRIDE' | sudo tee /etc/systemd/system/ollama.service.d/override.conf > /dev/null
[Service]
Environment="OLLAMA_HOST=0.0.0.0:11434"
OVERRIDE

# Reload systemd and restart Ollama
sudo systemctl daemon-reload
sudo systemctl restart ollama 2>/dev/null || echo "Note: Ollama service not running or not installed"
REMOTE_SCRIPT
print_success "Ollama configured"

# Start wg1 on hyperstack
echo "Starting wg1 on hyperstack..."
ssh "${SSH_USER}@${VM_IP}" "sudo systemctl start wg-quick@wg1 2>/dev/null || sudo wg-quick up wg1"
print_success "wg1 started on hyperstack"

echo ""
echo "=== Setting up earth (local) ==="

# Check if WireGuard is installed locally
if ! which wg >/dev/null 2>&1; then
    echo "Installing WireGuard locally..."
    sudo dnf install -y wireguard-tools
fi
print_success "WireGuard installed locally"

# Install client config locally
echo "Installing wg1.conf locally..."
sudo cp "$TMPDIR/client-wg1.conf" /etc/wireguard/wg1.conf
sudo chmod 600 /etc/wireguard/wg1.conf
print_success "Client config installed"

# Stop existing wg1 if running, then start fresh
echo "Starting wg1 locally..."
sudo systemctl stop wg-quick@wg1 2>/dev/null || true
sudo systemctl start wg-quick@wg1
print_success "wg1 started locally"

echo ""
echo "=============================================="
print_success "Setup complete!"
echo "=============================================="
echo ""
echo "WireGuard wg1 tunnel is now active."
echo ""
echo "Tunnel IPs:"
echo "  hyperstack (server): ${SERVER_WG_IP}"
echo "  earth (client):      ${CLIENT_WG_IP}"
echo ""
echo "=== Verification commands ==="
echo ""
echo "# Check tunnel status:"
echo "sudo wg show wg1"
echo ""
echo "# Ping hyperstack via tunnel:"
echo "ping -c 3 ${SERVER_WG_IP}"
echo ""
echo "# Verify default route is UNCHANGED:"
echo "ip route | grep default"
echo ""
echo "# Test Ollama access:"
echo "curl http://${SERVER_WG_IP}:11434/api/tags"
echo ""
echo "=== Manual start/stop commands ==="
echo ""
echo "# Stop tunnel:"
echo "sudo systemctl stop wg-quick@wg1"
echo ""
echo "# Start tunnel:"
echo "sudo systemctl start wg-quick@wg1"
echo ""
echo "# Restart on hyperstack (if VM rebooted):"
echo "ssh ${SSH_USER}@${VM_IP} 'sudo systemctl start wg-quick@wg1'"
echo ""
echo "=== Use Ollama remotely ==="
echo ""
echo "export OLLAMA_HOST=http://${SERVER_WG_IP}:11434"
echo "curl http://${SERVER_WG_IP}:11434/v1/models"