atc-agents/api/trino_federated.py

"""Trino federated business analytics + Hadoop lakehouse (external Iceberg tables).

Provides:
  - /api/federated/catalogs   : the federated catalog landscape (estimates, fast)
  - /api/federated/marquee    : a single cross-source SQL joining PG+MySQL+Mongo
                                (the federation proof) — computed in the background, cached
  - /api/federated/lake       : business analytics over the Hadoop Iceberg lake
                                (iceberg.hadoop external tables) — live & fast
  - /api/federated/materialize: (re)build the Hadoop external tables from the sources
  - /api/federated/dictionary : business data dictionary with masked / visible flags
                                (also fed to the LLM so it knows the data in detail)
"""

from __future__ import annotations

import json
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from fastapi import APIRouter, Body
from fastapi.responses import JSONResponse

router = APIRouter(prefix="/api/federated", tags=["federated"])

BUSINESS_CATALOGS = ["postgres_sales", "mysql_hr", "mongodb_supplychain", "cassandra_telemetry", "iceberg"]
ALL_CATALOGS = BUSINESS_CATALOGS + ["kafka", "system"]

# Hadoop external tables (Iceberg-on-HDFS) materialized from the federated sources.
# Each: target table name in iceberg.hadoop  ->  source SELECT (timestamps cast to
# timestamp(6) which Iceberg requires; measures cast to double).
LAKE_SPECS: list[tuple[str, str, str]] = [
    ("orders_ext", "Sales orders (PostgreSQL)",
     "SELECT order_id, customer_id, customer_name, customer_email, product_id, region, "
     "sales_channel, order_status, currency, CAST(amount AS double) amount, "
     "CAST(order_ts AS timestamp(6)) order_ts "
     "FROM postgres_sales.public.sales_orders WHERE customer_name IS NOT NULL LIMIT 120000"),
    ("employees_ext", "Employee events (MySQL)",
     "SELECT employee_id, employee_name, employee_email, employee_phone, department, role_name, "
     "region, event_type, CAST(salary_change AS double) salary_change, "
     "CAST(event_ts AS timestamp(6)) event_ts "
     "FROM mysql_hr.hr.employee_events WHERE employee_name IS NOT NULL LIMIT 120000"),
    ("supply_events_ext", "Supply chain events (MongoDB)",
     "SELECT event_id, type, region, source, CAST(amount AS double) amount, "
     "CAST(ts AS timestamp(6)) ts FROM mongodb_supplychain.supplychain.events LIMIT 120000"),
    ("telemetry_ext", "Device telemetry (Cassandra)",
     "SELECT device_id, CAST(metric_ts AS timestamp(6)) metric_ts, metric_type, "
     "CAST(metric_value AS double) metric_value FROM cassandra_telemetry.telemetry.device_metrics LIMIT 120000"),
    ("customers_ext", "Distinct customers (PostgreSQL)",
     "SELECT DISTINCT customer_id, customer_name, customer_email, region, sales_channel "
     "FROM postgres_sales.public.sales_orders WHERE customer_name IS NOT NULL LIMIT 60000"),
    ("products_ext", "Distinct products (PostgreSQL)",
     "SELECT DISTINCT product_id, region, sales_channel FROM postgres_sales.public.sales_orders "
     "WHERE product_id IS NOT NULL LIMIT 20000"),
]

MARQUEE_SQL = (
    "WITH o AS (SELECT region, count(*) orders, sum(amount) revenue "
    "FROM postgres_sales.public.sales_orders GROUP BY region),\n"
    "     e AS (SELECT region, count(*) hr_events FROM mysql_hr.hr.employee_events GROUP BY region),\n"
    "     s AS (SELECT region, count(*) supply_events FROM mongodb_supplychain.supplychain.events GROUP BY region)\n"
    "SELECT COALESCE(o.region, e.region, s.region) AS region,\n"
    "       o.orders, o.revenue, e.hr_events, s.supply_events\n"
    "FROM o FULL JOIN e ON o.region = e.region\n"
    "       FULL JOIN s ON COALESCE(o.region, e.region) = s.region\n"
    "ORDER BY o.revenue DESC NULLS LAST"
)

# Federated reach - ONE SQL touching every database/engine in the stack.
# Telemetry has no region dimension, so instead of a misleading join we
# summarise each source side-by-side in a single UNION ALL query.
MATRIX_SQL = (
    "SELECT 1 ord, 'PostgreSQL' source, 'postgres_sales' catalog, 'public.sales_orders' dataset,\n"
    "       count(*) records, CAST(sum(amount) AS double) metric, 'total revenue' metric_label\n"
    "FROM postgres_sales.public.sales_orders\n"
    "UNION ALL SELECT 2,'MySQL','mysql_hr','hr.employee_events',count(*),\n"
    "       CAST(count(DISTINCT employee_id) AS double),'distinct employees' FROM mysql_hr.hr.employee_events\n"
    "UNION ALL SELECT 3,'MongoDB','mongodb_supplychain','supplychain.events',count(*),\n"
    "       CAST(sum(amount) AS double),'event value' FROM mongodb_supplychain.supplychain.events\n"
    "UNION ALL SELECT 4,'Cassandra','cassandra_telemetry','telemetry.device_metrics',count(*),\n"
    "       CAST(avg(metric_value) AS double),'avg metric value' FROM cassandra_telemetry.telemetry.device_metrics\n"
    "UNION ALL SELECT 5,'Hadoop / Iceberg','iceberg','hadoop.orders_ext',count(*),\n"
    "       CAST(sum(amount) AS double),'lake revenue' FROM iceberg.hadoop.orders_ext\n"
    "ORDER BY ord"
)
MATRIX_CATALOGS = ["postgres_sales", "mysql_hr", "mongodb_supplychain", "cassandra_telemetry", "iceberg"]


def _trino(sql: str, limit: int = 500) -> dict[str, Any]:
    import sql_console as s
    return s._run_trino(sql, limit)


def _rows_as_dicts(res: dict) -> list[dict]:
    cols = res.get("columns", [])
    return [{cols[i]: r[i] for i in range(min(len(cols), len(r)))} for r in res.get("rows", [])]


# ──────────────────────────────────────────────────────────────────────────────
# Marquee federated query (cross-source) — background cached
# ──────────────────────────────────────────────────────────────────────────────
MARQUEE_PATH = Path("/data/federated_marquee.json")
_MARQUEE_TTL = 1800.0
_marquee: dict[str, Any] = {"running": False, "data": None}
_marquee_lock = threading.Lock()


def _load_marquee() -> None:
    try:
        if MARQUEE_PATH.exists():
            _marquee["data"] = json.loads(MARQUEE_PATH.read_text())
    except Exception:
        pass


def _marquee_worker() -> None:
    try:
        # Run the two federated queries concurrently so the total wait stays
        # close to the slower of the two (region scorecard ~ matrix).
        from concurrent.futures import ThreadPoolExecutor

        def timed(sql: str) -> tuple[dict, int]:
            a = time.time()
            r = _trino(sql, 50)
            return r, int((time.time() - a) * 1000)

        with ThreadPoolExecutor(max_workers=2) as ex:
            f_region = ex.submit(timed, MARQUEE_SQL)
            f_matrix = ex.submit(timed, MATRIX_SQL)
            res, elapsed = f_region.result()
            mres, melapsed = f_matrix.result()

        data = {
            "ok": res.get("ok", False),
            "sql": MARQUEE_SQL,
            "catalogs": ["postgres_sales", "mysql_hr", "mongodb_supplychain"],
            "elapsed_ms": elapsed,
            "rows": _rows_as_dicts(res) if res.get("ok") else [],
            "matrix": {
                "ok": mres.get("ok", False),
                "sql": MATRIX_SQL,
                "catalogs": MATRIX_CATALOGS,
                "elapsed_ms": melapsed,
                "rows": _rows_as_dicts(mres) if mres.get("ok") else [],
                "error": None if mres.get("ok") else str(mres.get("error", ""))[:300],
            },
            "error": None if res.get("ok") else str(res.get("error", ""))[:300],
            "generated_at": datetime.now(timezone.utc).isoformat(),
        }
        _marquee["data"] = data
        try:
            MARQUEE_PATH.write_text(json.dumps(data))
        except Exception:
            pass
    finally:
        _marquee["running"] = False


def _ensure_marquee(force: bool = False) -> None:
    with _marquee_lock:
        if _marquee["running"]:
            return
        data = _marquee["data"]
        fresh = data and (time.time() - _ts(data.get("generated_at")) < _MARQUEE_TTL)
        if fresh and not force:
            return
        _marquee["running"] = True
    threading.Thread(target=_marquee_worker, daemon=True).start()


def _ts(iso: str | None) -> float:
    if not iso:
        return 0.0
    try:
        return datetime.fromisoformat(iso).timestamp()
    except Exception:
        return 0.0


@router.get("/marquee")
async def get_marquee():
    if _marquee["data"] is None:
        _load_marquee()
    _ensure_marquee()
    data = _marquee["data"]
    return {
        "ok": True,
        "running": _marquee["running"],
        "marquee": data,
        "sql": MARQUEE_SQL,
    }


@router.post("/marquee/refresh")
async def refresh_marquee():
    _ensure_marquee(force=True)
    return {"ok": True, "running": True}


# ──────────────────────────────────────────────────────────────────────────────
# Catalog landscape (fast — SHOW + planner estimates)
# ──────────────────────────────────────────────────────────────────────────────
@router.get("/catalogs")
async def get_catalogs():
    import sql_console as s
    cat_res = _trino("SHOW CATALOGS", 100)
    catalogs = [r[0] for r in cat_res.get("rows", [])] if cat_res.get("ok") else ALL_CATALOGS
    totals = {
        "orders": s._table_row_count("postgres", "public.sales_orders"),
        "hr_events": s._table_row_count("mysql", "hr.employee_events"),
        "supply_events": s._table_row_count("mongodb", "supplychain.events"),
    }
    out = []
    labels = {
        "postgres_sales": ("PostgreSQL", "Sales / orders OLTP", "#336791"),
        "mysql_hr": ("MySQL", "HR / workforce", "#00758f"),
        "mongodb_supplychain": ("MongoDB", "Supply chain events", "#4db33d"),
        "cassandra_telemetry": ("Cassandra", "Device telemetry", "#1287b1"),
        "iceberg": ("Iceberg / Hadoop", "Lakehouse external tables", "#5b8def"),
        "kafka": ("Kafka", "CDC + streaming topics", "#231f20"),
        "system": ("Trino", "Engine metadata", "#dd00a1"),
    }
    for c in catalogs:
        lbl, desc, color = labels.get(c, (c, "", "#888"))
        item: dict[str, Any] = {"catalog": c, "label": lbl, "desc": desc, "color": color, "business": c in BUSINESS_CATALOGS}
        if c == "postgres_sales":
            item["rows"] = totals["orders"]
        elif c == "mysql_hr":
            item["rows"] = totals["hr_events"]
        elif c == "mongodb_supplychain":
            item["rows"] = totals["supply_events"]
        out.append(item)
    return {"ok": True, "catalogs": out, "source_totals": totals, "count": len(out)}


# ──────────────────────────────────────────────────────────────────────────────
# Realtime business dashboard — fast: instant source estimates + short-TTL
# cached aggregations over the small materialized Hadoop lake tables.
# ──────────────────────────────────────────────────────────────────────────────
_live_aggs: dict[str, Any] = {"ts": 0.0, "data": None}
_LIVE_AGG_TTL = 6.0
_avg_order_cache: dict[str, Any] = {"ts": 0.0, "val": 0.0}


def _avg_order_value() -> float:
    now = time.time()
    if now - _avg_order_cache["ts"] < 300 and _avg_order_cache["val"]:
        return _avg_order_cache["val"]
    res = _trino("SELECT avg(amount) FROM iceberg.hadoop.orders_ext", 1)
    val = 0.0
    if res.get("ok"):
        try:
            val = float((res.get("rows") or [[0]])[0][0] or 0)
        except Exception:
            val = 0.0
    _avg_order_cache["val"] = val
    _avg_order_cache["ts"] = now
    return val


def _region_matrix_from_lake() -> list[dict]:
    o = _terms("SELECT region k, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY region", "k", "c", "rev")
    e = _terms("SELECT region k, count(*) c FROM iceberg.hadoop.employees_ext GROUP BY region", "k", "c")
    sup = _terms("SELECT region k, count(*) c FROM iceberg.hadoop.supply_events_ext GROUP BY region", "k", "c")
    regions: dict[str, dict] = {}
    for x in o:
        regions.setdefault(x["key"], {})["orders"] = x["count"]
        regions[x["key"]]["revenue"] = x.get("value", 0)
    for x in e:
        regions.setdefault(x["key"], {})["hr_events"] = x["count"]
    for x in sup:
        regions.setdefault(x["key"], {})["supply_events"] = x["count"]
    rows = [{"region": k, "orders": v.get("orders", 0), "revenue": v.get("revenue", 0),
             "hr_events": v.get("hr_events", 0), "supply_events": v.get("supply_events", 0)}
            for k, v in regions.items() if k]
    rows.sort(key=lambda r: r.get("revenue") or 0, reverse=True)
    return rows


def _live_business_aggs() -> dict[str, Any]:
    now = time.time()
    if _live_aggs["data"] is not None and now - _live_aggs["ts"] < _LIVE_AGG_TTL:
        return _live_aggs["data"]
    data = {
        "orders_by_region": _terms("SELECT region, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY region ORDER BY rev DESC", "region", "c", "rev"),
        "orders_by_status": _terms("SELECT order_status k, count(*) c FROM iceberg.hadoop.orders_ext GROUP BY order_status ORDER BY c DESC", "k", "c"),
        "orders_by_channel": _terms("SELECT sales_channel k, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY sales_channel ORDER BY rev DESC", "k", "c", "rev"),
        "top_customers": _terms("SELECT customer_name k, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY customer_name ORDER BY rev DESC LIMIT 8", "k", "c", "rev"),
        "telemetry_by_metric": _terms("SELECT metric_type k, count(*) c, avg(metric_value) v FROM iceberg.hadoop.telemetry_ext GROUP BY metric_type ORDER BY c DESC", "k", "c", "v"),
        "supply_by_type": _terms("SELECT type k, count(*) c FROM iceberg.hadoop.supply_events_ext GROUP BY type ORDER BY c DESC", "k", "c"),
        "region_matrix": _region_matrix_from_lake(),
    }
    _live_aggs["data"] = data
    _live_aggs["ts"] = now
    return data


# ──────────────────────────────────────────────────────────────────────────────
# Continuous live generator — keeps the platform "alive": while the Live
# dashboard is open it streams small, randomly-sized batches of business rows
# into the real source databases (PostgreSQL / MySQL / MongoDB / Cassandra),
# which CDC then propagates downstream. Batch sizes fluctuate every tick so the
# throughput visibly goes up and down. It only runs while someone is watching
# (the /live poll refreshes a heartbeat) so the tables don't grow unbounded.
# ──────────────────────────────────────────────────────────────────────────────
import random as _rnd
from collections import deque as _deque

_GEN: dict[str, Any] = {
    "enabled": True,
    "running": False,
    "interval": 4.0,
    "last_seen": 0.0,
    "last_tick": 0.0,
    "counts": {"orders": 0, "hr_events": 0, "supply_events": 0, "telemetry": 0},
    "last_batch": {"orders": 0, "hr_events": 0, "supply_events": 0, "telemetry": 0},
    "by_region": {},
    "by_status": {},
    "tick_value": 0.0,
    "feed": _deque(maxlen=14),
    "base": None,
}
_gen_lock = threading.Lock()
_gen_conns: dict[str, Any] = {"pg": None, "mysql": None, "mongo": None, "cass": None}

_REGIONS = ["NA", "EU", "APAC", "LATAM", "EMEA", "MEA"]
_CHANNELS = ["B2B", "B2C", "ONLINE", "PARTNER", "RETAIL"]
_STATUSES = ["NEW", "PAID", "SHIPPED", "DELIVERED", "RETURNED", "CANCELLED"]
_CURR = ["EUR", "USD", "GBP", "JPY"]
_DEPTS = ["Engineering", "Sales", "Support", "Operations", "Finance", "HR", "Marketing"]
_ROLES = ["Analyst", "Engineer", "Manager", "Lead", "Specialist", "Director"]
_EVT = ["HIRE", "PROMOTION", "SALARY_CHANGE", "TRANSFER", "REVIEW", "EXIT"]
_SUPPLY = ["INSERT", "UPDATE", "REPLENISH", "SHIPMENT", "RETURN"]
_SRC = ["CRM", "ERP", "WMS", "API"]
_METRICS = ["temperature", "humidity", "pressure", "voltage", "current"]


def _gen_pg():
    import psycopg2
    import sql_console as s
    c = _gen_conns["pg"]
    if c is None or getattr(c, "closed", 1):
        c = psycopg2.connect(host=s.DB_HOST, port=s.PG_PORT, user=s.PG_USER,
                             password=s.PG_PASS, dbname=s.PG_DB, connect_timeout=6)
        c.autocommit = True
        _gen_conns["pg"] = c
    return c


def _gen_mysql():
    import pymysql
    import sql_console as s
    c = _gen_conns["mysql"]
    if c is None:
        c = pymysql.connect(host=s.DB_HOST, port=s.MYSQL_PORT, user=s.MYSQL_USER,
                            password=s.MYSQL_PASS, database=s.MYSQL_DB, connect_timeout=6,
                            autocommit=True)
        _gen_conns["mysql"] = c
    else:
        c.ping(reconnect=True)
    return c


def _gen_mongo():
    import sql_console as s
    c = _gen_conns["mongo"]
    if c is None:
        c = s._mongo_client()
        _gen_conns["mongo"] = c
    return c[s.MONGO_DB]


def _gen_cass():
    import sql_console as s
    sess = _gen_conns["cass"]
    if sess is None:
        cluster = s._cass_cluster()
        sess = cluster.connect()
        _gen_conns["cass"] = sess
    return sess


def _gen_reset(key: str):
    try:
        c = _gen_conns.get(key)
        if c is not None:
            c.close() if key != "cass" else c.cluster.shutdown()
    except Exception:
        pass
    _gen_conns[key] = None


# Row builders — shared by the background streamer and the on-demand "Generate
# data" button, so both produce identical, realistic business rows.
_PG_INSERT = ("INSERT INTO public.sales_orders "
              "(customer_id,product_id,region,sales_channel,order_ts,amount,currency,order_status) "
              "VALUES (%s,%s,%s,%s,%s,%s,%s,%s)")
_MYSQL_INSERT = ("INSERT INTO employee_events "
                 "(employee_id,department,role_name,region,event_type,salary_change,event_ts) "
                 "VALUES (%s,%s,%s,%s,%s,%s,%s)")


def _order_rows(n: int):
    import datetime as dt
    now = dt.datetime.utcnow()
    by_r: dict[str, int] = {}
    by_s: dict[str, int] = {}
    val = 0.0
    rows = []
    for _ in range(n):
        r = _rnd.choice(_REGIONS)
        st = _rnd.choices(_STATUSES, weights=[5, 6, 5, 8, 2, 2])[0]
        ch = _rnd.choice(_CHANNELS)
        amt = round(_rnd.uniform(15, 9500), 2)
        rows.append((_rnd.randint(1, 20000), _rnd.randint(1, 5000), r, ch, now, amt, _rnd.choice(_CURR), st))
        by_r[r] = by_r.get(r, 0) + 1
        by_s[st] = by_s.get(st, 0) + 1
        val += amt
    return rows, by_r, by_s, round(val, 2)


def _hr_rows(n: int):
    import datetime as dt
    now = dt.datetime.utcnow()
    return [(_rnd.randint(1, 100000), _rnd.choice(_DEPTS), _rnd.choice(_ROLES),
             _rnd.choice(_REGIONS), _rnd.choice(_EVT), round(_rnd.uniform(-2000, 6000), 2), now)
            for _ in range(n)]


def _supply_docs(n: int):
    import datetime as dt
    import uuid
    now = dt.datetime.utcnow()
    return [{"event_id": str(uuid.uuid4()), "type": _rnd.choice(_SUPPLY),
             "region": _rnd.choice(_REGIONS), "source": _rnd.choice(_SRC),
             "amount": round(_rnd.uniform(10, 40000), 2), "ts": now.isoformat()}
            for _ in range(n)]


def _tel_rows(n: int):
    import datetime as dt
    now = dt.datetime.utcnow()
    return [(f"device-{_rnd.randint(1, 99999)}", now, _rnd.choice(_METRICS),
             round(_rnd.uniform(0, 100), 3), "") for _ in range(n)]


_TEL_INSERT_TPL = ("INSERT INTO {ks}.device_metrics "
                   "(device_id, metric_ts, metric_type, metric_value, payload) VALUES (%s,%s,%s,%s,%s)")


def _gen_orders(n: int):
    rows, by_r, by_s, val = _order_rows(n)
    _gen_pg().cursor().executemany(_PG_INSERT, rows)
    return rows, by_r, by_s, val


def _gen_hr(n: int):
    rows = _hr_rows(n)
    _gen_mysql().cursor().executemany(_MYSQL_INSERT, rows)
    return rows


def _gen_supply(n: int):
    docs = _supply_docs(n)
    if docs:
        _gen_mongo()["events"].insert_many([dict(d) for d in docs])
    return docs


def _gen_tel(n: int):
    import sql_console as s
    sess = _gen_cass()
    cql = _TEL_INSERT_TPL.format(ks=s.CASS_KS)
    rows = _tel_rows(n)
    for row in rows:
        sess.execute(cql, row)
    return rows


def _generate_once(orders: int, hr: int, supply: int, tel: int) -> dict[str, Any]:
    """On-demand burst using FRESH short-lived connections (safe to run from a
    request thread alongside the background streamer). Returns inserted counts."""
    import sql_console as s
    out = {"orders": 0, "hr_events": 0, "supply_events": 0, "telemetry": 0}
    by_r: dict[str, int] = {}
    by_s: dict[str, int] = {}
    val = 0.0
    order_built: list = []
    hr_built: list = []
    supply_built: list = []
    tel_built: list = []
    if orders > 0:
        try:
            import psycopg2
            rows, by_r, by_s, val = _order_rows(orders)
            order_built = rows
            c = psycopg2.connect(host=s.DB_HOST, port=s.PG_PORT, user=s.PG_USER, password=s.PG_PASS, dbname=s.PG_DB, connect_timeout=8)
            try:
                c.autocommit = True
                c.cursor().executemany(_PG_INSERT, rows)
                out["orders"] = orders
            finally:
                c.close()
        except Exception:
            pass
    if hr > 0:
        try:
            import pymysql
            hr_built = _hr_rows(hr)
            c = pymysql.connect(host=s.DB_HOST, port=s.MYSQL_PORT, user=s.MYSQL_USER, password=s.MYSQL_PASS, database=s.MYSQL_DB, connect_timeout=8, autocommit=True)
            try:
                c.cursor().executemany(_MYSQL_INSERT, hr_built)
                out["hr_events"] = hr
            finally:
                c.close()
        except Exception:
            pass
    if supply > 0:
        try:
            supply_built = _supply_docs(supply)
            cli = s._mongo_client()
            try:
                cli[s.MONGO_DB]["events"].insert_many([dict(d) for d in supply_built])
                out["supply_events"] = supply
            finally:
                cli.close()
        except Exception:
            pass
    if tel > 0:
        try:
            tel_built = _tel_rows(tel)
            cluster = s._cass_cluster()
            sess = cluster.connect()
            try:
                cql = _TEL_INSERT_TPL.format(ks=s.CASS_KS)
                for row in tel_built:
                    sess.execute(cql, row)
                out["telemetry"] = tel
            finally:
                cluster.shutdown()
        except Exception:
            pass
    # Land the generated batch in S3 through the pipeline stages (Kafka→S3 CDC
    # archive + Spark→S3 curated masked) so Object Storage reflects it at once.
    try:
        from storage_s3 import archive_generated_batch
        archive_generated_batch(order_built if out["orders"] else [],
                                hr_built if out["hr_events"] else [],
                                supply_built if out["supply_events"] else [],
                                tel_built if out["telemetry"] else [], force=True)
    except Exception:
        pass
    # fold into the live counters + feed so the dashboard reflects it instantly
    with _gen_lock:
        c = _GEN["counts"]
        for k in out:
            c[k] += out[k]
        _GEN["last_batch"] = {"orders": out["orders"], "hr_events": out["hr_events"],
                              "supply_events": out["supply_events"], "telemetry": out["telemetry"]}
        _GEN["last_tick"] = time.time()  # marks recent activity → Data Flow CDC edges pulse
        if out["orders"]:
            _GEN["by_region"] = by_r
            _GEN["by_status"] = by_s
            _GEN["tick_value"] = val
            top = max(by_r, key=by_r.get) if by_r else "—"
            _GEN["feed"].appendleft({
                "ts": datetime.now(timezone.utc).isoformat(),
                "text": f"⚡ manual burst: +{out['orders']} orders · €{int(val):,} · top {top} · +{out['telemetry']} telemetry · +{out['hr_events']} HR · +{out['supply_events']} supply",
            })
    out["revenue"] = val
    return out


def _gen_tick():
    # fluctuating batch sizes, with the occasional spike, so throughput moves up & down
    no = _rnd.randint(2, 40)
    if _rnd.random() < 0.18:
        no += _rnd.randint(25, 70)
    nh = _rnd.randint(0, 18)
    ns = _rnd.randint(0, 16)
    nt = _rnd.randint(8, 55)
    by_r: dict[str, int] = {}
    by_s: dict[str, int] = {}
    val = 0.0
    order_built: list = []
    hr_built: list = []
    supply_built: list = []
    tel_built: list = []
    try:
        order_built, by_r, by_s, val = _gen_orders(no)
    except Exception:
        _gen_reset("pg"); no = 0
    try:
        hr_built = _gen_hr(nh)
    except Exception:
        _gen_reset("mysql"); nh = 0
    try:
        supply_built = _gen_supply(ns)
    except Exception:
        _gen_reset("mongo"); ns = 0
    try:
        tel_built = _gen_tel(nt)
    except Exception:
        _gen_reset("cass"); nt = 0
    # stream the batch into S3 (buffered like a Kafka-Connect S3 sink)
    try:
        from storage_s3 import archive_generated_batch
        archive_generated_batch(order_built, hr_built, supply_built, tel_built, force=False)
    except Exception:
        pass
    with _gen_lock:
        c = _GEN["counts"]
        c["orders"] += no
        c["hr_events"] += nh
        c["supply_events"] += ns
        c["telemetry"] += nt
        _GEN["last_batch"] = {"orders": no, "hr_events": nh, "supply_events": ns, "telemetry": nt}
        _GEN["by_region"] = by_r
        _GEN["by_status"] = by_s
        _GEN["tick_value"] = val
        _GEN["last_tick"] = time.time()
        if no:
            top = max(by_r, key=by_r.get) if by_r else "—"
            _GEN["feed"].appendleft({
                "ts": datetime.now(timezone.utc).isoformat(),
                "text": f"+{no} orders · €{int(val):,} · top {top} ({by_r.get(top, 0)}) · +{nt} telemetry · +{nh} HR",
            })


def _gen_loop():
    while True:
        try:
            if _GEN["enabled"] and (time.time() - _GEN["last_seen"] < 25):
                _GEN["running"] = True
                _gen_tick()
            else:
                _GEN["running"] = False
        except Exception:
            _GEN["running"] = False
        time.sleep(max(2.0, float(_GEN["interval"])))


threading.Thread(target=_gen_loop, daemon=True, name="live-generator").start()


def generator_active(window_s: float = 12.0) -> bool:
    """True when the live generator (continuous loop OR a manual 'Generate data'
    burst) wrote rows very recently. The Data Flow graph uses this to pulse the
    generator→source edges so generated data is visible flowing into the sources."""
    return (time.time() - float(_GEN.get("last_tick") or 0.0)) < window_s


@router.post("/live/generator")
async def toggle_generator(body: dict = Body(default={})):
    if "enabled" in body:
        _GEN["enabled"] = bool(body["enabled"])
    if "interval" in body:
        try:
            _GEN["interval"] = max(2.0, min(30.0, float(body["interval"])))
        except Exception:
            pass
    _GEN["last_seen"] = time.time()
    return {"ok": True, "enabled": _GEN["enabled"], "interval": _GEN["interval"], "running": _GEN["running"]}


@router.post("/live/heartbeat")
async def live_heartbeat(body: dict = Body(default={})):
    """Keep-alive for the continuous generator from any watching view (e.g. the
    Data Flow tab's Run control). active=true → generate; active=false → stop."""
    active = bool(body.get("active", True))
    if active:
        _GEN["enabled"] = True
        _GEN["last_seen"] = time.time()
    else:
        _GEN["last_seen"] = 0.0
    return {"ok": True, "running": _GEN["running"], "enabled": _GEN["enabled"]}


@router.post("/generate")
async def generate_now(body: dict = Body(default={})):
    """Manual one-shot burst into the source systems (the Data Flow "Generate
    data" button). `rows` controls the order volume; the other sources scale
    with it. CDC streams everything downstream automatically."""
    from starlette.concurrency import run_in_threadpool
    rows = int(body.get("rows", 500) or 500)
    rows = max(1, min(20000, rows))
    orders = rows
    hr = max(1, rows // 4)
    supply = max(1, rows // 4)
    tel = max(1, rows // 2)
    out = await run_in_threadpool(_generate_once, orders, hr, supply, tel)
    total = out["orders"] + out["hr_events"] + out["supply_events"] + out["telemetry"]
    return {"ok": True, "requested": rows, "inserted": out, "total": total}


@router.get("/live")
async def get_live():
    import sql_console as s
    _GEN["last_seen"] = time.time()  # heartbeat: keeps the generator running while watched

    # Cassandra has no cheap estimate — reuse the exact count from the cached
    # federated matrix query when available.
    if _marquee.get("data") is None:
        _load_marquee()
    mq = _marquee.get("data") or {}
    cass_base = 0
    for r in ((mq.get("matrix") or {}).get("rows") or []):
        if r.get("catalog") == "cassandra_telemetry":
            try:
                cass_base = int(r.get("records") or 0)
            except Exception:
                cass_base = 0

    # One-time base snapshot of source sizes; every subsequent reading is
    # base + rows the generator has streamed in, so the counters move smoothly
    # and in lock-step with the live activity feed.
    with _gen_lock:
        if _GEN["base"] is None:
            _GEN["base"] = {
                "orders": s._table_row_count("postgres", "public.sales_orders") or 0,
                "hr_events": s._table_row_count("mysql", "hr.employee_events") or 0,
                "supply_events": s._table_row_count("mongodb", "supplychain.events") or 0,
                "telemetry": cass_base,
            }
        elif cass_base and not _GEN["base"].get("telemetry"):
            _GEN["base"]["telemetry"] = cass_base
        base = dict(_GEN["base"])
        gc = dict(_GEN["counts"])
        gen_view = {
            "enabled": _GEN["enabled"],
            "running": _GEN["running"],
            "interval": _GEN["interval"],
            "counts": dict(_GEN["counts"]),
            "last_batch": dict(_GEN["last_batch"]),
            "tick_value": _GEN["tick_value"],
            "by_region": [{"key": k, "count": v} for k, v in sorted(_GEN["by_region"].items(), key=lambda kv: -kv[1])],
            "by_status": [{"key": k, "count": v} for k, v in sorted(_GEN["by_status"].items(), key=lambda kv: -kv[1])],
            "feed": list(_GEN["feed"]),
        }

    orders = base["orders"] + gc["orders"]
    hr = base["hr_events"] + gc["hr_events"]
    supply = base["supply_events"] + gc["supply_events"]
    telemetry = base["telemetry"] + gc["telemetry"]

    avg_order = _avg_order_value()
    sources = [
        {"key": "orders", "label": "Orders", "engine": "PostgreSQL", "catalog": "postgres_sales", "rows": orders, "added": gc["orders"], "color": "#fbbf24"},
        {"key": "hr_events", "label": "HR events", "engine": "MySQL", "catalog": "mysql_hr", "rows": hr, "added": gc["hr_events"], "color": "#60a5fa"},
        {"key": "supply_events", "label": "Supply events", "engine": "MongoDB", "catalog": "mongodb_supplychain", "rows": supply, "added": gc["supply_events"], "color": "#a78bfa"},
        {"key": "telemetry", "label": "Telemetry", "engine": "Cassandra", "catalog": "cassandra_telemetry", "rows": telemetry, "added": gc["telemetry"], "color": "#22d3ee"},
    ]
    return {
        "ok": True,
        "ts": datetime.now(timezone.utc).isoformat(),
        "sources": sources,
        "generator": gen_view,
        "totals": {
            "records": orders + hr + supply + telemetry,
            "revenue_est": round(orders * avg_order, 2),
            "avg_order": round(avg_order, 2),
        },
        "business": _live_business_aggs(),
    }


# ──────────────────────────────────────────────────────────────────────────────
# Hadoop lake analytics (live over the small materialized external tables)
# ──────────────────────────────────────────────────────────────────────────────
def _terms(sql: str, key_col: str, count_col: str = "c", val_col: str | None = None) -> list[dict]:
    res = _trino(sql, 100)
    if not res.get("ok"):
        return []
    cols = res.get("columns", [])
    idx = {c: i for i, c in enumerate(cols)}
    out = []
    for r in res.get("rows", []):
        item = {"key": r[idx[key_col]] if key_col in idx else None,
                "count": r[idx[count_col]] if count_col in idx else 0}
        if val_col and val_col in idx:
            item["value"] = round(float(r[idx[val_col]] or 0), 2)
        out.append(item)
    return out


def _lake_table_exists(name: str) -> bool:
    res = _trino(f"SELECT count(*) FROM iceberg.hadoop.{name}", 1)
    return res.get("ok", False)


@router.get("/lake")
async def get_lake():
    # which external tables are present
    tbl_res = _trino("SHOW TABLES FROM iceberg.hadoop", 200)
    tables = [r[0] for r in tbl_res.get("rows", [])] if tbl_res.get("ok") else []
    table_stats = []
    for t in tables:
        cnt = _trino(f"SELECT count(*) FROM iceberg.hadoop.{t}", 1)
        n = cnt.get("rows", [[0]])[0][0] if cnt.get("ok") else 0
        table_stats.append({"table": t, "rows": n})

    has_orders = "orders_ext" in tables
    has_emp = "employees_ext" in tables
    has_sup = "supply_events_ext" in tables
    has_tel = "telemetry_ext" in tables

    orders = {}
    if has_orders:
        orders = {
            "by_region": _terms("SELECT region, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY region ORDER BY rev DESC", "region", "c", "rev"),
            "by_channel": _terms("SELECT sales_channel k, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY sales_channel ORDER BY rev DESC", "k", "c", "rev"),
            "by_status": _terms("SELECT order_status k, count(*) c FROM iceberg.hadoop.orders_ext GROUP BY order_status ORDER BY c DESC", "k", "c"),
            "top_customers": _terms("SELECT customer_name k, count(*) c, sum(amount) rev FROM iceberg.hadoop.orders_ext GROUP BY customer_name ORDER BY rev DESC LIMIT 10", "k", "c", "rev"),
            "revenue": (_trino("SELECT sum(amount), avg(amount) FROM iceberg.hadoop.orders_ext", 1).get("rows") or [[0, 0]])[0],
        }
    hr = {}
    if has_emp:
        hr = {
            "by_department": _terms("SELECT department k, count(*) c FROM iceberg.hadoop.employees_ext GROUP BY department ORDER BY c DESC", "k", "c"),
            "by_role": _terms("SELECT role_name k, count(*) c FROM iceberg.hadoop.employees_ext GROUP BY role_name ORDER BY c DESC LIMIT 12", "k", "c"),
        }
    supply = {}
    if has_sup:
        supply = {"by_type": _terms("SELECT type k, count(*) c FROM iceberg.hadoop.supply_events_ext GROUP BY type ORDER BY c DESC", "k", "c")}
    telemetry = {}
    if has_tel:
        telemetry = {"by_metric": _terms("SELECT metric_type k, count(*) c, avg(metric_value) v FROM iceberg.hadoop.telemetry_ext GROUP BY metric_type ORDER BY c DESC", "k", "c", "v")}

    return {
        "ok": True,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "tables": table_stats,
        "total_rows": sum(t["rows"] or 0 for t in table_stats),
        "orders": orders,
        "hr": hr,
        "supply": supply,
        "telemetry": telemetry,
    }


# ──────────────────────────────────────────────────────────────────────────────
# Materialize the Hadoop external tables from the sources
# ──────────────────────────────────────────────────────────────────────────────
_mat_state: dict[str, Any] = {"running": False, "current": None, "done": [], "errors": [], "started_at": None, "finished_at": None}
_mat_lock = threading.Lock()


def _materialize_worker() -> None:
    try:
        for name, label, sql in LAKE_SPECS:
            _mat_state["current"] = f"{name} ({label})"
            try:
                _trino(f"DROP TABLE IF EXISTS iceberg.hadoop.{name}", 1)
                res = _trino(f"CREATE TABLE iceberg.hadoop.{name} AS {sql}", 5)
                if res.get("ok"):
                    cnt = _trino(f"SELECT count(*) FROM iceberg.hadoop.{name}", 1)
                    n = cnt.get("rows", [[0]])[0][0] if cnt.get("ok") else 0
                    _mat_state["done"].append({"table": name, "rows": n})
                else:
                    _mat_state["errors"].append(f"{name}: {str(res.get('error',''))[:160]}")
            except Exception as exc:  # noqa: BLE001
                _mat_state["errors"].append(f"{name}: {str(exc)[:160]}")
    finally:
        _mat_state["current"] = None
        _mat_state["running"] = False
        _mat_state["finished_at"] = datetime.now(timezone.utc).isoformat()


@router.post("/materialize")
async def materialize():
    with _mat_lock:
        if _mat_state["running"]:
            return {"ok": True, "already_running": True, "state": _mat_state}
        _mat_state.update({"running": True, "current": "starting…", "done": [], "errors": [],
                           "started_at": datetime.now(timezone.utc).isoformat(), "finished_at": None})
    threading.Thread(target=_materialize_worker, daemon=True).start()
    return {"ok": True, "started": True, "tables": [s[0] for s in LAKE_SPECS]}


@router.get("/materialize/status")
async def materialize_status():
    return {"ok": True, **_mat_state}


# ──────────────────────────────────────────────────────────────────────────────
# Business data dictionary (columns + masked/visible) — UI + LLM
# ──────────────────────────────────────────────────────────────────────────────
DICT_TABLES = [
    ("postgres_sales", "public", "sales_orders", "PostgreSQL", "Sales orders (OLTP, CDC source)"),
    ("mysql_hr", "hr", "employee_events", "MySQL", "Employee / HR events (CDC source)"),
    ("mongodb_supplychain", "supplychain", "events", "MongoDB", "Supply chain events (CDC source)"),
    ("cassandra_telemetry", "telemetry", "device_metrics", "Cassandra", "IoT device telemetry"),
    ("iceberg", "curated_masked", "sales_orders_masked", "Iceberg", "Curated masked layer (physically masked)"),
    ("iceberg", "hadoop", "orders_ext", "Hadoop", "Lake external table (orders)"),
]

_dict_cache: dict[str, Any] = {"at": 0.0, "data": None}
_DICT_TTL = 120.0


def build_dictionary() -> dict[str, Any]:
    if _dict_cache["data"] and time.time() - _dict_cache["at"] < _DICT_TTL:
        return _dict_cache["data"]
    import sql_console as s
    # masking map from the PII catalog
    masked_map: dict[tuple[str, str], dict] = {}
    pii_names: dict[str, str] = {}  # column-name -> category (any dataset)
    try:
        from pii_catalog import get_pii
        pii_data = get_pii()
        for d in pii_data.get("datasets", []):
            tname = (d.get("table") or "").split(".")[-1]
            for c in d.get("pii_columns", []):
                masked_map[(tname, c["name"])] = {"masked": c.get("masked"), "category": c.get("category")}
                pii_names.setdefault(c["name"], c.get("category"))
    except Exception:
        pass

    tables_out = []
    for catalog, schema, table, engine, desc in DICT_TABLES:
        cols_res = _trino(f'SHOW COLUMNS FROM "{catalog}"."{schema}"."{table}"', 200)
        if not cols_res.get("ok"):
            continue
        masked_layer = schema in ("curated_masked", "curated")
        cols = []
        for r in cols_res.get("rows", []):
            cname, ctype = r[0], r[1]
            direct = masked_map.get((table, cname))
            if direct:
                is_pii, is_masked, cat = True, bool(direct["masked"]), direct["category"]
            elif cname in pii_names:
                # PII column name found in a derived/lake table — masked only if it
                # is a physically-masked layer; otherwise raw PII is visible there.
                is_pii, is_masked, cat = True, masked_layer, pii_names[cname]
            else:
                is_pii, is_masked, cat = False, False, None
            cols.append({
                "name": cname, "type": ctype,
                "pii": is_pii, "masked": is_masked, "category": cat,
            })
        tables_out.append({
            "catalog": catalog, "schema": schema, "table": table, "engine": engine, "desc": desc,
            "fqn": f"{catalog}.{schema}.{table}",
            "columns": cols,
            "pii_count": sum(1 for c in cols if c["pii"]),
            "masked_count": sum(1 for c in cols if c["masked"]),
        })
    data = {
        "ok": True,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "tables": tables_out,
        "summary": {
            "tables": len(tables_out),
            "columns": sum(len(t["columns"]) for t in tables_out),
            "pii_columns": sum(t["pii_count"] for t in tables_out),
            "masked_columns": sum(t["masked_count"] for t in tables_out),
        },
    }
    _dict_cache.update({"at": time.time(), "data": data})
    return data


@router.get("/dictionary")
async def get_dictionary():
    try:
        return build_dictionary()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)[:300]}, status_code=502)