"""
第 1 章 · Hello Kafka —— 你与 Kafka 的「第一次握手」

完整闭环：
    1. 用 AdminClient 创建 Topic `learn.01.hello`（3 partition × 3 replica）。
       如果 Topic 已经存在则跳过（幂等）。
    2. 用 Producer 同步生产 5 条消息（带 key，便于观察分区路由）。
    3. 用 Consumer 从头消费，打印 partition / offset / key / value。
    4. 程序自动结束（最多等 10 秒）。

前置条件：
    cd learnNote/kafka && docker compose up -d   # 启动集群
    pip install -r requirements.txt              # 安装依赖

运行：
    python 01_intro/code/hello_kafka.py

预期输出（partition 顺序可能不同，offset 必定从 0 起）：
    [INFO ] Topic learn.01.hello 已就绪 (partitions=3, RF=3)
    [SEND ] key=0 → partition=2 offset=0
    [SEND ] key=1 → partition=0 offset=0
    [SEND ] key=2 → partition=1 offset=0
    [SEND ] key=3 → partition=2 offset=1
    [SEND ] key=4 → partition=0 offset=1
    [RECV ] partition=0 offset=0 key=1 value=hello-1
    [RECV ] partition=0 offset=1 key=4 value=hello-4
    ...
    [DONE ] 共消费 5 条消息

阅读建议：
    - 关注「同 key 必定到同一分区」（Murmur2 Hash）
    - 关注 partition 内部 offset 单调递增
    - 关注 group.id + auto.offset.reset=earliest 的组合是怎么"从头消费"的
"""
from __future__ import annotations

import logging
import sys
import time
from typing import Optional

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

# ----------------------- 全局配置 -----------------------
BOOTSTRAP_SERVERS = "localhost:9092"
TOPIC_NAME = "learn.01.hello"
NUM_PARTITIONS = 3
REPLICATION_FACTOR = 3
MESSAGE_COUNT = 5
CONSUMER_GROUP = "learn.01.hello.group"

logging.basicConfig(
    level=logging.INFO,
    format="[%(levelname)-5s] %(message)s",
    stream=sys.stdout,
)
log = logging.getLogger("hello_kafka")


# ----------------------- 1) 建 Topic -----------------------
def ensure_topic(admin: AdminClient, name: str, partitions: int, rf: int) -> None:
    """幂等创建 Topic：已存在则跳过。"""
    existing = admin.list_topics(timeout=10).topics
    if name in existing:
        meta = existing[name]
        log.info(
            "Topic %s 已存在 (partitions=%d, partitions_meta=%s)",
            name,
            len(meta.partitions),
            sorted(meta.partitions.keys()),
        )
        return

    log.info("Topic %s 不存在，正在创建 ...", name)
    new_topic = NewTopic(
        topic=name,
        num_partitions=partitions,
        replication_factor=rf,
        config={
            # 教学用：保留 1 小时即可，方便反复重置
            "retention.ms": str(60 * 60 * 1000),
            # acks=all 时至少 2 副本同步成功才算写入成功
            "min.insync.replicas": "2",
        },
    )
    futures = admin.create_topics([new_topic])

    # create_topics 返回 dict[str, Future]，必须遍历 .result() 才会真正阻塞等结果
    for topic, fut in futures.items():
        try:
            fut.result(timeout=15)
            log.info("Topic %s 已就绪 (partitions=%d, RF=%d)", topic, partitions, rf)
        except KafkaException as e:
            # 并发创建时可能撞到 TOPIC_ALREADY_EXISTS，可视为成功
            if e.args and e.args[0].code() == KafkaError.TOPIC_ALREADY_EXISTS:
                log.info("Topic %s 已存在（并发创建撞车），继续", topic)
            else:
                raise


# ----------------------- 2) 生产消息 -----------------------
def delivery_callback(err: Optional[KafkaError], msg) -> None:
    """每条消息发送结果的异步回调（在 producer.poll/flush 时被触发）。"""
    if err is not None:
        log.error("发送失败: %s", err)
    else:
        log.info(
            "[SEND ] key=%s → partition=%d offset=%d",
            msg.key().decode() if msg.key() else "<none>",
            msg.partition(),
            msg.offset(),
        )


def produce_messages(producer: Producer, topic: str, count: int) -> None:
    log.info("开始生产 %d 条消息到 %s ...", count, topic)
    for i in range(count):
        # key 决定分区：confluent-kafka 默认 Murmur2 Hash 算法
        # 同 key → 同 partition；这里能观察到 partition 路由分布
        producer.produce(
            topic=topic,
            key=str(i).encode(),
            value=f"hello-{i}".encode(),
            on_delivery=delivery_callback,
        )
        # 每次 produce 后调用 poll 处理回调队列；不调用回调就不会触发
        producer.poll(0)

    # flush 阻塞直到所有消息都收到 ack 或 timeout（30s）
    remaining = producer.flush(timeout=30)
    if remaining > 0:
        raise RuntimeError(f"还有 {remaining} 条消息未发送成功，请检查 broker 状态")


# ----------------------- 3) 消费消息 -----------------------
def consume_messages(consumer: Consumer, topic: str, expected: int) -> int:
    """从头消费 expected 条，打印后返回实际拿到的条数。"""
    log.info("订阅 %s，开始消费（最多等待 10s）...", topic)
    consumer.subscribe([topic])

    received = 0
    deadline = time.time() + 10.0

    while received < expected and time.time() < deadline:
        # poll(timeout) 拉取一条消息；timeout 单位秒
        msg = consumer.poll(1.0)
        if msg is None:
            continue
        if msg.error():
            # _PARTITION_EOF 是「分区没新消息」的友好提示，不是错误
            if msg.error().code() == KafkaError._PARTITION_EOF:
                continue
            log.error("消费错误: %s", msg.error())
            continue

        log.info(
            "[RECV ] partition=%d offset=%d key=%s value=%s",
            msg.partition(),
            msg.offset(),
            msg.key().decode() if msg.key() else "<none>",
            msg.value().decode(),
        )
        received += 1

    return received


# ----------------------- main -----------------------
def main() -> int:
    log.info("Hello Kafka! bootstrap.servers = %s", BOOTSTRAP_SERVERS)

    common = {"bootstrap.servers": BOOTSTRAP_SERVERS}

    # ---- AdminClient ----
    admin = AdminClient(common)
    try:
        ensure_topic(admin, TOPIC_NAME, NUM_PARTITIONS, REPLICATION_FACTOR)
    except Exception as e:
        log.error("建 Topic 失败：%s", e)
        log.error("请确认 docker compose up -d 已启动，并能 ping 通 localhost:9092")
        return 1

    # ---- Producer ----
    producer_conf = {
        **common,
        "client.id": "hello-kafka-producer",
        "acks": "all",                  # 等所有 ISR 副本确认（最强一致）
        "enable.idempotence": True,     # 幂等 Producer：失败重试不会重复
        "linger.ms": 5,                 # 攒批 5ms 提升吞吐
        "compression.type": "lz4",      # Producer 端压缩（教学使用 lz4 兼顾速度与压缩比）
        "max.in.flight.requests.per.connection": 5,
    }
    producer = Producer(producer_conf)
    try:
        produce_messages(producer, TOPIC_NAME, MESSAGE_COUNT)
    except Exception as e:
        log.error("生产失败：%s", e)
        return 2

    # ---- Consumer ----
    consumer_conf = {
        **common,
        "group.id": CONSUMER_GROUP,
        "client.id": "hello-kafka-consumer",
        "enable.auto.commit": False,    # 教学环境关掉自动提交，主动控制
        "auto.offset.reset": "earliest",  # 没有 offset 时从最早开始
    }
    consumer = Consumer(consumer_conf)
    try:
        got = consume_messages(consumer, TOPIC_NAME, MESSAGE_COUNT)
        log.info("[DONE ] 共消费 %d 条消息", got)
        if got < MESSAGE_COUNT:
            log.warning("预期 %d 条，实际只拿到 %d 条 —— 检查是否之前 Group 已提交过 offset", MESSAGE_COUNT, got)
            log.warning("可手动重置：kafka-consumer-groups.sh --bootstrap-server %s "
                        "--group %s --reset-offsets --to-earliest --topic %s --execute",
                        BOOTSTRAP_SERVERS, CONSUMER_GROUP, TOPIC_NAME)
    finally:
        consumer.close()

    log.info("👋 完成。下一步可以打开 http://localhost:8080 在 Kafka UI 里查看 Topic 与 Group。")
    return 0


if __name__ == "__main__":
    sys.exit(main())