#!/usr/bin/env python3
"""
compaction_demo.py
==================
往 cleanup.policy=compact 的 topic 写入大量「同 Key 不同 Value」，
等几秒让 Cleaner 触发，然后从头消费看实际剩了多少条。

结论：物理记录数 << 写入次数（因为相同 Key 的旧值被压缩掉），但每个 Key
的最新 Value 仍然能读到。

用法：
  pip install confluent-kafka
  bash ../init.sh                         # 先建好 compact topic
  python compaction_demo.py write 10 100  # 10 个 Key，每个写 100 次
  python compaction_demo.py wait 30       # 等 30 秒让 Cleaner 工作
  python compaction_demo.py read          # 从头消费，看剩多少条
"""

import os
import sys
import time
import json
from collections import Counter

from confluent_kafka import Producer, Consumer

BOOTSTRAP = os.environ.get("KAFKA_BOOTSTRAP", "127.0.0.1:9092")
TOPIC = os.environ.get("TOPIC", "learn.14.compact")


def write(num_keys, repeat):
    p = Producer({"bootstrap.servers": BOOTSTRAP, "linger.ms": 5})
    total = num_keys * repeat
    print(f"writing {num_keys} keys × {repeat} updates = {total} records")
    t0 = time.time()
    for r in range(repeat):
        for k in range(num_keys):
            key = f"item-{k:03d}".encode()
            val = json.dumps({"version": r, "ts": time.time()}).encode()
            p.produce(TOPIC, key=key, value=val)
        if r % 10 == 0:
            p.poll(0)
    p.flush(30)
    print(f"done in {time.time()-t0:.1f}s; latest version per key = {repeat - 1}")


def wait_cleaner(secs):
    print(f"sleeping {secs}s to let Log Cleaner do its work …")
    print("（init.sh 把 segment.ms=5s, dirty ratio=0.1, max.lag=60s，"
          "所以一般 30s 内就能看到 compaction 效果）")
    for i in range(secs, 0, -1):
        sys.stdout.write(f"\r  remaining: {i:3d}s ")
        sys.stdout.flush()
        time.sleep(1)
    sys.stdout.write("\n")


def read():
    c = Consumer({
        "bootstrap.servers": BOOTSTRAP,
        "group.id": "compaction-reader-" + str(os.getpid()),
        "auto.offset.reset": "earliest",
        "enable.auto.commit": False,
    })
    c.subscribe([TOPIC])

    n = 0
    cnt = Counter()
    latest = {}
    end_t = time.time() + 5
    while time.time() < end_t:
        msg = c.poll(1.0)
        if msg is None: continue
        if msg.error(): continue
        n += 1
        end_t = time.time() + 2   # 每读到一条就续 2 秒
        k = msg.key().decode() if msg.key() else "<null>"
        cnt[k] += 1
        try:
            latest[k] = json.loads(msg.value())["version"] if msg.value() else None
        except Exception:
            pass
    c.close()

    print(f"\n=== 物理读到 {n} 条记录 ===")
    print(f"涉及 {len(cnt)} 个 Key")
    print()
    if cnt:
        avg = sum(cnt.values()) / len(cnt)
        print(f"平均每个 Key 还剩 {avg:.2f} 条记录")
        print(f"  ↑ 没压缩前每个 Key 应该有 N 条；压缩后理论 1 条；中间值表示压缩进行中")
        print()
        print("Key 现在的「最新版本号」:")
        for k in sorted(latest.keys())[:10]:
            print(f"  {k}: version={latest[k]} (occurrences={cnt[k]})")
        if len(latest) > 10:
            print(f"  ... ({len(latest)} keys total)")


def main():
    if len(sys.argv) < 2:
        print(__doc__); return
    cmd = sys.argv[1]
    if cmd == "write":
        num = int(sys.argv[2]) if len(sys.argv) > 2 else 10
        rep = int(sys.argv[3]) if len(sys.argv) > 3 else 100
        write(num, rep)
    elif cmd == "wait":
        secs = int(sys.argv[2]) if len(sys.argv) > 2 else 30
        wait_cleaner(secs)
    elif cmd == "read":
        read()
    else:
        print(__doc__)


if __name__ == "__main__":
    main()