#!/usr/bin/env python3
"""
大消息分片发送范式（应对案例 7）。

思路：
- 把 10MB 的大消息切成多个 ≤ 512KB 的分片
- 给每片打上同一个 message_id 与 (idx, total)
- 用 message_id 作为 Kafka Key，保证所有分片落同一分区（顺序）
- 消费端按 message_id 收齐后再 reassemble

依赖：pip install confluent-kafka
"""

import os
import json
import uuid
import hashlib
from collections import defaultdict
from confluent_kafka import Producer, Consumer

BOOTSTRAP = os.getenv("BOOTSTRAP", "localhost:9092")
TOPIC = os.getenv("TOPIC", "learn.20.large")
CHUNK_SIZE = 512 * 1024   # 512KB / 片


# =================== Producer 端 ===================
def send_large(producer: Producer, payload: bytes, meta: dict = None):
    """把大 payload 切片发送。"""
    msg_id = uuid.uuid4().hex
    digest = hashlib.md5(payload).hexdigest()
    total = (len(payload) + CHUNK_SIZE - 1) // CHUNK_SIZE

    for idx in range(total):
        chunk = payload[idx * CHUNK_SIZE:(idx + 1) * CHUNK_SIZE]
        envelope = {
            "msg_id": msg_id,
            "idx": idx,
            "total": total,
            "digest": digest,         # MD5 校验
            "size": len(payload),
            "meta": meta or {},
            "data_b64": chunk.hex(),  # 用 hex 简化（生产可用 base64 + 直接 bytes）
        }
        producer.produce(
            topic=TOPIC,
            key=msg_id.encode(),       # 同一 message_id 保证同分区有序
            value=json.dumps(envelope).encode(),
        )
        producer.poll(0)
    producer.flush(10)
    print(f"  ✓ 发送 msg_id={msg_id} 共 {total} 片，原大小 {len(payload)} bytes")


# =================== Consumer 端 ===================
def consume_large(consumer: Consumer):
    """收齐所有分片后 reassemble。"""
    buffers = defaultdict(dict)   # msg_id -> {idx: chunk}
    metas = {}

    print("等待消息 ... Ctrl+C 退出")
    while True:
        msg = consumer.poll(1.0)
        if msg is None or msg.error():
            continue

        env = json.loads(msg.value())
        mid = env["msg_id"]
        buffers[mid][env["idx"]] = bytes.fromhex(env["data_b64"])
        if mid not in metas:
            metas[mid] = env

        # 收齐了？
        if len(buffers[mid]) == env["total"]:
            data = b"".join(buffers[mid][i] for i in sorted(buffers[mid]))
            ok = hashlib.md5(data).hexdigest() == env["digest"]
            print(f"  ✓ 收齐 msg_id={mid} size={len(data)} 校验={'OK' if ok else 'FAIL'}")
            del buffers[mid]
            del metas[mid]
            consumer.commit(msg)
            yield data, env["meta"]


# =================== 演示入口 ===================
def demo():
    p = Producer({
        "bootstrap.servers": BOOTSTRAP,
        "compression.type": "zstd",   # 大消息一定开压缩
        "enable.idempotence": True,
        "acks": "all",
    })

    # 造一个 5MB 的「大消息」
    big = b"X" * (5 * 1024 * 1024)
    print("Sending one 5MB message via chunking ...")
    send_large(p, big, meta={"filename": "huge.bin"})


if __name__ == "__main__":
    if os.getenv("ROLE", "producer") == "consumer":
        c = Consumer({
            "bootstrap.servers": BOOTSTRAP,
            "group.id": "large-msg-demo",
            "auto.offset.reset": "earliest",
            "enable.auto.commit": False,
            "fetch.max.bytes": 10 * 1024 * 1024,
        })
        c.subscribe([TOPIC])
        try:
            for data, meta in consume_large(c):
                pass
        except KeyboardInterrupt:
            c.close()
    else:
        demo()