"""
第 5 章 - seek / seek_to_beginning / offsets_for_times 演示
==================================================

3 种典型用法：
  1. seek_to_beginning：在 on_assign 回调里强制从头消费
  2. seek 到指定 offset
  3. offsets_for_times：按时间戳跳转

注意：seek 必须在 partition 已分配后才能调用，所以一般写在 on_assign 回调里。

运行：
    bash ../init.sh
    python seek_demo.py beginning          # 从头
    python seek_demo.py offset 5           # 跳到 offset 5
    python seek_demo.py time 60            # 跳到 60 秒前
"""

from __future__ import annotations

import sys
import time

from confluent_kafka import Consumer, KafkaError, TopicPartition

BOOTSTRAP = "127.0.0.1:9092"
TOPIC = "learn.05.seek_demo"
GROUP = f"ch5-seek-{int(time.time())}"

mode = sys.argv[1] if len(sys.argv) > 1 else "beginning"
arg = sys.argv[2] if len(sys.argv) > 2 else "0"


def make_on_assign(mode: str, arg: str):
    """返回 on_assign 回调，根据 mode 决定 seek 策略。"""

    def on_assign(consumer, partitions):
        print(f"[on_assign] mode={mode}, partitions={[(p.topic, p.partition) for p in partitions]}")

        if mode == "beginning":
            for p in partitions:
                p.offset = 0
            consumer.assign(partitions)

        elif mode == "offset":
            target = int(arg)
            for p in partitions:
                p.offset = target
            consumer.assign(partitions)
            print(f"  -> seek to offset {target} on every partition")

        elif mode == "time":
            seconds_ago = int(arg)
            ts = int((time.time() - seconds_ago) * 1000)
            tps_with_ts = [TopicPartition(p.topic, p.partition, ts) for p in partitions]
            offsets = consumer.offsets_for_times(tps_with_ts, timeout=10)
            for o in offsets:
                print(f"  P{o.partition} @ ts={ts} -> offset={o.offset}")
            consumer.assign(offsets)

        else:
            consumer.assign(partitions)
            print("  no seek, using committed/auto.offset.reset")

    return on_assign


def main() -> None:
    consumer = Consumer(
        {
            "bootstrap.servers": BOOTSTRAP,
            "group.id": GROUP,
            "auto.offset.reset": "latest",
            "enable.auto.commit": False,
        }
    )
    consumer.subscribe([TOPIC], on_assign=make_on_assign(mode, arg))

    deadline = time.time() + 8
    consumed = 0
    while time.time() < deadline:
        msg = consumer.poll(1.0)
        if msg is None:
            continue
        if msg.error():
            if msg.error().code() != KafkaError._PARTITION_EOF:
                print(f"err: {msg.error()}")
            continue
        consumed += 1
        print(f"[{consumed:>3}] P{msg.partition()}@{msg.offset()} value={msg.value().decode()}")

    print(f"\nconsumed {consumed} msgs in 8s")
    consumer.close()


if __name__ == "__main__":
    main()