#!/bin/bash
source ~/.bashrc
# Setup NAT for VMs AND Reconfigure with SSH Keys
# Combines NAT setup with cloud-init SSH key injection

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}"
SSH_KEY_FILE="$SSH_KEY.pub"
PVE_USERNAME="${PVE_USERNAME:-root@pam}"
PVE_PASSWORD="${PVE_ROOT_PASS:-}"
PROXMOX_URL="${PROXMOX_ML110_URL:-https://192.168.1.206:8006}"
PROXMOX_NODE="${PROXMOX_NODE:-pve}"

# 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"
)

get_api_token() {
    local response=$(curl -s -k --connect-timeout 10 --max-time 15 \
        -d "username=$PVE_USERNAME&password=$PVE_PASSWORD" \
        "$PROXMOX_URL/api2/json/access/ticket" 2>&1)
    
    if echo "$response" | grep -q '"data"'; then
        local ticket=$(echo "$response" | grep -o '"ticket":"[^"]*' | cut -d'"' -f4)
        local csrf_token=$(echo "$response" | grep -o '"CSRFPreventionToken":"[^"]*' | cut -d'"' -f4)
        echo "$ticket|$csrf_token"
    else
        echo ""
    fi
}

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_with_ssh() {
    local vmid=$1
    local name=$2
    local nat_ip=$3
    
    log_info "Configuring VM $vmid ($name) with NAT IP $nat_ip and SSH keys..."
    
    if [ ! -f "$SSH_KEY_FILE" ]; then
        log_error "SSH key file not found: $SSH_KEY_FILE"
        return 1
    fi
    
    local ssh_key_content=$(cat "$SSH_KEY_FILE")
    local ssh_key_b64=$(echo "$ssh_key_content" | base64 -w 0)
    
    local tokens=$(get_api_token)
    if [ -z "$tokens" ]; then
        log_error "Failed to authenticate with Proxmox"
        return 1
    fi
    
    local ticket=$(echo "$tokens" | cut -d'|' -f1)
    local csrf_token=$(echo "$tokens" | cut -d'|' -f2)
    
    # Update VM network to use NAT bridge
    curl -s -k -X POST \
        -H "Cookie: PVEAuthCookie=$ticket" \
        -H "CSRFPreventionToken: $csrf_token" \
        -d "net0=virtio,bridge=$NAT_BRIDGE" \
        "$PROXMOX_URL/api2/json/nodes/$PROXMOX_NODE/qemu/$vmid/config" > /dev/null
    
    # Configure cloud-init with NAT IP and SSH keys
    curl -s -k -X POST \
        -H "Cookie: PVEAuthCookie=$ticket" \
        -H "CSRFPreventionToken: $csrf_token" \
        --data-urlencode "ipconfig0=ip=$nat_ip/24,gw=$NAT_GATEWAY" \
        --data-urlencode "ciuser=ubuntu" \
        --data-urlencode "sshkeys=$ssh_key_b64" \
        --data-urlencode "agent=1" \
        "$PROXMOX_URL/api2/json/nodes/$PROXMOX_NODE/qemu/$vmid/config" > /dev/null
    
    log_info "✓ VM $vmid configured for NAT with SSH keys"
}

setup_port_forwarding() {
    log_step "Step 3: Setting up Port Forwarding"
    
    log_info "Setting up port forwarding rules..."
    
    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"
}

main() {
    log_step "Setup NAT with SSH Keys"
    
    if [ ! -f "$SSH_KEY_FILE" ]; then
        log_error "SSH key file not found: $SSH_KEY_FILE"
        exit 1
    fi
    
    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. Inject SSH keys via cloud-init"
    log_warn "  4. 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 with SSH Keys"
    for vm_spec in "${VMS[@]}"; do
        read -r vmid name nat_ip <<< "$vm_spec"
        configure_vm_nat_with_ssh "$vmid" "$name" "$nat_ip" || log_warn "Failed to configure VM $vmid"
    done
    
    setup_port_forwarding
    
    log_info "Rebooting VMs to apply network and SSH key 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 90 seconds for VMs to reboot and apply cloud-init..."
    sleep 90
    
    log_step "Testing Access"
    log_info "Testing SSH via port forwarding..."
    local all_ok=true
    for port in 2222 2223 2224 2225; do
        local vm_name=$(echo $port | sed 's/2222/cloudflare-tunnel/;s/2223/k3s-master/;s/2224/git-server/;s/2225/observability/')
        if ssh -i "$SSH_KEY" -o ConnectTimeout=10 -p $port ubuntu@$PROXMOX_HOST "echo 'SSH OK' && hostname" &>/dev/null; then
            log_info "  ✓ $vm_name (port $port): SSH working"
        else
            log_warn "  ✗ $vm_name (port $port): SSH not working yet (may need more time)"
            all_ok=false
        fi
    done
    
    if [ "$all_ok" = true ]; then
        log_info ""
        log_info "✓ All VMs accessible via NAT with SSH!"
    else
        log_warn ""
        log_warn "Some VMs may need more time. Wait a few minutes and test again."
    fi
    
    log_step "Access Information"
    log_info "VM Access:"
    echo "  VM 100: ssh -i $SSH_KEY -p 2222 ubuntu@$PROXMOX_HOST"
    echo "  VM 101: ssh -i $SSH_KEY -p 2223 ubuntu@$PROXMOX_HOST"
    echo "  VM 102: ssh -i $SSH_KEY -p 2224 ubuntu@$PROXMOX_HOST"
    echo "  VM 103: ssh -i $SSH_KEY -p 2225 ubuntu@$PROXMOX_HOST"
    echo ""
    log_info "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 "$@"