#!/usr/bin/env python3
"""
manual_commit_correct.py
========================
手动提交三种正确姿势：
  1) commit_sync_each_msg : 每条消息处理完同步提交（最稳，吞吐最低）
  2) commit_async_each_msg: 每条消息异步提交（高吞吐，关闭时同步兜底）
  3) commit_batch         : 批量处理 + 批量提交（生产推荐）

用法：
  pip install confluent-kafka
  # 先用 auto_commit_pitfall.py 的 produce 命令灌些数据
  python manual_commit_correct.py sync
  python manual_commit_correct.py async
  python manual_commit_correct.py batch

每次跑都用同一个 group（demo-manual），可以观察连续运行
不会丢/重——相比 auto 模式有显著区别。

关键 API：
  - consumer.commit(message=msg, asynchronous=False)  # 同步、提交单条
  - consumer.commit(asynchronous=True)                # 异步、提交当前 highwater
  - consumer.commit(offsets=[TopicPartition(t, p, off)])  # 精细化
"""

import os
import sys
import time
import json
import logging

from confluent_kafka import Consumer, TopicPartition

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

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


def make_consumer(group):
    return Consumer({
        "bootstrap.servers": BOOTSTRAP,
        "group.id": group,
        "enable.auto.commit": False,           # 关键：关闭自动提交
        "auto.offset.reset": "earliest",
        "session.timeout.ms": 10000,
        "max.poll.interval.ms": 60000,
    })


def commit_sync_each_msg():
    c = make_consumer("demo-manual-sync")
    c.subscribe([TOPIC])
    try:
        while True:
            msg = c.poll(1.0)
            if msg is None: continue
            if msg.error():
                log.error(msg.error()); continue
            payload = json.loads(msg.value())
            try:
                process(payload)
            except Exception:
                log.exception("process failed, will retry next round")
                continue
            c.commit(message=msg, asynchronous=False)   # 同步提交，确认成功
            log.info(f"  committed p{msg.partition()}@{msg.offset() + 1}")
    finally:
        c.close()


def commit_async_each_msg():
    c = make_consumer("demo-manual-async")
    c.subscribe([TOPIC])

    def on_commit(err, partitions):
        if err:
            log.error(f"async commit failed: {err}")
        else:
            for tp in partitions:
                log.debug(f"async commit ok p{tp.partition}@{tp.offset}")

    try:
        while True:
            msg = c.poll(1.0)
            if msg is None: continue
            if msg.error():
                log.error(msg.error()); continue
            payload = json.loads(msg.value())
            try:
                process(payload)
            except Exception:
                log.exception("process failed, skip commit")
                continue
            c.commit(message=msg, asynchronous=True)
            log.info(f"  asynchronously committed p{msg.partition()}@{msg.offset() + 1}")
    except KeyboardInterrupt:
        pass
    finally:
        # 关闭前同步提交一次「兜底」
        try:
            c.commit(asynchronous=False)
            log.info("final sync commit done")
        finally:
            c.close()


def commit_batch():
    c = make_consumer("demo-manual-batch")
    c.subscribe([TOPIC])
    BATCH = 50

    try:
        while True:
            buf = []
            t0 = time.time()
            while len(buf) < BATCH and time.time() - t0 < 1.0:
                msg = c.poll(0.5)
                if msg is None: continue
                if msg.error():
                    log.error(msg.error()); continue
                buf.append(msg)

            if not buf:
                continue

            # 业务批处理：通常是一次性 INSERT INTO ... VALUES (?,?), ...
            try:
                process_batch([json.loads(m.value()) for m in buf])
            except Exception:
                log.exception("batch failed, will retry")
                continue

            # 取每个分区的最大 offset，统一提交（避免漏提交）
            tp_max = {}
            for m in buf:
                key = (m.topic(), m.partition())
                if m.offset() > tp_max.get(key, -1):
                    tp_max[key] = m.offset()
            offsets = [TopicPartition(t, p, off + 1) for (t, p), off in tp_max.items()]
            c.commit(offsets=offsets, asynchronous=False)
            log.info(f"  committed batch: {[(o.partition, o.offset) for o in offsets]}")
    finally:
        c.close()


def process(payload):
    # 业务！务必幂等
    time.sleep(0.05)


def process_batch(payloads):
    time.sleep(0.05)
    log.info(f"  process_batch size={len(payloads)}")


def main():
    if len(sys.argv) < 2:
        print(__doc__); return
    cmd = sys.argv[1]
    {"sync": commit_sync_each_msg,
     "async": commit_async_each_msg,
     "batch": commit_batch}[cmd]()


if __name__ == "__main__":
    main()