"""
Radionica3D API Throughput Benchmark
=====================================
Measures: Order creation, status update, and deletion throughput.

Usage:
    cd backend
    python scratch/bench.py

Requires: pip install httpx
"""
import asyncio
import httpx
import time
import statistics
import os

BASE_URL = "http://localhost:8000"

# ──────────────────────────────────────────────
# CONFIG — fill in a valid admin token or leave
# empty to auto-login (requires correct creds).
# ──────────────────────────────────────────────
ADMIN_EMAIL    = "admin@radionica3d.com"
ADMIN_PASSWORD = "admin123"  # change if needed
MATERIAL_ID    = 1        # must exist in the DB
CONCURRENCY    = 10       # parallel requests per batch
ITERATIONS     = 30       # total operations per scenario
# ──────────────────────────────────────────────


async def get_admin_token(client: httpx.AsyncClient) -> str:
    resp = await client.post(
        f"{BASE_URL}/auth/login",
        json={"email": ADMIN_EMAIL, "password": ADMIN_PASSWORD},
    )
    resp.raise_for_status()
    token = resp.json().get("access_token")
    if not token:
        raise RuntimeError(f"No token in response: {resp.json()}")
    print("[OK] Logged in as admin")
    return token


async def create_one_order(client: httpx.AsyncClient, token: str, idx: int) -> tuple[float, int]:
    """Returns (latency_ms, status_code)"""
    # Use admin-created order (no file slicing, no rate limit)
    form = {
        "first_name":       ("", f"Bench{idx}"),
        "last_name":        ("", "Test"),
        "phone":            ("", "+38269000000"),
        "email":            ("", f"bench_{idx}@test.local"),
        "shipping_address": ("", "Bench Street 1"),
        "material_id":      ("", str(MATERIAL_ID)),
        "quantity":         ("", "1"),
        "color_name":       ("", "White"),
        "file_ids":         ("", "[]"),
        "file_quantities":  ("", "[]"),
    }
    t0 = time.monotonic()
    resp = await client.post(
        f"{BASE_URL}/orders",
        headers={"Authorization": f"Bearer {token}"},
        files=form,
    )
    latency = (time.monotonic() - t0) * 1000
    return latency, resp.status_code, resp.json().get("order_id") or resp.json().get("id")


async def update_order_status(client: httpx.AsyncClient, token: str, order_id: int) -> tuple[float, int]:
    t0 = time.monotonic()
    resp = await client.patch(
        f"{BASE_URL}/orders/{order_id}",
        headers={"Authorization": f"Bearer {token}"},
        json={"status": "processing"},
    )
    latency = (time.monotonic() - t0) * 1000
    body = resp.text[:200]
    return latency, resp.status_code, body


async def delete_order(client: httpx.AsyncClient, token: str, order_id: int) -> tuple[float, int]:
    t0 = time.monotonic()
    resp = await client.delete(
        f"{BASE_URL}/orders/{order_id}/admin",
        headers={"Authorization": f"Bearer {token}"},
    )
    latency = (time.monotonic() - t0) * 1000
    return latency, resp.status_code


def print_stats(label: str, latencies: list[float], duration: float, count: int):
    ok = len(latencies)
    rps = ok / duration
    rpm = rps * 60
    print(f"\n-- {label} --")
    print(f"  Requests sent   : {count}")
    print(f"  Succeeded (2xx) : {ok}")
    print(f"  Failed          : {count - ok}")
    if latencies:
        print(f"  Avg latency     : {statistics.mean(latencies):.0f} ms")
        print(f"  P50             : {statistics.median(latencies):.0f} ms")
        print(f"  P95             : {sorted(latencies)[int(len(latencies)*0.95)]:.0f} ms")
        print(f"  Max             : {max(latencies):.0f} ms")
    print(f"  Total time      : {duration:.2f}s")
    print(f"  Throughput      : {rps:.1f} req/s  ~  {rpm:.0f} req/min")


async def run_concurrent(tasks):
    """Run tasks in batches of CONCURRENCY"""
    results = []
    for i in range(0, len(tasks), CONCURRENCY):
        batch = tasks[i:i + CONCURRENCY]
        batch_results = await asyncio.gather(*batch, return_exceptions=True)
        results.extend(batch_results)
    return results


async def bench_create(client, token) -> list[int]:
    """Benchmark creation, return list of created order IDs."""
    tasks = [create_one_order(client, token, i) for i in range(ITERATIONS)]
    
    print(f"\n>>  Creating {ITERATIONS} orders ({CONCURRENCY} concurrent)...")
    t_start = time.monotonic()
    raw = await run_concurrent(tasks)
    duration = time.monotonic() - t_start

    latencies, order_ids = [], []
    for r in raw:
        if isinstance(r, Exception):
            print(f"  [error] {r}")
            continue
        lat, code, order_id = r
        if 200 <= code < 300 and order_id:
            latencies.append(lat)
            order_ids.append(order_id)
        else:
            print(f"  [fail] status={code}")

    print_stats("CREATE", latencies, duration, ITERATIONS)
    return order_ids


async def bench_update(client, token, order_ids: list[int]):
    """Benchmark status update for all created orders."""
    tasks = [update_order_status(client, token, oid) for oid in order_ids]

    print(f"\n>>  Updating {len(order_ids)} orders ({CONCURRENCY} concurrent)...")
    t_start = time.monotonic()
    raw = await run_concurrent(tasks)
    duration = time.monotonic() - t_start

    latencies = []
    for r in raw:
        if isinstance(r, Exception):
            print(f"  [error] {r}")
            continue
        lat, code, body = r
        if 200 <= code < 300:
            latencies.append(lat)
        else:
            print(f"  [fail] status={code} body={body}")

    print_stats("UPDATE STATUS", latencies, duration, len(order_ids))


async def bench_delete(client, token, order_ids: list[int]):
    """Benchmark deletion for all created orders."""
    tasks = [delete_order(client, token, oid) for oid in order_ids]

    print(f"\n>>  Deleting {len(order_ids)} orders ({CONCURRENCY} concurrent)...")
    t_start = time.monotonic()
    raw = await run_concurrent(tasks)
    duration = time.monotonic() - t_start

    latencies = []
    for r in raw:
        if isinstance(r, Exception):
            continue
        lat, code = r
        if 200 <= code < 300:
            latencies.append(lat)

    print_stats("DELETE", latencies, duration, len(order_ids))


async def main():
    print("=" * 50)
    print(" Radionica3D API Throughput Benchmark")
    print(f" Target : {BASE_URL}")
    print(f" Concurrency : {CONCURRENCY}   Iterations : {ITERATIONS}")
    print("=" * 50)

    limits = httpx.Limits(max_connections=CONCURRENCY + 5, max_keepalive_connections=CONCURRENCY)
    async with httpx.AsyncClient(timeout=30.0, limits=limits) as client:
        try:
            token = await get_admin_token(client)
        except Exception as e:
            print(f"Login failed: {e}")
            return

        # 1. Create
        order_ids = await bench_create(client, token)
        if not order_ids:
            print("\nNo orders created, stopping.")
            return

        # 2. Update
        await bench_update(client, token, order_ids)

        # 3. Delete
        await bench_delete(client, token, order_ids)

    print("\n[OK] Benchmark complete. Check server logs for any errors.")


if __name__ == "__main__":
    asyncio.run(main())
