"""
Players, not askers: each virtual student loops home -> island quiz -> answers -> lesson -> summary, with a
random 1..THINK_MAX s pause between requests. Stages of N users for D seconds each; optional AI asks.
usage: python play.py SOCKET users.json OUT.json "300:40,600:40,1000:60" THINK_MAX ASKS_IN_LAST_STAGE
"""
import asyncio, json, random, sys, time
import httpx

SOCK, USERS, OUT, STAGES, THINK, ASKS = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], float(sys.argv[5]), int(sys.argv[6])
PROC, NPROC = int(sys.argv[7]), int(sys.argv[8])
users = json.load(open(USERS))
stages = [tuple(map(int, s.split(':'))) for s in STAGES.split(',')]
H = {'Host': 'mathapi.tadrisino2.ir', 'X-Forwarded-Proto': 'https'}
recs = []


async def player(c, u, stop, stage):
    rnd = random.Random(u['userId'])
    hdr = {**H, 'Authorization': f"Bearer {u['token']}"}
    steps = ['/api/home', f"/api/nodes/{u['nodeId']}/quiz"] + ['answer'] * 4 + [f"/api/lessons/{u['lessonId']}/questions", '/api/user/summary']
    i = rnd.randrange(len(steps))
    await asyncio.sleep(rnd.uniform(0, THINK))
    while not stop.is_set():
        p = steps[i % len(steps)]; i += 1
        t = time.monotonic()
        try:
            if p == 'answer':
                q = rnd.choice(u['questions'])
                r = await c.post(f"http://x/api/questions/{q['id']}/answer", json={'selectedOptionId': rnd.choice(q['options'])}, headers=hdr)
                name = 'answer (POST)'
            else:
                r = await c.get(f'http://x{p}', headers=hdr)
                name = p.split('/')[2] if p.count('/') <= 3 else p.split('/')[2] + '/' + p.split('/')[-1]
            recs.append((stage, name, r.status_code, time.monotonic() - t))
        except Exception as e:
            recs.append((stage, p if p == 'answer' else p.split('/')[2], type(e).__name__, time.monotonic() - t))
        await asyncio.sleep(rnd.uniform(1, THINK))


async def ask(c, u, delay, stage, out):
    await asyncio.sleep(delay)
    t = time.monotonic()
    try:
        r = await c.post('http://x/api/ai/ask', json={'userQuestion': f'این را قدم به قدم توضیح بده: سوالِ شماره‌ی {u["userId"]} درباره‌ی همین درس', 'lessonId': str(u['lessonId'])},
                         headers={**H, 'Authorization': f"Bearer {u['token']}"}, timeout=120)
        out.append((stage, 'ai/ask', r.status_code, time.monotonic() - t))
    except Exception as e:
        out.append((stage, 'ai/ask', type(e).__name__, time.monotonic() - t))


async def main():
    tr = httpx.AsyncHTTPTransport(uds=SOCK, retries=0)
    lim = httpx.Limits(max_connections=200, max_keepalive_connections=200)
    async with httpx.AsyncClient(transport=tr, limits=lim, timeout=60) as c:
        summary = []
        for n, dur in stages:
            stop = asyncio.Event(); label = f'{n}'
            start = len(recs); t0 = time.monotonic()
            tasks = [asyncio.create_task(player(c, u, stop, label)) for i, u in enumerate(users[:n]) if i % NPROC == PROC]
            ai = []
            if ASKS and (n, dur) == stages[-1]:
                tasks += [asyncio.create_task(ask(c, u, k * (dur - 10) / ASKS, label, ai)) for k, u in enumerate(users[-ASKS:]) if k % NPROC == PROC]
            await asyncio.sleep(dur); stop.set()
            await asyncio.gather(*tasks)
            recs.extend(ai)
            el = time.monotonic() - t0
            mine = recs[start:]
            lat = sorted(x[3] for x in mine if x[1] != 'ai/ask')
            q = lambda a, p: round(a[min(len(a) - 1, int(p * len(a)))] * 1000) if a else None
            codes = {}
            for x in mine: codes[str(x[2])] = codes.get(str(x[2]), 0) + 1
            s = {'users': n, 'seconds': round(el), 'requests': len(mine), 'rps': round(len([x for x in mine if x[1] != 'ai/ask']) / el, 1),
                 'p50_ms': q(lat, .5), 'p95_ms': q(lat, .95), 'p99_ms': q(lat, .99), 'max_ms': q(lat, 1), 'codes': codes}
            per = {}
            for x in mine:
                per.setdefault(x[1], []).append(x[3])
            s['per'] = {k: {'n': len(v), 'p50': q(sorted(v), .5), 'p95': q(sorted(v), .95)} for k, v in per.items()}
            summary.append(s)
            await asyncio.sleep(5)
        json.dump(recs, open(OUT, 'w'))

asyncio.run(main())