#!/usr/bin/env python3
"""
UCall Load & Concurrency Test
──────────────────────────────
Targets from the spec:

  - 100–500 concurrent WebSocket connections
  - 1,000+ signaling operations
  - Concurrent Universal ID allocation
  - Concurrent presence heartbeats
  - Call-routing contention

Records:
  - p50 / p95 / p99 latency
  - throughput
  - error rate
  - WebSocket connection stability

Usage:
    python3 tests/load_test.py [--connections 200] [--ops 1000]
"""

import asyncio
import json
import sys
import time
import uuid
import argparse
import statistics
import random
import urllib.request
from dataclasses import dataclass, field
from typing import Optional

try:
    import websockets
except ImportError:
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets", "-q"])
    import websockets

BASE = "http://127.0.0.1:8080"
WS   = "ws://127.0.0.1:8080/api/v1/ws"


# ─── Helpers ──────────────────────────────────────────────────────────────────

def percentile(data: list[float], p: int) -> float:
    if not data:
        return 0.0
    sorted_data = sorted(data)
    idx = max(0, int(len(sorted_data) * p / 100) - 1)
    return sorted_data[idx]

def rand_phone() -> str:
    return f"+1{random.randint(2000000000, 9999999999)}"

def register_sync(phone: str) -> Optional[dict]:
    data = json.dumps({"phone": phone, "display_name": "Load"}).encode()
    req  = urllib.request.Request(
        f"{BASE}/api/v1/auth/register", data=data,
        headers={"Content-Type": "application/json"}
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as r:
            return json.loads(r.read()).get("data")
    except Exception:
        return None

def ws_url(token: str) -> str:
    return f"{WS}?token={token}&endpoint_id={uuid.uuid4()}"

def print_latency_report(label: str, latencies: list[float], errors: int, total: int):
    if not latencies:
        print(f"  {label}: NO DATA (errors={errors}/{total})")
        return
    print(f"  {label}:")
    print(f"    count={len(latencies)}, errors={errors}/{total}")
    print(f"    p50={percentile(latencies, 50):.1f}ms "
          f"p95={percentile(latencies, 95):.1f}ms "
          f"p99={percentile(latencies, 99):.1f}ms")
    print(f"    min={min(latencies):.1f}ms max={max(latencies):.1f}ms "
          f"avg={statistics.mean(latencies):.1f}ms")
    error_rate = errors / total * 100 if total else 0
    print(f"    error_rate={error_rate:.1f}%")


# ─────────────────────────────────────────────────────────────────────────────
# TEST 1: Concurrent Universal ID Allocation
# ─────────────────────────────────────────────────────────────────────────────

async def load_uid_allocation(count: int) -> tuple[list[float], int, set]:
    """Register `count` users concurrently and verify UIDs are unique and valid."""
    latencies = []
    errors = 0
    uid_set = set()

    async def register_one(session):
        phone = rand_phone()
        data  = json.dumps({"phone": phone, "display_name": "Load"}).encode()
        headers = {"Content-Type": "application/json"}
        t0 = time.perf_counter()
        try:
            req = urllib.request.Request(f"{BASE}/api/v1/auth/register", data=data, headers=headers)
            with urllib.request.urlopen(req, timeout=10) as r:
                body = json.loads(r.read())
                ms = (time.perf_counter() - t0) * 1000
                uid = body.get("data", {}).get("universal_id", "")
                return ms, uid
        except Exception:
            return None, None

    tasks = [asyncio.get_event_loop().run_in_executor(None, register_sync, rand_phone())
             for _ in range(count)]

    t_start = time.perf_counter()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    elapsed = time.perf_counter() - t_start

    for r in results:
        if isinstance(r, Exception) or r is None:
            errors += 1
            continue
        uid = r.get("universal_id", "") if r else ""
        if not uid or len(uid) != 12 or not uid.isdigit():
            errors += 1
        else:
            if uid in uid_set:
                errors += 1  # Collision
            uid_set.add(uid)
        latencies.append(0)  # sync registration — track errors only

    print(f"  Throughput: {count/elapsed:.1f} registrations/sec over {elapsed:.2f}s")
    return latencies, errors, uid_set


# ─────────────────────────────────────────────────────────────────────────────
# TEST 2: Concurrent WebSocket Connections
# ─────────────────────────────────────────────────────────────────────────────

async def load_ws_connections(token: str, num_connections: int) -> tuple[list[float], int]:
    """Open `num_connections` WebSockets concurrently and measure connect latency."""
    latencies = []
    errors = 0
    semaphore = asyncio.Semaphore(50)  # max 50 concurrent connect attempts

    async def connect_one():
        nonlocal errors
        url = ws_url(token)
        t0  = time.perf_counter()
        async with semaphore:
            try:
                ws = await asyncio.wait_for(websockets.connect(url), timeout=10)
                ms = (time.perf_counter() - t0) * 1000
                latencies.append(ms)
                await asyncio.sleep(0.5)  # hold connection briefly
                await ws.close()
            except Exception:
                errors += 1

    print(f"  Opening {num_connections} concurrent WebSocket connections...")
    t_start = time.perf_counter()
    await asyncio.gather(*[connect_one() for _ in range(num_connections)])
    elapsed = time.perf_counter() - t_start
    print(f"  Throughput: {num_connections/elapsed:.1f} connects/sec")

    return latencies, errors


# ─────────────────────────────────────────────────────────────────────────────
# TEST 3: Concurrent Presence Heartbeats
# ─────────────────────────────────────────────────────────────────────────────

async def load_presence_heartbeats(token: str, num_clients: int, heartbeats_each: int):
    """Open `num_clients` WS connections and each sends `heartbeats_each` heartbeats."""
    latencies = []
    errors = 0
    semaphore = asyncio.Semaphore(30)

    async def heartbeat_client():
        nonlocal errors
        url = ws_url(token)
        async with semaphore:
            try:
                async with websockets.connect(url, open_timeout=10) as ws:
                    for _ in range(heartbeats_each):
                        t0 = time.perf_counter()
                        await ws.send(json.dumps({"method": "presence.heartbeat", "params": {}}))
                        latencies.append((time.perf_counter() - t0) * 1000)
                        await asyncio.sleep(0.05)
            except Exception:
                errors += 1

    print(f"  {num_clients} clients × {heartbeats_each} heartbeats = {num_clients*heartbeats_each} ops...")
    t_start = time.perf_counter()
    await asyncio.gather(*[heartbeat_client() for _ in range(num_clients)])
    elapsed = time.perf_counter() - t_start
    total_ops = num_clients * heartbeats_each
    print(f"  Throughput: {total_ops/elapsed:.1f} heartbeats/sec")
    return latencies, errors, total_ops


# ─────────────────────────────────────────────────────────────────────────────
# TEST 4: Concurrent Call Signaling
# ─────────────────────────────────────────────────────────────────────────────

async def load_call_signaling(pairs: int) -> tuple[list[float], int]:
    """Run `pairs` call invite→accept flows concurrently. Returns RTTs and errors."""
    rtts   = []
    errors = 0
    semaphore = asyncio.Semaphore(20)

    async def one_call():
        nonlocal errors
        async with semaphore:
            # Fresh users for each pair
            user_a = await asyncio.get_event_loop().run_in_executor(None, register_sync, rand_phone())
            user_b = await asyncio.get_event_loop().run_in_executor(None, register_sync, rand_phone())
            if not user_a or not user_b:
                errors += 1
                return

            tok_a  = user_a["access_token"]
            tok_b  = user_b["access_token"]
            uid_b  = user_b["universal_id"]

            try:
                async with websockets.connect(ws_url(tok_a), open_timeout=10) as ws_a, \
                           websockets.connect(ws_url(tok_b), open_timeout=10) as ws_b:

                    t0 = time.perf_counter()
                    await ws_a.send(json.dumps({
                        "method": "call.invite",
                        "params": {
                            "target_universal_id": uid_b,
                            "call_type": "audio",
                            "sdp_offer": "v=0\r\n"
                        }
                    }))

                    # Wait for call.incoming on B
                    call_id = None
                    deadline = time.perf_counter() + 5
                    while time.perf_counter() < deadline:
                        try:
                            raw = await asyncio.wait_for(ws_b.recv(), timeout=1)
                            msg = json.loads(raw)
                            if msg.get("method") == "call.incoming":
                                call_id = msg["params"]["call_id"]
                                break
                        except asyncio.TimeoutError:
                            break

                    if not call_id:
                        errors += 1
                        return

                    await ws_b.send(json.dumps({
                        "method": "call.accept",
                        "params": {"call_id": call_id, "sdp_answer": "v=0\r\n"}
                    }))

                    # Wait for call.connected on A
                    deadline2 = time.perf_counter() + 5
                    while time.perf_counter() < deadline2:
                        try:
                            raw = await asyncio.wait_for(ws_a.recv(), timeout=1)
                            msg = json.loads(raw)
                            if msg.get("method") == "call.connected":
                                rtt = (time.perf_counter() - t0) * 1000
                                rtts.append(rtt)
                                return
                        except asyncio.TimeoutError:
                            break

                    errors += 1  # call.connected never received

            except Exception:
                errors += 1

    print(f"  Running {pairs} concurrent call flows...")
    t_start = time.perf_counter()
    await asyncio.gather(*[one_call() for _ in range(pairs)])
    elapsed = time.perf_counter() - t_start
    total   = pairs
    successful = len(rtts)
    print(f"  Completed: {successful}/{total} calls in {elapsed:.2f}s")
    if successful:
        print(f"  Throughput: {successful/elapsed:.2f} calls/sec")
    return rtts, errors


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

async def main(connections: int, ops: int, call_pairs: int):
    print()
    print("╔══════════════════════════════════════════════════════════════╗")
    print("║         UCall — Load & Concurrency Test                     ║")
    print("╚══════════════════════════════════════════════════════════════╝")
    print(f"  Target: {BASE}")
    print(f"  Params: connections={connections}, ops={ops}, call_pairs={call_pairs}")
    print()

    # Seed a shared token (we'll reuse for WS tests to avoid registration overhead)
    seed_user = register_sync(rand_phone())
    if not seed_user:
        print("❌  Server not reachable or registration failed. Is ucall-server running?")
        return 1
    shared_token = seed_user["access_token"]

    all_pass = True

    # ── 1. UID Allocation ─────────────────────────────────────────────────
    print("━━  1. Concurrent Universal ID Allocation  ━━")
    uid_count = min(ops, 200)
    _, uid_errors, uid_set = await load_uid_allocation(uid_count)
    ok = uid_errors == 0 and len(uid_set) == uid_count
    print(f"  UIDs allocated: {len(uid_set)}/{uid_count}, "
          f"collisions/errors={uid_errors}, unique={len(uid_set)==uid_count}")
    if not ok:
        all_pass = False
        print(f"  ❌ UID allocation: {uid_errors} errors detected")
    else:
        print(f"  ✅ UID allocation: zero collisions")

    # ── 2. Concurrent WS Connections ─────────────────────────────────────
    print()
    print("━━  2. Concurrent WebSocket Connections  ━━")
    ws_lats, ws_errors = await load_ws_connections(shared_token, connections)
    print_latency_report("WS Connect Latency", ws_lats, ws_errors, connections)
    ws_error_rate = ws_errors / connections * 100 if connections else 0
    if ws_error_rate > 5:
        all_pass = False
        print(f"  ❌ WS error rate {ws_error_rate:.1f}% exceeds 5% threshold")
    else:
        print(f"  ✅ WS connection stability: error_rate={ws_error_rate:.1f}%")

    # ── 3. Concurrent Presence Heartbeats ────────────────────────────────
    print()
    print("━━  3. Concurrent Presence Heartbeats  ━━")
    hb_clients = min(50, connections // 2)
    hb_each    = max(5, ops // hb_clients)
    hb_lats, hb_errors, hb_total = await load_presence_heartbeats(
        shared_token, hb_clients, hb_each)
    print_latency_report("Heartbeat Send Latency", hb_lats, hb_errors, hb_total)
    hb_error_rate = hb_errors / hb_clients * 100 if hb_clients else 0
    if hb_error_rate > 5:
        all_pass = False
    else:
        print(f"  ✅ Presence heartbeats: error_rate={hb_error_rate:.1f}%")

    # ── 4. Concurrent Call Signaling ──────────────────────────────────────
    print()
    print("━━  4. Concurrent Call Signaling  ━━")
    call_rtts, call_errors = await load_call_signaling(call_pairs)
    print_latency_report("Call Signaling RTT", call_rtts, call_errors, call_pairs)
    call_error_rate = call_errors / call_pairs * 100 if call_pairs else 0
    if call_error_rate > 10:
        all_pass = False
        print(f"  ❌ Call signaling error rate {call_error_rate:.1f}% exceeds 10%")
    else:
        print(f"  ✅ Call signaling: error_rate={call_error_rate:.1f}%")

    # ── Summary ───────────────────────────────────────────────────────────
    print()
    print("══════════════════════════════════════════════════════════════")
    if all_pass:
        print("  Load Test: ✅  ALL THRESHOLDS MET")
    else:
        print("  Load Test: ⚠️   SOME THRESHOLDS EXCEEDED")
    print("══════════════════════════════════════════════════════════════")

    return 0 if all_pass else 1


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="UCall Load Test")
    parser.add_argument("--connections", type=int, default=100,
                        help="Concurrent WebSocket connections (default: 100)")
    parser.add_argument("--ops", type=int, default=500,
                        help="Total signaling operations (default: 500)")
    parser.add_argument("--call-pairs", type=int, default=20,
                        help="Concurrent call signaling pairs (default: 20)")
    args = parser.parse_args()

    sys.exit(asyncio.run(main(args.connections, args.ops, args.call_pairs)))