#!/usr/bin/env python3
"""
seek_by_timestamp.py
====================
按时间戳定位 offset，是线上排障的「时光机」：
> 「线上昨天 14:00~14:05 这一段订单出问题了，能不能只回放这 5 分钟的消息？」

底层原理：Broker 利用 .timeindex 文件（每条 (timestamp, offset) 索引项）
做二分查找，返回 >= 给定 ts 的最早消息 offset。

用法：
  pip install confluent-kafka
  python seek_by_timestamp.py produce 200      # 灌 200 条带 timestamp 的消息
  python seek_by_timestamp.py replay "2026-04-17 12:00:00" "2026-04-17 12:00:05"
                                               # 只回放某 5 秒内的消息
"""

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

from confluent_kafka import Producer, Consumer, TopicPartition

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

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


def produce(n):
    p = Producer({"bootstrap.servers": BOOTSTRAP, "linger.ms": 5})
    base = int(time.time() * 1000)
    for i in range(n):
        body = json.dumps({"id": i, "ts": base + i * 1000})
        # 显式指定 timestamp，让消费者能按时间戳定位
        p.produce(TOPIC, key=str(i % 3), value=body.encode(), timestamp=base + i * 1000)
        if i % 50 == 0:
            p.poll(0)
    p.flush(10)
    log.info(f"produced {n} msgs to {TOPIC} (timestamp 1s/条递增)")


def parse_ts(s):
    """把 'YYYY-MM-DD HH:MM:SS' 解析成毫秒时间戳（本地时区）"""
    dt = datetime.strptime(s, "%Y-%m-%d %H:%M:%S").astimezone()
    return int(dt.timestamp() * 1000)


def replay(start_str, end_str):
    start_ts = parse_ts(start_str)
    end_ts = parse_ts(end_str)
    log.info(f"replay window: {start_str} ({start_ts}) ~ {end_str} ({end_ts})")

    c = Consumer({
        "bootstrap.servers": BOOTSTRAP,
        "group.id": "demo-replay-by-ts",
        "enable.auto.commit": False,
        "auto.offset.reset": "earliest",
    })
    c.subscribe([TOPIC])

    # 必须先 poll 一次（或者 assign），否则 seek 报「No current assignment」
    c.poll(2.0)

    md = c.list_topics(TOPIC, timeout=5.0).topics[TOPIC]
    tps_query = [TopicPartition(TOPIC, p, start_ts) for p in md.partitions]
    log.info(f"querying offsets_for_times for {len(tps_query)} partitions...")
    offsets = c.offsets_for_times(tps_query, timeout=5.0)

    # offsets 列表里每个 tp.offset 已经是「>= start_ts 的最早 offset」
    for tp in offsets:
        if tp.offset == -1:
            log.info(f"  P{tp.partition}: 无消息 >= start_ts (整个分区都比该时间早或没有数据)")
            continue
        log.info(f"  P{tp.partition}: seek to offset {tp.offset}")
        c.seek(tp)

    received = 0
    try:
        while True:
            msg = c.poll(1.0)
            if msg is None:
                # 等待 3s 没有新消息就结束
                continue
            if msg.error():
                log.error(msg.error()); continue
            ts_type, ts = msg.timestamp()
            if ts > end_ts:
                log.info(f"  P{msg.partition()}@{msg.offset()} ts={ts} > end_ts → 该分区窗口结束")
                # 简单起见全部退出（生产里应跟踪每个分区分别结束）
                break
            received += 1
            payload = json.loads(msg.value())
            log.info(f"  REPLAY p{msg.partition()}@{msg.offset()} ts={ts} id={payload['id']}")
            if received >= 1000:
                break
    finally:
        c.close()
    log.info(f"window replayed: {received} msgs")


def main():
    if len(sys.argv) < 2:
        print(__doc__); return
    cmd = sys.argv[1]
    if cmd == "produce":
        n = int(sys.argv[2]) if len(sys.argv) > 2 else 200
        produce(n)
    elif cmd == "replay":
        start = sys.argv[2]; end = sys.argv[3]
        replay(start, end)
    else:
        print("unknown:", cmd); print(__doc__)


if __name__ == "__main__":
    main()