#!/bin/bash
source ~/.bashrc
# Setup NAT for VMs - Make VMs accessible via Proxmox host
# Creates a NAT network so VMs can be accessed via Proxmox host IP

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

# Load environment variables
if [ -f "$PROJECT_ROOT/.env" ]; then
    set -a
    source <(grep -v '^#' "$PROJECT_ROOT/.env" | grep -v '^$' | sed 's/#.*$//' | grep '=')
    set +a
fi

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo ""
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}========================================${NC}"
    echo ""
}

PROXMOX_HOST="${PROXMOX_ML110_IP:-192.168.1.206}"
SSH_KEY="${SSH_KEY:-$HOME/.ssh/id_ed25519_proxmox}"

# NAT network configuration
NAT_NETWORK="10.0.0.0/24"
NAT_BRIDGE="vmbr1"
NAT_GATEWAY="10.0.0.1"

# VM definitions: vmid name nat_ip
VMS=(
    "100 cloudflare-tunnel 10.0.0.10"
    "101 k3s-master 10.0.0.11"
    "102 git-server 10.0.0.12"
    "103 observability 10.0.0.13"
)

setup_nat_bridge() {
    log_step "Step 1: Setting up NAT Bridge"
    
    log_info "Creating NAT bridge $NAT_BRIDGE on Proxmox host..."
    
    ssh -i "$SSH_KEY" root@$PROXMOX_HOST <<EOF
set -e

# Check if bridge already exists
if ip link show $NAT_BRIDGE &>/dev/null; then
    echo "Bridge $NAT_BRIDGE already exists"
else
    # Create bridge
    cat >> /etc/network/interfaces <<INTERFACES

# NAT bridge for VMs
auto $NAT_BRIDGE
iface $NAT_BRIDGE inet static
    address $NAT_GATEWAY
    netmask 255.255.255.0
    bridge_ports none
    bridge_stp off
    bridge_fd 0
    post-up echo 1 > /proc/sys/net/ipv4/ip_forward
    post-up iptables -t nat -A POSTROUTING -s $NAT_NETWORK -o vmbr0 -j MASQUERADE
    post-up iptables -A FORWARD -s $NAT_NETWORK -j ACCEPT
    post-up iptables -A FORWARD -d $NAT_NETWORK -j ACCEPT
INTERFACES

    # Bring up bridge
    ifup $NAT_BRIDGE
    
    echo "✓ NAT bridge $NAT_BRIDGE created"
fi

# Enable IP forwarding
echo 1 > /proc/sys/net/ipv4/ip_forward

# Setup iptables rules (idempotent)
iptables -t nat -C POSTROUTING -s $NAT_NETWORK -o vmbr0 -j MASQUERADE 2>/dev/null || \
    iptables -t nat -A POSTROUTING -s $NAT_NETWORK -o vmbr0 -j MASQUERADE

iptables -C FORWARD -s $NAT_NETWORK -j ACCEPT 2>/dev/null || \
    iptables -A FORWARD -s $NAT_NETWORK -j ACCEPT

iptables -C FORWARD -d $NAT_NETWORK -j ACCEPT 2>/dev/null || \
    iptables -A FORWARD -d $NAT_NETWORK -j ACCEPT

echo "✓ NAT rules configured"
EOF
    
    log_info "✓ NAT bridge configured"
}

configure_vm_nat() {
    local vmid=$1
    local name=$2
    local nat_ip=$3
    
    log_info "Configuring VM $vmid ($name) with NAT IP $nat_ip..."
    
    ssh -i "$SSH_KEY" root@$PROXMOX_HOST <<EOF
# Update VM network to use NAT bridge
qm set $vmid --net0 virtio,bridge=$NAT_BRIDGE

# Configure cloud-init with NAT IP
qm set $vmid --ipconfig0 ip=$nat_ip/24,gw=$NAT_GATEWAY

echo "✓ VM $vmid configured for NAT"
EOF
}

setup_port_forwarding() {
    log_step "Step 3: Setting up Port Forwarding"
    
    log_info "Setting up port forwarding rules..."
    
    # Port mappings: external_port -> vm_nat_ip:internal_port
    # Format: vmid external_port internal_port description
    PORT_MAPPINGS=(
        "100 2222 22 cloudflare-tunnel-ssh"
        "101 2223 22 k3s-master-ssh"
        "102 2224 22 git-server-ssh"
        "103 2225 22 observability-ssh"
        "102 3000 3000 gitea-web"
        "103 9090 9090 prometheus"
        "103 3001 3000 grafana"
    )
    
    ssh -i "$SSH_KEY" root@$PROXMOX_HOST <<'EOF'
set -e

# Get NAT IPs for VMs
declare -A VM_NAT_IPS=(
    ["100"]="10.0.0.10"
    ["101"]="10.0.0.11"
    ["102"]="10.0.0.12"
    ["103"]="10.0.0.13"
)

# Port forwarding rules
# Format: vmid external_port internal_port
PORT_MAPPINGS=(
    "100 2222 22"
    "101 2223 22"
    "102 2224 22"
    "103 2225 22"
    "102 3000 3000"
    "103 9090 9090"
    "103 3001 3000"
)

for mapping in "${PORT_MAPPINGS[@]}"; do
    read -r vmid ext_port int_port <<< "$mapping"
    nat_ip="${VM_NAT_IPS[$vmid]}"
    
    # Check if rule exists
    if iptables -t nat -C PREROUTING -p tcp --dport $ext_port -j DNAT --to-destination $nat_ip:$int_port 2>/dev/null; then
        echo "Port forwarding $ext_port -> $nat_ip:$int_port already exists"
    else
        # Add port forwarding
        iptables -t nat -A PREROUTING -p tcp --dport $ext_port -j DNAT --to-destination $nat_ip:$int_port
        iptables -A FORWARD -p tcp -d $nat_ip --dport $int_port -j ACCEPT
        echo "✓ Port forwarding: $PROXMOX_HOST:$ext_port -> $nat_ip:$int_port"
    fi
done

# Save iptables rules
if command -v netfilter-persistent &>/dev/null; then
    netfilter-persistent save
elif [ -f /etc/iptables/rules.v4 ]; then
    iptables-save > /etc/iptables/rules.v4
fi

echo "✓ Port forwarding configured"
EOF
    
    log_info "✓ Port forwarding configured"
}

show_access_info() {
    log_step "Access Information"
    
    log_info "VM Access via NAT:"
    echo ""
    echo "  VM 100 (cloudflare-tunnel):"
    echo "    SSH: ssh -i $SSH_KEY ubuntu@$PROXMOX_HOST -p 2222"
    echo "    Direct NAT: ssh -i $SSH_KEY ubuntu@10.0.0.10 (from Proxmox host)"
    echo ""
    echo "  VM 101 (k3s-master):"
    echo "    SSH: ssh -i $SSH_KEY ubuntu@$PROXMOX_HOST -p 2223"
    echo "    Direct NAT: ssh -i $SSH_KEY ubuntu@10.0.0.11 (from Proxmox host)"
    echo ""
    echo "  VM 102 (git-server):"
    echo "    SSH: ssh -i $SSH_KEY ubuntu@$PROXMOX_HOST -p 2224"
    echo "    Gitea: http://$PROXMOX_HOST:3000"
    echo "    Direct NAT: ssh -i $SSH_KEY ubuntu@10.0.0.12 (from Proxmox host)"
    echo ""
    echo "  VM 103 (observability):"
    echo "    SSH: ssh -i $SSH_KEY ubuntu@$PROXMOX_HOST -p 2225"
    echo "    Prometheus: http://$PROXMOX_HOST:9090"
    echo "    Grafana: http://$PROXMOX_HOST:3001"
    echo "    Direct NAT: ssh -i $SSH_KEY ubuntu@10.0.0.13 (from Proxmox host)"
    echo ""
    log_info "To access VMs from Proxmox host:"
    echo "  ssh -i $SSH_KEY ubuntu@10.0.0.10  # VM 100"
    echo "  ssh -i $SSH_KEY ubuntu@10.0.0.11  # VM 101"
    echo "  ssh -i $SSH_KEY ubuntu@10.0.0.12  # VM 102"
    echo "  ssh -i $SSH_KEY ubuntu@10.0.0.13  # VM 103"
}

main() {
    log_step "Setup NAT for VMs"
    
    log_warn "This will:"
    log_warn "  1. Create a NAT bridge (vmbr1) on Proxmox host"
    log_warn "  2. Reconfigure VMs to use NAT network"
    log_warn "  3. Setup port forwarding for SSH and services"
    echo ""
    read -p "Continue? (yes/no): " confirm
    
    if [ "$confirm" != "yes" ]; then
        log_info "Cancelled"
        exit 0
    fi
    
    setup_nat_bridge
    
    log_step "Step 2: Configuring VMs for NAT"
    for vm_spec in "${VMS[@]}"; do
        read -r vmid name nat_ip <<< "$vm_spec"
        configure_vm_nat "$vmid" "$name" "$nat_ip" || log_warn "Failed to configure VM $vmid"
    done
    
    setup_port_forwarding
    
    log_info "Rebooting VMs to apply network changes..."
    ssh -i "$SSH_KEY" root@$PROXMOX_HOST "for vmid in 100 101 102 103; do qm reboot \$vmid 2>/dev/null || true; done"
    
    log_info "Waiting 60 seconds for VMs to reboot..."
    sleep 60
    
    show_access_info
    
    log_step "Testing NAT Access"
    log_info "Testing SSH via port forwarding..."
    if ssh -i "$SSH_KEY" -o ConnectTimeout=10 -p 2222 ubuntu@$PROXMOX_HOST "echo 'SSH OK' && hostname" &>/dev/null; then
        log_info "✓ SSH via NAT is working!"
    else
        log_warn "SSH may need more time. Wait a few minutes and test again."
    fi
}

main "$@"