"""
Ch7 配套代码 3 / 4 —— Lua 库存扣减实战

演示：
  1. 用纯客户端「GET → 判断 → DECR」会怎么超卖
  2. 用 Lua 原子脚本扣减，1000 并发也精确不超卖
  3. 演示 EVALSHA 缓存、KEYS / ARGV 用法
"""

import threading
import time
import redis

POOL = redis.ConnectionPool(host="127.0.0.1", port=6379, decode_responses=True)


def section(title: str) -> None:
    print("\n" + "=" * 62)
    print(title)
    print("=" * 62)


# ================ Lua 脚本 ================
# KEYS[1]: 库存 key
# ARGV[1]: 要扣的数量
# 返回:    -1 表示库存不足，否则返回剩余库存
LUA_DECR_STOCK = """
local stock = tonumber(redis.call('GET', KEYS[1]) or '0')
local need  = tonumber(ARGV[1])
if stock < need then
    return -1
end
redis.call('DECRBY', KEYS[1], need)
return stock - need
"""

# 滑动窗口限流：1 秒内最多 N 次
# KEYS[1]: 限流 key
# ARGV[1]: 限制次数  ARGV[2]: 窗口秒
LUA_RATE_LIMIT = """
local cur = tonumber(redis.call('GET', KEYS[1]) or '0')
if cur >= tonumber(ARGV[1]) then
    return 0
end
redis.call('INCR', KEYS[1])
if cur == 0 then
    redis.call('EXPIRE', KEYS[1], ARGV[2])
end
return 1
"""


def demo_unsafe_oversell() -> None:
    section("Demo 1: 不用 Lua —— 客户端 GET → 判断 → DECR 会超卖")
    r = redis.Redis(connection_pool=POOL)
    INIT_STOCK = 100
    THREADS = 1000
    BUY_PER = 1

    r.set("stock:item", INIT_STOCK)
    success = []
    lock = threading.Lock()

    def buyer():
        rr = redis.Redis(connection_pool=POOL)
        cur = int(rr.get("stock:item") or 0)
        if cur >= BUY_PER:
            time.sleep(0.0005)  # 模拟一点业务处理时间，制造窗口期
            rr.decrby("stock:item", BUY_PER)
            with lock: success.append(1)

    threads = [threading.Thread(target=buyer) for _ in range(THREADS)]
    for t in threads: t.start()
    for t in threads: t.join()

    final = int(r.get("stock:item"))
    sold = sum(success)
    print(f"  初始库存 {INIT_STOCK}，{THREADS} 并发各买 {BUY_PER} 件")
    print(f"  实际成功购买 = {sold}")
    print(f"  剩余库存     = {final}")
    if final < 0 or sold > INIT_STOCK:
        print(f"  ❌ 超卖！（多卖了 {sold - INIT_STOCK} 件）")
    else:
        print(f"  本次没复现超卖（高并发下大概率会发生）")


def demo_lua_safe() -> None:
    section("Demo 2: 用 Lua 原子扣减 —— 1000 并发也不超卖")
    r = redis.Redis(connection_pool=POOL)
    INIT_STOCK = 100
    THREADS = 1000

    r.set("stock:item", INIT_STOCK)
    sha = r.script_load(LUA_DECR_STOCK)
    print(f"  脚本已缓存 SHA = {sha[:16]}...")

    success = []
    lock = threading.Lock()

    def buyer():
        rr = redis.Redis(connection_pool=POOL)
        try:
            res = rr.evalsha(sha, 1, "stock:item", 1)
        except redis.exceptions.NoScriptError:
            res = rr.eval(LUA_DECR_STOCK, 1, "stock:item", 1)
        if res != -1:
            with lock: success.append(1)

    threads = [threading.Thread(target=buyer) for _ in range(THREADS)]
    t0 = time.time()
    for t in threads: t.start()
    for t in threads: t.join()
    elapsed = time.time() - t0

    final = int(r.get("stock:item"))
    sold = sum(success)
    print(f"\n  耗时 {elapsed:.2f}s")
    print(f"  实际成功购买 = {sold}    (期望 {INIT_STOCK})")
    print(f"  剩余库存     = {final}    (期望 0)")
    if sold == INIT_STOCK and final == 0:
        print("  ✅ 完美：原子扣减，零超卖")
    else:
        print("  ❌ 异常")


def demo_rate_limit() -> None:
    section("Demo 3: Lua 滑动计数限流 —— 1 秒内最多 5 次")
    r = redis.Redis(connection_pool=POOL)
    r.delete("rl:user:1001")

    sha = r.script_load(LUA_RATE_LIMIT)
    for i in range(8):
        ok = r.evalsha(sha, 1, "rl:user:1001", 5, 1)
        flag = "✅ 通过" if ok else "🚫 限流"
        print(f"  第 {i+1} 次请求 → {flag}")
        time.sleep(0.05)


def demo_keys_argv() -> None:
    section("Demo 4: KEYS / ARGV 区别 + 多 key 操作")
    r = redis.Redis(connection_pool=POOL)
    script = """
        return {
          'KEYS[1]=' .. KEYS[1],
          'KEYS[2]=' .. KEYS[2],
          'ARGV[1]=' .. ARGV[1],
          'ARGV[2]=' .. ARGV[2],
        }
    """
    res = r.eval(script, 2, "user:1", "user:2", "hello", "world")
    for line in res:
        print(f"  {line}")
    print("\n  💡 numkeys=2 → 前 2 个参数是 KEYS，剩下是 ARGV")
    print("  💡 Cluster 模式必须把 Key 通过 KEYS[] 传，否则路由失败")


if __name__ == "__main__":
    try:
        demo_unsafe_oversell()
        demo_lua_safe()
        demo_rate_limit()
        demo_keys_argv()
        redis.Redis(connection_pool=POOL).delete(
            "stock:item", "rl:user:1001", "user:1", "user:2"
        )
    except redis.ConnectionError as e:
        print(f"❌ Redis 连接失败: {e}")