#!/usr/bin/env python3
"""Real-world load balancing test for vpnem API."""
import urllib.request
import json
import sys

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

def req(path, method="GET", data=None, headers=None):
    """Make HTTP request."""
    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 urllib.error.HTTPError as e:
        print(f"  ERROR {e.code}: {e.read().decode()[:200]}")
        return None
    except Exception as e:
        print(f"  ERROR: {e}")
        return None

def get_recommend(client_ip):
    """Get recommendation for a client IP."""
    return req("/api/v1/recommend", headers={"X-Forwarded-For": client_ip})

def connect(client_ip, server_ip, node_id):
    """Report connection."""
    return req("/api/v1/connect", method="POST", data={
        "server_ip": server_ip, "node_id": node_id
    }, headers={"X-Forwarded-For": client_ip})

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

def main():
    print("=" * 70)
    print("VPNEM Load Balancing — Real Server Test")
    print("=" * 70)
    print()

    # Step 1: Test recommendations for multiple "studios"
    print("[1] Testing recommendations for 5 different studios...")
    print()
    studios = [
        ("195.10.20.1", "Barnaul Studio 1"),
        ("195.10.20.2", "Barnaul Studio 2"),
        ("91.50.60.1",  "Moscow Studio"),
        ("46.30.20.1",  "Novosibirsk Studio"),
        ("178.120.1.1", "Test Studio"),
    ]

    recommendations = {}
    for ip, name in studios:
        rec = get_recommend(ip)
        if rec:
            server = rec.get("recommended_server_ip", "none")
            reason = rec.get("reason", "unknown")
            load = rec.get("load_info", "no info")
            recommendations[ip] = server
            print(f"  {name:25s} → {server:15s} ({reason})")
            print(f"                            Load: {load}")
        else:
            print(f"  {name:25s} → ERROR")

    print()

    # Step 2: Check distribution — do different studios get different servers?
    print("[2] Distribution check...")
    servers_used = list(recommendations.values())
    unique = set(servers_used)
    print(f"  Studios: {len(studios)}, Recommended IPs: {len(unique)}")
    print(f"  Servers used: {', '.join(unique)}")
    if len(unique) > 1:
        print("  ✅ GOOD — different studios get different servers")
    else:
        print("  ⚠️  All studios got the same server (possible if load is equal)")

    print()

    # Step 3: Simulate connections
    print("[3] Simulating connections...")
    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",
    }
    for ip, name in studios:
        srv = recommendations.get(ip)
        nid = nodes.get(srv, "nl-multi-181")
        resp = connect(ip, srv, nid)
        if resp:
            rec_srv = resp.get("recommended_server_ip", "none")
            print(f"  {name:25s} connected → {srv:15s}  (next rec: {rec_srv})")
        else:
            print(f"  {name:25s} connect FAILED")

    print()

    # Step 4: Check recommendations AFTER connections (should shift)
    print("[4] Recommendations AFTER connections (should re-balance)...")
    new_ip = "100.100.100.1"
    rec = get_recommend(new_ip)
    if rec:
        srv = rec.get("recommended_server_ip", "none")
        reason = rec.get("reason", "unknown")
        load = rec.get("load_info", "no info")
        print(f"  New Studio ({new_ip}) → {srv:15s} ({reason})")
        print(f"                            Load: {load}")
        # The new studio should NOT get the most loaded server
    else:
        print(f"  ERROR getting recommendation")

    print()

    # Step 5: Disconnect one studio, check if recommendation changes
    print("[5] Disconnecting first studio and checking rebalance...")
    disc_ip = studios[0][0]
    disconnect(disc_ip)
    print(f"  Disconnected {studios[0][1]} ({disc_ip})")

    rec = get_recommend("200.200.200.1")
    if rec:
        srv = rec.get("recommended_server_ip", "none")
        load = rec.get("load_info", "no info")
        print(f"  New Studio (200.200.200.1) → {srv:15s}")
        print(f"                            Load: {load}")
    else:
        print("  ERROR")

    print()

    # Step 6: Clean up
    print("[6] Cleaning up — disconnecting all test studios...")
    for ip, name in studios:
        disconnect(ip)
    disconnect(new_ip)
    disconnect("200.200.200.1")
    print("  Done")

    print()
    print("=" * 70)
    print("Test complete.")
    print("=" * 70)

if __name__ == "__main__":
    main()