#!/usr/bin/env python3
"""Collect Debezium, Kafka CDC, and Spark metrics into postgres monitor schema."""
import json
import shlex
import subprocess
import urllib.request
from collections import Counter
from datetime import datetime, timezone

import psycopg2

PG_DSN = "host=10.0.21.51 dbname=postgres user=mo password=Dell2026!"
KAFKA = "10.0.21.36:9092"
DEBEZIUM = "http://localhost:8083"  # Kafka Connect on kafka01; fallback lake01 :8083
SPARK_MASTER = "http://10.0.21.50:8080"

CDC_TOPICS = [
    ("PostgreSQL", "postgres-sales.public.sales_orders"),
    ("MongoDB", "mongodb-supplychain.supplychain.events"),
]

OP_LABELS = {"c": "INSERT", "u": "UPDATE", "d": "DELETE", "r": "SNAPSHOT", "i": "INSERT"}


def fetch_json(url, timeout=10):
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read().decode())


def collect_debezium(cur):
    connectors = fetch_json(f"{DEBEZIUM}/connectors")
    cur.execute("DELETE FROM monitor.debezium_connectors")
    now = datetime.now(timezone.utc)
    for name in connectors:
        try:
            st = fetch_json(f"{DEBEZIUM}/connectors/{name}/status")
        except Exception as e:
            cur.execute(
                """INSERT INTO monitor.debezium_connectors
                   (connector_name, state, task_state, worker_id, checked_at)
                   VALUES (%s,%s,%s,%s,%s)""",
                (name, "ERROR", str(e)[:32], "", now),
            )
            continue
        conn_state = st.get("connector", {}).get("state", "UNKNOWN")
        tasks = st.get("tasks") or []
        task_state = tasks[0].get("state", "NONE") if tasks else "NONE"
        worker = st.get("connector", {}).get("worker_id", "")
        cur.execute(
            """INSERT INTO monitor.debezium_connectors
               (connector_name, state, task_state, worker_id, checked_at)
               VALUES (%s,%s,%s,%s,%s)""",
            (name, conn_state, task_state, worker, now),
        )


KAFKA_BIN = "/opt/kafka/bin"
USE_SSH_KAFKA = False  # set True when running off-host


def _kafka_cmd(bin_name, args):
    parts = [f"{KAFKA_BIN}/{bin_name}"] + list(args)
    if USE_SSH_KAFKA:
        remote = " ".join(shlex.quote(p) for p in parts)
        full = f"ssh -o StrictHostKeyChecking=no root@10.0.21.36 {remote}"
        return subprocess.check_output(full, shell=True, stderr=subprocess.DEVNULL, timeout=90, text=True)
    return subprocess.check_output(parts, stderr=subprocess.DEVNULL, timeout=90, text=True)


def kafka_end_offsets(topic):
    try:
        out = _kafka_cmd(
            "kafka-run-class.sh",
            [
                "kafka.tools.GetOffsetShell",
                "--broker-list",
                "localhost:9092",
                "--topic",
                topic,
            ],
        )
    except Exception:
        return []
    rows = []
    for line in out.strip().splitlines():
        parts = line.split(":")
        if len(parts) >= 3:
            rows.append((int(parts[1]), int(parts[2])))
    return rows


def sample_topic_messages(topic, max_msgs=3000, tail=5000):
    """Sample recent messages using kafka-console-consumer from tail."""
    offsets = kafka_end_offsets(topic)
    if not offsets:
        return []
    # Pick partition 0 for sampling
    part, end = offsets[0]
    start = max(0, end - tail)
    try:
        out = _kafka_cmd(
            "kafka-console-consumer.sh",
            [
                "--bootstrap-server",
                "localhost:9092",
                "--topic",
                topic,
                "--partition",
                str(part),
                "--offset",
                str(start),
                "--max-messages",
                str(min(max_msgs, tail)),
                "--timeout-ms",
                "15000",
            ],
        )
    except Exception:
        return []
    return [ln for ln in out.strip().split("\n") if ln.strip()]


def parse_debezium_line(line):
    try:
        doc = json.loads(line)
        payload = doc.get("payload") or doc
        op = payload.get("op") or payload.get("operationType") or "?"
        src = payload.get("source") or {}
        table = src.get("table") or src.get("collection") or ""
        ts_ms = payload.get("ts_ms") or src.get("ts_ms")
        after = payload.get("after") or {}
        before = payload.get("before") or {}
        row = after if after else before
        key = str(row.get("order_id") or row.get("event_id") or row.get("_id") or "")[:200]
        detail = str(row.get("region") or row.get("type") or row.get("department") or "")[:200]
        event_ts = None
        if ts_ms:
            event_ts = datetime.fromtimestamp(int(ts_ms) / 1000, tz=timezone.utc)
        return op, table, key, detail, event_ts
    except Exception:
        return None


def collect_kafka_cdc(cur):
    now = datetime.now(timezone.utc)
    cur.execute("DELETE FROM monitor.kafka_topics")
    cur.execute("DELETE FROM monitor.cdc_operations")
    cur.execute("DELETE FROM monitor.cdc_recent_events")

    for source, topic in CDC_TOPICS:
        for part, end in kafka_end_offsets(topic):
            cur.execute(
                """INSERT INTO monitor.kafka_topics (topic, partition_id, end_offset, checked_at)
                   VALUES (%s,%s,%s,%s)""",
                (topic, part, end, now),
            )

        lines = sample_topic_messages(topic, max_msgs=2000, tail=3000)
        ops = Counter()
        recent = []
        for line in lines:
            parsed = parse_debezium_line(line)
            if not parsed:
                continue
            op, table, key, detail, event_ts = parsed
            ops[op] += 1
            if len(recent) < 100:
                recent.append((op, table, key, detail, event_ts))

        for op, cnt in ops.items():
            cur.execute(
                """INSERT INTO monitor.cdc_operations
                   (source_system, topic, operation, operation_label, event_count, checked_at)
                   VALUES (%s,%s,%s,%s,%s,%s)""",
                (source, topic, op, OP_LABELS.get(op, op), cnt, now),
            )

        for op, table, key, detail, event_ts in recent[:50]:
            cur.execute(
                """INSERT INTO monitor.cdc_recent_events
                   (source_system, topic, operation, operation_label, table_name,
                    record_key, detail, event_ts, sampled_at)
                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                (
                    source,
                    topic,
                    op,
                    OP_LABELS.get(op, op),
                    table,
                    key,
                    detail,
                    event_ts,
                    now,
                ),
            )


def collect_spark(cur):
    now = datetime.now(timezone.utc)
    cur.execute("DELETE FROM monitor.spark_applications")
    try:
        data = fetch_json(f"{SPARK_MASTER}/json/", timeout=5)
        apps = []
        if isinstance(data, dict):
            # Standalone master JSON
            for a in data.get("activeapps", []) or []:
                apps.append(a)
            for a in data.get("completedapps", []) or []:
                apps.append(a)
        for a in apps[:20]:
            cur.execute(
                """INSERT INTO monitor.spark_applications
                   (app_id, app_name, state, cores, memory_mb, duration_sec, checked_at)
                   VALUES (%s,%s,%s,%s,%s,%s,%s)
                   ON CONFLICT (app_id) DO UPDATE SET
                     app_name=EXCLUDED.app_name, state=EXCLUDED.state,
                     cores=EXCLUDED.cores, memory_mb=EXCLUDED.memory_mb,
                     duration_sec=EXCLUDED.duration_sec, checked_at=EXCLUDED.checked_at""",
                (
                    a.get("id", "unknown"),
                    a.get("name", "Spark App"),
                    "RUNNING" if "attempts" not in a else "COMPLETED",
                    int(a.get("cores", 0) or 0),
                    int((a.get("memory", 0) or 0) / 1024 / 1024),
                    int(a.get("duration", 0) / 1000) if a.get("duration") else 0,
                    now,
                ),
            )
    except Exception as e:
        # Placeholder row so dashboard shows Spark host status
        cur.execute(
            """INSERT INTO monitor.spark_applications
               (app_id, app_name, state, cores, memory_mb, duration_sec, checked_at)
               VALUES (%s,%s,%s,%s,%s,%s,%s)
               ON CONFLICT (app_id) DO UPDATE SET state=EXCLUDED.state, checked_at=EXCLUDED.checked_at""",
            (
                "spark-master",
                f"Spark Master @ {SPARK_MASTER}",
                "REACHABLE" if "Connection" not in str(e) else "UNREACHABLE",
                0,
                0,
                0,
                now,
            ),
        )


def main():
    conn = psycopg2.connect(PG_DSN)
    conn.autocommit = True
    cur = conn.cursor()
    print("Collecting Debezium...")
    collect_debezium(cur)
    print("Collecting Kafka CDC samples...")
    collect_kafka_cdc(cur)
    print("Collecting Spark...")
    collect_spark(cur)
    cur.close()
    conn.close()
    print("Done.")


if __name__ == "__main__":
    main()