#!/usr/bin/env python3
"""Sequential load balancing test — each studio connects before next query."""
import urllib.request, json, sys

BASE = "https://vpn.em-sysadmin.xyz"

def req(path, method="GET", data=None, headers=None):
    url = BASE + path
    h = headers or {}
    h["Content-Type"] = "application/json"
    body = json.dumps(data).encode() if data else None
    rq = urllib.request.Request(url, data=body, headers=h, method=method)
    try:
        resp = urllib.request.urlopen(rq, timeout=10)
        return json.loads(resp.read().decode())
    except Exception as e:
        print(f"  ERROR: {e}")
        return None

def recommend(ip):
    return req("/api/v1/recommend", headers={"X-Forwarded-For": ip})

def connect(ip, srv, nid):
    return req("/api/v1/connect", method="POST", data={"server_ip": srv, "node_id": nid},
               headers={"X-Forwarded-For": ip})

def disconnect(ip):
    return req("/api/v1/disconnect", method="POST", data={"server_ip": ""},
               headers={"X-Forwarded-For": ip})

def main():
    print("=" * 70)
    print("Sequential Load Balancing Test")
    print("Each studio connects BEFORE the next one queries")
    print("=" * 70)
    print()

    studios = [
        ("195.10.20.1", "Barnaul-1"),
        ("195.10.20.2", "Barnaul-2"),
        ("91.50.60.1",  "Moscow"),
        ("46.30.20.1",  "Novosibirsk"),
        ("178.120.1.1", "Test"),
    ]
    nodes = {"5.180.97.181": "nl-multi-181", "5.180.97.197": "nl-multi-197",
             "5.180.97.198": "nl-multi-198", "5.180.97.199": "nl-multi-199"}

    assigned = {}

    # Step 1: Sequential recommend + connect
    print("[1] Sequential recommend + connect...")
    print()
    for ip, name in studios:
        rec = recommend(ip)
        if not rec:
            print(f"  {name:15s} → ERROR")
            continue
        srv = rec.get("recommended_server_ip", "none")
        reason = rec.get("reason", "unknown")
        load = rec.get("load_info", "")
        nid = nodes.get(srv, "nl-multi-181")
        assigned[ip] = srv

        print(f"  {name:15s} recommends {srv:15s} ({reason})")
        print(f"                load: {load}")

        # Connect immediately
        resp = connect(ip, srv, nid)
        if resp:
            next_srv = resp.get("recommended_server_ip", "none")
            print(f"                → connected. Next recommendation: {next_srv}")
        print()

    # Step 2: Check distribution
    print("[2] Distribution check...")
    servers_used = list(assigned.values())
    unique = set(servers_used)
    print(f"  Studios: {len(studios)}, Unique servers used: {len(unique)}")
    print(f"  Servers: {', '.join(sorted(unique))}")

    counts = {}
    for s in servers_used:
        counts[s] = counts.get(s, 0) + 1
    for srv, cnt in sorted(counts.items()):
        print(f"    {srv}: {cnt} studio(s)")

    if len(unique) > 1:
        print("  ✅ GOOD — studios distributed across servers")
    else:
        print("  ❌ All studios got the same server")

    print()

    # Step 3: Disconnect all and verify clean
    print("[3] Cleanup...")
    for ip, name in studios:
        disconnect(ip)
    print("  All test studios disconnected")

    # Step 4: Final check
    print()
    print("[4] Final recommendation (all clean)...")
    rec = recommend("999.999.999.1")
    if rec:
        srv = rec.get("recommended_server_ip", "none")
        load = rec.get("load_info", "no info")
        print(f"  New studio → {srv:15s}")
        print(f"               load: {load}")
    else:
        print("  ERROR")

    print()
    print("=" * 70)

if __name__ == "__main__":
    main()