#!/bin/bash
# PVE RAM-Based Rebalancer
# Runs via cron. Checks RAM overcommit on all PVE nodes.
# Migrates HA-managed VMs from overcommitted nodes to nodes with headroom.
# Respects HA node-affinity rules (strict rules constrain target nodes).
# Only outputs when it performs a migration (silent otherwise).
#
# Overcommit = sum(running VM maxmem on node) / node physical RAM
# A node is "overcommitted" when this ratio exceeds THRESHOLD_PCT.

THRESHOLD_PCT=100
MIN_FREE_GB=4
MAX_MIGRATIONS=2
SSH_KEY=~/.ssh/id_ed25519_proxmox
PVE_API_NODE="10.0.20.10"

TMP_DIR=$(mktemp -d)
trap "rm -rf $TMP_DIR" EXIT

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -o ConnectTimeout=10 root@$PVE_API_NODE \
  'pvesh get /cluster/resources --type vm --output-format json 2>/dev/null' \
  > "$TMP_DIR/vms.json" 2>/dev/null

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -o ConnectTimeout=10 root@$PVE_API_NODE \
  'pvesh get /cluster/resources --type node --output-format json 2>/dev/null' \
  > "$TMP_DIR/nodes.json" 2>/dev/null

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -o ConnectTimeout=10 root@$PVE_API_NODE \
  'cat /etc/pve/ha/rules.cfg 2>/dev/null' \
  > "$TMP_DIR/rules.cfg" 2>/dev/null

ssh -i $SSH_KEY -o StrictHostKeyChecking=no -o ConnectTimeout=10 root@$PVE_API_NODE \
  'ha-manager config 2>/dev/null | grep "^vm:" | cut -d: -f2' \
  > "$TMP_DIR/ha_vms.txt" 2>/dev/null

export REAL_TMP_DIR="$TMP_DIR"
python3 << PYEOF
import json, sys, os

TMP_DIR = os.environ.get("REAL_TMP_DIR", "/tmp")
THRESHOLD_PCT = 100
MIN_FREE_GB = 4
MAX_MIGRATIONS = 2

with open(f"{TMP_DIR}/vms.json") as f:
    vms_raw = json.load(f)
with open(f"{TMP_DIR}/nodes.json") as f:
    nodes_raw = json.load(f)

# Parse HA rules — extract strict node-affinity constraints
vm_node_constraints = {}
current_rule = None
rules_path = f"{TMP_DIR}/rules.cfg"
if os.path.exists(rules_path):
    with open(rules_path) as f:
        for line in f:
            line = line.strip()
            if line.startswith("node-affinity:"):
                if current_rule and current_rule.get("resources") and current_rule.get("strict") == 1:
                    for vid in current_rule["resources"]:
                        vm_node_constraints[vid] = set(current_rule["nodes"])
                current_rule = {"nodes": [], "resources": [], "strict": 0}
            elif current_rule is not None:
                if line.startswith("nodes "):
                    current_rule["nodes"] = [n.strip() for n in line.split(None,1)[1].split(",")]
                elif line.startswith("resources "):
                    for p in line.split(None,1)[1].split(","):
                        p = p.strip()
                        if ":" in p:
                            typ, vid = p.split(":")
                            if typ.strip() == "vm":
                                current_rule["resources"].append(int(vid.strip()))
                elif line.startswith("strict "):
                    current_rule["strict"] = int(line.split()[1])
        if current_rule and current_rule.get("resources") and current_rule.get("strict") == 1:
            for vid in current_rule["resources"]:
                vm_node_constraints[vid] = set(current_rule["nodes"])

ha_vms = set()
ha_path = f"{TMP_DIR}/ha_vms.txt"
if os.path.exists(ha_path):
    with open(ha_path) as f:
        for line in f:
            vid = line.strip()
            if vid.isdigit():
                ha_vms.add(int(vid))

nodes = {}
for n in nodes_raw:
    name = n.get("node", "")
    maxmem = n.get("maxmem", 0)
    if maxmem > 0:
        nodes[name] = {
            "maxmem": maxmem,
            "maxmem_gb": round(maxmem / 1024**3, 1),
        }

vms = []
node_allocated = {}
for v in vms_raw:
    if v.get("status") != "running":
        continue
    vmid = v.get("vmid")
    nodename = v.get("node", "")
    maxmem = v.get("maxmem", 0)
    name = v.get("name", "?")
    if nodename in nodes and maxmem > 0:
        vms.append({
            "vmid": vmid,
            "name": name,
            "node": nodename,
            "maxmem_mb": maxmem // 1024**2,
            "maxmem_gb": round(maxmem / 1024**3, 1),
            "ha_managed": vmid in ha_vms,
            "constrained_nodes": vm_node_constraints.get(vmid),
        })
        node_allocated[nodename] = node_allocated.get(nodename, 0) + maxmem

for nodename, info in nodes.items():
    allocated = node_allocated.get(nodename, 0)
    info["allocated_mb"] = allocated // 1024**2
    info["allocated_gb"] = round(allocated / 1024**3, 1)
    max_mb = info["maxmem"] // 1024**2
    info["overcommit_pct"] = round(info["allocated_mb"] / max_mb * 100) if max_mb > 0 else 0
    info["free_for_vms_gb"] = round((info["maxmem"] - allocated) / 1024**3, 1)

actions = []
for nodename, info in sorted(nodes.items(), key=lambda x: x[1]["overcommit_pct"], reverse=True):
    if info["overcommit_pct"] <= THRESHOLD_PCT:
        break
    node_vms = sorted([v for v in vms if v["node"] == nodename and v["ha_managed"]],
                     key=lambda x: x["maxmem_mb"], reverse=True)
    for vm in node_vms:
        allowed_targets = set(nodes.keys()) - {nodename}
        if vm["constrained_nodes"]:
            allowed_targets = vm["constrained_nodes"] - {nodename}
        candidates = []
        for n, ni in nodes.items():
            if n not in allowed_targets:
                continue
            free_after = ni["free_for_vms_gb"] - vm["maxmem_gb"]
            if ni["free_for_vms_gb"] >= MIN_FREE_GB and free_after >= 0:
                new_target_alloc = ni["allocated_mb"] + vm["maxmem_mb"]
                new_target_pct = round(new_target_alloc / (ni["maxmem"] // 1024**2) * 100)
                if new_target_pct <= THRESHOLD_PCT:
                    candidates.append((n, ni, new_target_pct))
        if candidates:
            candidates.sort(key=lambda x: x[1]["free_for_vms_gb"], reverse=True)
            target, target_info, new_target_pct = candidates[0]
            new_source_pct = round((info["allocated_mb"] - vm["maxmem_mb"]) / (info["maxmem"] // 1024**2) * 100)
            actions.append({
                "vmid": vm["vmid"],
                "name": vm["name"],
                "from": nodename,
                "to": target,
                "ram_gb": vm["maxmem_gb"],
                "source_before": info["overcommit_pct"],
                "source_after": new_source_pct,
                "target_before": target_info["overcommit_pct"],
                "target_after": new_target_pct,
            })
            info["allocated_mb"] -= vm["maxmem_mb"]
            info["free_for_vms_gb"] += vm["maxmem_gb"]
            target_info["allocated_mb"] += vm["maxmem_mb"]
            target_info["free_for_vms_gb"] -= vm["maxmem_gb"]
            break

if not actions:
    sys.exit(0)

for a in actions[:MAX_MIGRATIONS]:
    print(f"MIGRATE {a['vmid']} {a['from']} {a['to']}")
    print(f"DETAILS VM {a['vmid']} ({a['name']}, {a['ram_gb']}GB): "
          f"{a['from']} ({a['source_before']}%\u2192{a['source_after']}%) \u2192 "
          f"{a['to']} ({a['target_before']}%\u2192{a['target_after']}%)")
PYEOF
