import paramiko
import json

TARGET_VIPS = {
    ("monitoring", "uptime-lb"): "140.44.4.28",
    ("infra", "authentik-lb"): "140.44.4.29",
    ("awx", "awx-lb"): "140.44.4.30",
    ("infra", "technitium-npm-vip-service"): "140.44.4.31",
    ("default", "grafana-lb"): "140.44.4.32",
    ("infra", "unified-npm-vip-service"): "140.44.4.33",
    ("kube-system", "traefik"): "140.44.4.34",
    ("infra", "ak-outpost-authentik-ldap-outpost"): "140.44.4.35",
    ("longhorn-system", "longhorn-frontend-lb"): "140.44.4.36",
}

def main():
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect('140.44.4.11', username='root', key_filename=r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet', timeout=10)

    print("=== PINNING AERODROME LOADBALANCER VIPS ===")
    for (ns, name), target_ip in TARGET_VIPS.items():
        patch = json.dumps({
            "metadata": {
                "annotations": {
                    "metallb.universe.tf/loadBalancerIP": target_ip
                }
            },
            "spec": {
                "loadBalancerIP": target_ip
            }
        })
        
        cmd = f"kubectl patch svc {name} -n {ns} --type=merge -p '{patch}'"
        stdin, stdout, stderr = ssh.exec_command(cmd)
        out = stdout.read().decode().strip()
        err = stderr.read().decode().strip()
        print(f"[{ns}/{name}] -> {target_ip}: {out} {err}")

    print("\n=== VERIFYING FINAL ALLOCATION STATE ===")
    for (ns, name), target_ip in TARGET_VIPS.items():
        stdin, stdout, stderr = ssh.exec_command(f"kubectl get svc {name} -n {ns} -o json")
        svc = json.loads(stdout.read().decode())
        current_ingress = [x.get('ip') for x in svc.get('status', {}).get('loadBalancer', {}).get('ingress', [])]
        spec_ip = svc.get('spec', {}).get('loadBalancerIP')
        ann_ip = svc.get('metadata', {}).get('annotations', {}).get('metallb.universe.tf/loadBalancerIP')
        status = "LOCKED" if (spec_ip == target_ip and ann_ip == target_ip and target_ip in current_ingress) else "MISMATCH"
        print(f"[{status}] {ns}/{name}: Ingress={current_ingress}, Spec={spec_ip}, Annotation={ann_ip}")

    ssh.close()

if __name__ == '__main__':
    main()