"""
Ch4 配套代码 3 / 3 —— 手写一个简易 bigkeys 扫描器

  - 用 SCAN 渐进式遍历所有 Key
  - 对每种类型分别记录最大的 N 个
  - 输出汇总报告，定位需要治理的 Key
"""

import heapq
import redis

r = redis.Redis(host="127.0.0.1", port=6379, decode_responses=True)
TOP_N = 5
SCAN_COUNT = 200


def section(t): print("\n" + "=" * 60 + f"\n{t}\n" + "=" * 60)


def gen_demo_data() -> None:
    """造一些大小不一的 Key，方便看效果。"""
    section("[setup] 造测试数据：String/List/Hash/Set/ZSet 各 5 个，大小不一")
    for i in range(5):
        r.set(f"demo:str:{i}", "x" * (10 ** (i + 1)))
        r.delete(f"demo:list:{i}")
        if i > 0: r.rpush(f"demo:list:{i}", *[f"v{j}" for j in range(10 * (i + 1))])
        r.delete(f"demo:hash:{i}")
        if i > 0: r.hset(f"demo:hash:{i}", mapping={f"f{j}": f"v{j}" for j in range(20 * (i + 1))})
        r.delete(f"demo:set:{i}")
        if i > 0: r.sadd(f"demo:set:{i}", *[f"m{j}" for j in range(15 * (i + 1))])
        r.delete(f"demo:zset:{i}")
        if i > 0: r.zadd(f"demo:zset:{i}", {f"m{j}": j for j in range(25 * (i + 1))})


def cleanup() -> None:
    cursor = 0
    while True:
        cursor, keys = r.scan(cursor, match="demo:*", count=200)
        if keys: r.delete(*keys)
        if cursor == 0: break


def find_big_keys() -> None:
    section("[scan] 扫描 demo:* 找出每种类型 Top 5 大 Key")

    # 每种类型一个最小堆，存 (size, key)
    heaps = {t: [] for t in ["string", "list", "hash", "set", "zset"]}
    type_handlers = {
        "string": lambda k: r.strlen(k),
        "list":   lambda k: r.llen(k),
        "hash":   lambda k: r.hlen(k),
        "set":    lambda k: r.scard(k),
        "zset":   lambda k: r.zcard(k),
    }

    cursor = 0
    total_scanned = 0
    while True:
        cursor, keys = r.scan(cursor, match="demo:*", count=SCAN_COUNT)
        for k in keys:
            t = r.type(k)
            if t not in type_handlers: continue
            try:
                size = type_handlers[t](k)
                heap = heaps[t]
                if len(heap) < TOP_N:
                    heapq.heappush(heap, (size, k))
                elif size > heap[0][0]:
                    heapq.heapreplace(heap, (size, k))
                total_scanned += 1
            except redis.RedisError:
                continue
        if cursor == 0: break

    print(f"  共扫描 {total_scanned} 个 Key\n")
    for t, heap in heaps.items():
        if not heap: continue
        print(f"  📦 {t} 类型 Top {len(heap)}：")
        for size, k in sorted(heap, reverse=True):
            mem = r.memory_usage(k) or 0
            unit = "字节" if t == "string" else "元素"
            print(f"      {k:<25} {size:>6} {unit}    内存约 {mem} 字节")
        print()


if __name__ == "__main__":
    try:
        cleanup()
        gen_demo_data()
        find_big_keys()
    finally:
        cleanup()