#!/usr/bin/env python3
"""
transactional_producer.py
=========================
完整的事务 Producer 演示：往两个 Topic 写消息，成功则 commit，失败则 abort。
观察 Read Committed 消费者的行为：commit 之前看不到，commit 之后才出现；abort
的消息永远看不到。

用法：
  pip install confluent-kafka
  python transactional_producer.py commit   # 一个事务，最后 commit
  python transactional_producer.py abort    # 一个事务，写一半 abort
  python transactional_producer.py loop     # 100 个事务循环跑（看吞吐）

  # 同时另一个终端跑 read_committed_consumer.py 看效果

⚠️ 启动两个相同 transactional.id 的实例会触发 ProducerFencedException —— 这是
PID Fencing 的真实效果。
"""

import os
import sys
import time
import json

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

BOOTSTRAP = os.environ.get("KAFKA_BOOTSTRAP", "127.0.0.1:9092")
TX_ID = os.environ.get("TX_ID", "tx-order-1")
T1 = "learn.13.tx_orders"
T2 = "learn.13.tx_outbox"


def ensure_topics():
    a = AdminClient({"bootstrap.servers": BOOTSTRAP})
    existing = a.list_topics(timeout=5).topics
    new_topics = [NewTopic(t, num_partitions=3, replication_factor=1)
                  for t in (T1, T2) if t not in existing]
    if new_topics:
        a.create_topics(new_topics)
        time.sleep(1)


def make_tx_producer():
    p = Producer({
        "bootstrap.servers": BOOTSTRAP,
        "enable.idempotence": True,        # 事务必须开幂等
        "acks": "all",
        "transactional.id": TX_ID,
        "transaction.timeout.ms": 60000,
        "linger.ms": 5,
    })
    p.init_transactions()                  # 一次性初始化，向 Coordinator 申请 PID/Epoch
    return p


def commit_demo():
    ensure_topics()
    p = make_tx_producer()
    p.begin_transaction()
    print(f"== commit demo == tx_id={TX_ID}")
    try:
        for i in range(5):
            payload = json.dumps({"order_id": f"O-{i}", "amount": 99 + i, "ts": time.time()})
            p.produce(T1, key=f"O-{i}", value=payload.encode())
            p.produce(T2, key=f"O-{i}", value=("audit-" + payload).encode())
            print(f"  → produced {i} to {T1} & {T2}  (still pending, read_committed 看不到)")
            time.sleep(0.3)
        p.commit_transaction()
        print("✅ commit_transaction() done. read_committed Consumer 现在能看到这 10 条")
    except KafkaException as e:
        print(f"❌ tx error: {e} → abort")
        p.abort_transaction()


def abort_demo():
    ensure_topics()
    p = make_tx_producer()
    p.begin_transaction()
    print(f"== abort demo == tx_id={TX_ID}")
    try:
        for i in range(3):
            payload = json.dumps({"order_id": f"BAD-{i}"})
            p.produce(T1, key=f"BAD-{i}", value=payload.encode())
            print(f"  → produced BAD-{i}（pending, 即将被 abort）")
        raise RuntimeError("simulated business failure!!!")
    except RuntimeError as e:
        print(f"💥 simulated: {e}")
        p.abort_transaction()
        print("✅ abort_transaction() done. read_committed Consumer 永远看不到这 3 条")


def loop_demo():
    ensure_topics()
    p = make_tx_producer()
    print(f"== loop demo == tx_id={TX_ID}, 100 个事务，每个 100 条")
    t0 = time.time()
    for tx in range(100):
        p.begin_transaction()
        try:
            for i in range(100):
                p.produce(T1, key=f"L-{tx}-{i}", value=f"v{tx}-{i}".encode())
            p.commit_transaction()
        except KafkaException as e:
            err = e.args[0]
            if err.txn_requires_abort():
                p.abort_transaction()
            else:
                raise
        if tx % 10 == 0:
            print(f"  done {tx + 1} txns")
    print(f"== finished in {time.time() - t0:.2f}s ==")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__); sys.exit(0)
    {"commit": commit_demo, "abort": abort_demo, "loop": loop_demo}[sys.argv[1]]()