"""
Ch14 配套代码 6 / 7 —— 压测脚本
================================================================

用途：
  对完整秒杀方案做压测，输出：
    - 总耗时 / QPS
    - 成功 / 失败分类计数
    - 是否超卖
    - Stream 消息数 vs 抢购成功数（一致性校验）

运行：
  python 01_init_stock.py 1001 1000
  python 06_load_test.py [threads] [per_thread]
  # 默认 200 线程 × 100 次 = 2 万请求
"""

import os
import sys
import time
import uuid
import random
import threading
from collections import Counter
import redis

ITEM_ID    = 1001
STOCK_KEY  = f"seckill:stock:{ITEM_ID}"
USERS_KEY  = f"seckill:users:{ITEM_ID}"
STREAM_KEY = "seckill:orders"

THREADS    = int(sys.argv[1]) if len(sys.argv) > 1 else 200
PER_THREAD = int(sys.argv[2]) if len(sys.argv) > 2 else 100
USER_POOL  = THREADS * 5

LUA_FILE = os.path.join(os.path.dirname(__file__), "lua", "seckill.lua")

results = Counter()
latencies = []
lat_lock = threading.Lock()


def worker(script):
    rr = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)
    local = Counter()
    local_lat = []
    for _ in range(PER_THREAD):
        user_id = f"u_{random.randint(1, USER_POOL)}"
        order_id = uuid.uuid4().hex
        t = time.perf_counter()
        ret = script(keys=[STOCK_KEY, USERS_KEY, STREAM_KEY],
                     args=[user_id, ITEM_ID, order_id],
                     client=rr)
        local_lat.append((time.perf_counter() - t) * 1000)
        local[{1: "ok", 0: "no_stock", -1: "dup", -2: "no_item"}[ret]] += 1
    with lat_lock:
        results.update(local)
        latencies.extend(local_lat)


def percentile(data, p):
    if not data: return 0
    data = sorted(data)
    idx = int(len(data) * p)
    return data[min(idx, len(data) - 1)]


def main():
    r = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)
    initial = int(r.get(STOCK_KEY) or 0)
    if initial == 0:
        print("❌ 请先运行: python 01_init_stock.py 1001 1000")
        sys.exit(1)

    script = r.register_script(open(LUA_FILE).read())

    print("=" * 64)
    print(f"  压测开始 · {THREADS} 线程 × {PER_THREAD} 请求 = {THREADS*PER_THREAD}")
    print(f"  初始库存  : {initial}")
    print(f"  用户池    : {USER_POOL}")
    print("=" * 64)

    t0 = time.time()
    ts = [threading.Thread(target=worker, args=(script,)) for _ in range(THREADS)]
    for t in ts: t.start()
    for t in ts: t.join()
    cost = time.time() - t0

    total = sum(results.values())
    final = int(r.get(STOCK_KEY) or 0)
    sold = initial - final
    msgs = r.xlen(STREAM_KEY)

    print()
    print("--- 性能 ---")
    print(f"  总请求    : {total}")
    print(f"  总耗时    : {cost:.2f} s")
    print(f"  QPS       : {total/cost:.0f}")
    print(f"  延迟 P50  : {percentile(latencies, 0.50):.2f} ms")
    print(f"  延迟 P95  : {percentile(latencies, 0.95):.2f} ms")
    print(f"  延迟 P99  : {percentile(latencies, 0.99):.2f} ms")

    print()
    print("--- 结果分类 ---")
    print(f"  抢购成功  : {results['ok']}")
    print(f"  库存不足  : {results['no_stock']}")
    print(f"  重复购买  : {results['dup']}")
    print(f"  商品不存在 : {results['no_item']}")

    print()
    print("--- 库存校验 ---")
    print(f"  剩余库存   : {final}")
    print(f"  实际卖出   : {sold}")
    print(f"  Stream 消息: {msgs}")
    if final < 0:
        print(f"  ❌ 出现负库存（超卖 {-final}）！")
    elif sold == results['ok'] == msgs:
        print(f"  ✅ 完美一致：库存 / 成功 / 消息 三方对齐")
    else:
        print(f"  ⚠️  数据有偏差，请检查")


if __name__ == "__main__":
    try:
        main()
    except redis.ConnectionError as e:
        print(f"❌ Redis 连接失败: {e}")