#!/usr/bin/env python3
"""
auto_commit_pitfall.py
======================
演示 enable.auto.commit=true 的两种陷阱：
  1) 处理慢 -> 已提交 offset 跑到「真处理位点」前 -> 进程崩溃后丢消息
  2) 处理快 -> 未到 commit 周期就崩溃 -> 重启后已处理消息会被重复

用法（本地需要 Kafka 在 127.0.0.1:9092）：
  pip install confluent-kafka

  # 终端 A —— 先生产 50 条 0..49
  python auto_commit_pitfall.py produce

  # 终端 B —— 演示「丢」（处理 8s，commit interval 5s）
  python auto_commit_pitfall.py lose

  # 终端 C —— 演示「重」（处理 0.05s，commit interval 5s）
  python auto_commit_pitfall.py dup

每个 demo 模式都会在「关键时刻」抛 SystemExit 模拟进程崩溃，
然后请重新跑一次 lose/dup（同一 group），观察 CURRENT-OFFSET 与
真实处理日志的差距：
  kafka-consumer-groups.sh --bootstrap-server 127.0.0.1:9092 \\
      --describe --group demo-pitfall-lose
"""

import os
import sys
import time
import json
import logging
from datetime import datetime

from confluent_kafka import Producer, Consumer
from confluent_kafka.admin import AdminClient, NewTopic

BOOTSTRAP = os.environ.get("KAFKA_BOOTSTRAP", "127.0.0.1:9092")
TOPIC = "learn.12.pitfall"
N = 50

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("pitfall")


def ensure_topic():
    admin = AdminClient({"bootstrap.servers": BOOTSTRAP})
    md = admin.list_topics(timeout=5).topics
    if TOPIC not in md:
        log.info(f"creating topic {TOPIC} (3 partitions)")
        fs = admin.create_topics([NewTopic(TOPIC, num_partitions=3, replication_factor=1)])
        for t, f in fs.items():
            try: f.result()
            except Exception as e: log.warning(f"create topic: {e}")


def produce():
    ensure_topic()
    p = Producer({"bootstrap.servers": BOOTSTRAP, "linger.ms": 5})
    for i in range(N):
        body = json.dumps({"id": i, "ts": datetime.utcnow().isoformat()})
        p.produce(TOPIC, key=str(i % 3), value=body.encode())
    p.flush(10)
    log.info(f"produced {N} msgs to {TOPIC}")


def consume(group, process_seconds, commit_interval_ms, crash_after_seconds):
    ensure_topic()
    c = Consumer({
        "bootstrap.servers": BOOTSTRAP,
        "group.id": group,
        "enable.auto.commit": True,
        "auto.commit.interval.ms": commit_interval_ms,
        "auto.offset.reset": "earliest",
        # 不打开手动 fetch.min/max 等，保持默认让效果更明显
    })
    c.subscribe([TOPIC])

    started = time.time()
    handled_ids = []
    log.info(
        f"consumer group={group} process={process_seconds}s commit_int={commit_interval_ms}ms "
        f"crash_after={crash_after_seconds}s"
    )
    try:
        while True:
            msg = c.poll(timeout=1.0)
            now = time.time()
            if now - started >= crash_after_seconds:
                # 模拟进程被 kill -9：直接 os._exit，不让 client flush/commit
                log.warning("💥 SIMULATED CRASH (os._exit) — last commit may NOT be the truth!")
                os._exit(137)
            if msg is None:
                continue
            if msg.error():
                log.error(msg.error())
                continue
            payload = json.loads(msg.value())
            log.info(
                f"  poll p{msg.partition()}@offset={msg.offset()} id={payload['id']} "
                f"-> processing for {process_seconds}s"
            )
            time.sleep(process_seconds)        # 模拟业务
            handled_ids.append(payload["id"])
            log.info(f"  done id={payload['id']}, total handled={len(handled_ids)}")
    finally:
        c.close()


def main():
    if len(sys.argv) < 2:
        print(__doc__); return
    cmd = sys.argv[1]
    if cmd == "produce":
        produce()
    elif cmd == "lose":
        # 处理 8s + commit interval 5s + 第 6 秒崩溃
        # → poll 第一条 -> commit 立即提交 offset -> 处理 8s 还没完 -> 6s 时崩
        # → 重启从「已提交+1」继续，第 0 条永远没处理（实际上更糟：第 1~K 条全部丢）
        consume(
            group="demo-pitfall-lose",
            process_seconds=8.0,
            commit_interval_ms=5000,
            crash_after_seconds=6.0,
        )
    elif cmd == "dup":
        # 处理 0.05s + commit interval 5s + 第 3 秒崩溃
        # → 3s 内已处理几十条，但 commit 还没触发 -> 重启后从 0 重新拉
        consume(
            group="demo-pitfall-dup",
            process_seconds=0.05,
            commit_interval_ms=5000,
            crash_after_seconds=3.0,
        )
    else:
        print("unknown:", cmd); print(__doc__)


if __name__ == "__main__":
    main()