#!/usr/bin/env python3
"""
Consumer Lag 监控 + Webhook 告警。

依赖：
    pip install confluent-kafka requests

示例：
    BOOTSTRAP=localhost:9092 LAG_THRESHOLD=10000 \
    WEBHOOK_URL=https://hooks.example.com/xxx python lag_alert.py
"""

import os
import time
import requests
from confluent_kafka.admin import AdminClient, ConsumerGroupTopicPartitions
from confluent_kafka import Consumer, TopicPartition

BOOTSTRAP = os.getenv("BOOTSTRAP", "localhost:9092")
LAG_THRESHOLD = int(os.getenv("LAG_THRESHOLD", "10000"))
INTERVAL = int(os.getenv("INTERVAL", "30"))
WEBHOOK = os.getenv("WEBHOOK_URL", "")


def list_groups(admin):
    fut = admin.list_consumer_groups(request_timeout=10)
    res = fut.result()
    return [g.group_id for g in res.valid if not g.group_id.startswith("_")]


def get_group_lag(group_id):
    admin = AdminClient({"bootstrap.servers": BOOTSTRAP})
    consumer = Consumer({
        "bootstrap.servers": BOOTSTRAP,
        "group.id": group_id,
        "enable.auto.commit": False,
    })
    try:
        fut = admin.list_consumer_group_offsets(
            [ConsumerGroupTopicPartitions(group_id, None)]
        )
        offsets_res = list(fut.values())[0].result()
        committed = offsets_res.topic_partitions

        rows = []
        for tp in committed:
            if tp.error:
                continue
            tp_q = TopicPartition(tp.topic, tp.partition)
            low, high = consumer.get_watermark_offsets(tp_q, timeout=5, cached=False)
            cur = tp.offset if tp.offset >= 0 else low
            lag = max(0, high - cur)
            rows.append((tp.topic, tp.partition, cur, high, lag))
        return rows
    finally:
        consumer.close()


def alert(group, summary, details):
    text = "[Kafka Lag Alert] group={}\n{}\n\n{}".format(group, summary, details)
    print(text)
    if WEBHOOK:
        try:
            requests.post(
                WEBHOOK,
                json={"msgtype": "text", "text": {"content": text}},
                timeout=5,
            )
        except Exception as e:
            print("  webhook fail:", e)


def main():
    print("Lag Monitor: bootstrap={} threshold={} interval={}s".format(
        BOOTSTRAP, LAG_THRESHOLD, INTERVAL))
    admin = AdminClient({"bootstrap.servers": BOOTSTRAP})
    while True:
        try:
            groups = list_groups(admin)
            print("\n[{}] checking {} groups".format(
                time.strftime("%H:%M:%S"), len(groups)))
            for g in groups:
                rows = get_group_lag(g)
                if not rows:
                    continue
                total_lag = sum(r[4] for r in rows)
                worst = max(rows, key=lambda r: r[4])
                print("  {:30s} total={:>10d} worst={}-{} lag={}".format(
                    g, total_lag, worst[0], worst[1], worst[4]))
                if worst[4] >= LAG_THRESHOLD:
                    top = sorted(rows, key=lambda r: -r[4])[:5]
                    detail = "\n".join(
                        "  {}-{}: cur={} end={} lag={}".format(t, p, c, e, l)
                        for (t, p, c, e, l) in top
                    )
                    alert(
                        group=g,
                        summary="max-lag={} total={}".format(worst[4], total_lag),
                        details="Top 5 lagged partitions:\n" + detail,
                    )
        except Exception as e:
            print("  ERR:", e)
        time.sleep(INTERVAL)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nbye")