#!/usr/bin/env python3
"""
observe_isr.py — 实时打印每个分区的 Leader / ISR / Replicas

用法：
    pip install confluent-kafka
    python observe_isr.py --bootstrap localhost:9092 --topic learn.09.orders
    python observe_isr.py --bootstrap localhost:9092 --topic learn.09.orders --interval 2

观察重点：
    1) 在另一个终端 `docker compose stop kafka1`，看本脚本输出里 Leader/ISR 的变化；
    2) 起回来后看 ISR 重新扩张；
    3) 配合 09_replication.md 中「故障场景剧本」对照阅读。

依赖：
    confluent-kafka >= 2.0（AdminClient.describe_topics 在 2.2 之后行为更直观；
    若只有旧版本，可改用 list_topics(...) 的 cluster_metadata，本脚本两种 API 都兼容）。
"""

from __future__ import annotations

import argparse
import sys
import time
from datetime import datetime

from confluent_kafka.admin import AdminClient


COLOR_RED = "\033[31m"
COLOR_YEL = "\033[33m"
COLOR_GRN = "\033[32m"
COLOR_DIM = "\033[2m"
COLOR_RST = "\033[0m"


def parse_args():
    p = argparse.ArgumentParser(description="实时观察 Kafka Topic 的 ISR / Leader 变化")
    p.add_argument("--bootstrap", default="localhost:9092",
                   help="bootstrap.servers")
    p.add_argument("--topic", required=True, help="要观察的 Topic 名")
    p.add_argument("--interval", type=float, default=1.0,
                   help="刷新间隔秒数（默认 1.0）")
    p.add_argument("--once", action="store_true",
                   help="只查一次后退出（用于脚本编排）")
    return p.parse_args()


def fetch_topic_meta(admin: AdminClient, topic: str):
    """优先使用 list_topics()（AdminClient 与 Producer/Consumer 都通用）。"""
    md = admin.list_topics(topic=topic, timeout=10)
    if topic not in md.topics:
        raise RuntimeError(f"Topic {topic!r} 不存在")
    t = md.topics[topic]
    if t.error is not None:
        raise RuntimeError(f"获取 Topic {topic} 元数据失败：{t.error}")
    parts = []
    for pid in sorted(t.partitions.keys()):
        p = t.partitions[pid]
        parts.append({
            "id": pid,
            "leader": p.leader,
            "replicas": list(p.replicas),
            "isr": list(p.isrs),
        })
    cluster = {
        "broker_ids": sorted(md.brokers.keys()),
        "controller_id": md.controller_id,
        "cluster_id": md.cluster_id,
    }
    return cluster, parts


def fmt_replicas(replicas, isr, leader):
    """把副本列表用颜色标注：Leader=黄、ISR 内=绿、OSR=红。"""
    chips = []
    for r in replicas:
        if r == leader:
            chips.append(f"{COLOR_YEL}{r}*{COLOR_RST}")
        elif r in isr:
            chips.append(f"{COLOR_GRN}{r}{COLOR_RST}")
        else:
            chips.append(f"{COLOR_RED}{r}{COLOR_RST}")
    return ",".join(chips)


def render(cluster, parts, topic):
    sys.stdout.write("\033[2J\033[H")  # 清屏
    print(f"=== Kafka ISR Watcher === {datetime.now().strftime('%H:%M:%S')}")
    print(f"cluster_id     : {cluster['cluster_id']}")
    print(f"controller     : {cluster['controller_id']}")
    print(f"alive brokers  : {cluster['broker_ids']}")
    print(f"topic          : {topic}")
    print()
    print(f"{'Part':<6}{'Leader':<8}{'Replicas (黄=Leader, 绿=ISR, 红=OSR)':<55}{'ISR 数'}")
    print("-" * 85)
    for p in parts:
        leader = p["leader"]
        leader_str = f"{COLOR_YEL}{leader}{COLOR_RST}" if leader >= 0 else f"{COLOR_RED}NONE{COLOR_RST}"
        rep_str = fmt_replicas(p["replicas"], p["isr"], leader)
        n_isr = len(p["isr"])
        n_rep = len(p["replicas"])
        warn = ""
        if n_isr < n_rep:
            warn = f"  {COLOR_YEL}⚠ under-replicated{COLOR_RST}"
        if leader < 0:
            warn = f"  {COLOR_RED}⚠ NO LEADER{COLOR_RST}"
        print(f"{p['id']:<6}{leader_str:<16}{rep_str:<70}{n_isr}/{n_rep}{warn}")
    print()
    print(f"{COLOR_DIM}按 Ctrl+C 退出{COLOR_RST}")


def main():
    args = parse_args()
    admin = AdminClient({"bootstrap.servers": args.bootstrap})

    try:
        while True:
            try:
                cluster, parts = fetch_topic_meta(admin, args.topic)
                render(cluster, parts, args.topic)
            except Exception as e:
                sys.stdout.write("\033[2J\033[H")
                print(f"{COLOR_RED}[ERROR] {e}{COLOR_RST}")
                print(f"{COLOR_DIM}下一次刷新：{args.interval}s 后重试…{COLOR_RST}")
            if args.once:
                break
            time.sleep(args.interval)
    except KeyboardInterrupt:
        print("\nbye.")


if __name__ == "__main__":
    main()