atc-agents/api/elasticsearch_api.py

"""Elasticsearch + Kibana API for the Command Center.

Provides:
  - cluster / kibana health
  - rich index listing, field mappings, full-text + filtered search, aggregations
  - a Trino-federated indexer that pushes every source's data into ES (atc-* indices)
  - Kibana provisioning (data views + an overview dashboard)
"""

from __future__ import annotations

import json
import os
import re
import threading
import time
from datetime import datetime, timezone
from typing import Any

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

ELASTICSEARCH_URL = os.getenv("ELASTICSEARCH_URL", "https://10.0.21.46:9200").rstrip("/")
KIBANA_URL = os.getenv("KIBANA_URL", "http://10.0.21.46:5601").rstrip("/")
ELASTIC_USER = os.getenv("ELASTIC_USER", "elastic")
ELASTIC_PASSWORD = os.getenv("ELASTIC_PASSWORD", "")

INDEX_PREFIX = "atc-"
# Trino catalogs to index (skip system + streaming kafka).
INDEX_CATALOGS = ("postgres_sales", "mysql_hr", "mongodb_supplychain", "cassandra_telemetry", "iceberg")
SKIP_SCHEMAS = {"information_schema", "sys", "performance_schema", "config", "local", "admin"}
ROWS_PER_TABLE = int(os.getenv("ES_ROWS_PER_TABLE", "1000"))
# big source tables we want broadly searchable (row-level)
ROWS_BUSINESS = int(os.getenv("ES_ROWS_BUSINESS", "15000"))
BUSINESS_TABLES = {
    ("postgres_sales", "public", "sales_orders"),
    ("mysql_hr", "hr", "employee_events"),
    ("mongodb_supplychain", "supplychain", "events"),
    ("cassandra_telemetry", "telemetry", "device_metrics"),
}
# prefer rows that actually carry the business identifiers (names are sparse in raw tables)
BUSINESS_WHERE = {
    ("postgres_sales", "public", "sales_orders"): "customer_name IS NOT NULL",
    ("mysql_hr", "hr", "employee_events"): "employee_name IS NOT NULL",
}
# de-duplicated business entities so users can search people/customers by name + id
ENTITY_QUERIES = [
    (
        "atc-customers", "customers", "customer",
        'SELECT DISTINCT customer_id, customer_name, customer_email, region, sales_channel, currency '
        'FROM "postgres_sales"."public"."sales_orders" WHERE customer_id IS NOT NULL LIMIT 50000',
        "customer_id",
    ),
    (
        "atc-employees", "employees", "employee",
        'SELECT DISTINCT employee_id, employee_name, employee_email, employee_phone, department, role_name, region '
        'FROM "mysql_hr"."hr"."employee_events" WHERE employee_id IS NOT NULL LIMIT 50000',
        "employee_id",
    ),
    (
        "atc-products", "products", "product",
        'SELECT DISTINCT product_id, region, sales_channel FROM "postgres_sales"."public"."sales_orders" '
        'WHERE product_id IS NOT NULL LIMIT 50000',
        "product_id",
    ),
]

# business index names (federated source data + entity rollups)
IDX_ORDERS = "atc-postgres-sales-public-sales-orders"
IDX_HR = "atc-mysql-hr-hr-employee-events"
IDX_SUPPLY = "atc-mongodb-supplychain-supplychain-events"
IDX_TELEMETRY = "atc-cassandra-telemetry-telemetry-device-metrics"
IDX_CUSTOMERS = "atc-customers"
IDX_EMPLOYEES = "atc-employees"
IDX_PRODUCTS = "atc-products"

# lenient multi-format date parsing for the timestamp columns in the data
_DATE_FMT = (
    "yyyy-MM-dd HH:mm:ss.SSS zzz||yyyy-MM-dd HH:mm:ss zzz||"
    "yyyy-MM-dd HH:mm:ss.SSS||yyyy-MM-dd HH:mm:ss||yyyy-MM-dd||"
    "strict_date_optional_time||epoch_millis"
)
_NUM = {"type": "double", "ignore_malformed": True}
_DATE = {"type": "date", "format": _DATE_FMT, "ignore_malformed": True}
# index template so measures are numeric (sum/avg) and timestamps are real dates
ATC_TEMPLATE = {
    "index_patterns": [f"{INDEX_PREFIX}*"],
    "priority": 200,
    "template": {
        "settings": {"number_of_replicas": 0, "refresh_interval": "5s"},
        "mappings": {
            "properties": {
                "amount": _NUM, "total_amount": _NUM, "salary_change": _NUM,
                "unit_price": _NUM, "metric_value": _NUM, "quantity": _NUM,
                "revenue": _NUM, "order_count": _NUM,
                "order_ts": _DATE, "event_ts": _DATE, "metric_ts": _DATE,
                "order_date": _DATE, "ts": _DATE,
            }
        },
    },
}

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


def _auth():
    return (ELASTIC_USER, ELASTIC_PASSWORD) if ELASTIC_PASSWORD else None


def _es_client(timeout: float = 30.0) -> httpx.Client:
    return httpx.Client(timeout=timeout, verify=False, auth=_auth())


def _sanitize_index(*parts: str) -> str:
    raw = "-".join(str(p) for p in parts if p)
    raw = raw.lower()
    raw = re.sub(r"[^a-z0-9]+", "-", raw).strip("-")
    raw = re.sub(r"-+", "-", raw)
    return (INDEX_PREFIX + raw)[:240]


# ──────────────────────────────────────────────────────────────────────────────
# Health
# ──────────────────────────────────────────────────────────────────────────────
async def _probe_es() -> dict[str, Any]:
    try:
        async with httpx.AsyncClient(timeout=6.0, verify=False) as client:
            r = await client.get(f"{ELASTICSEARCH_URL}/", auth=_auth())
            if r.status_code >= 400:
                return {"ok": False, "status_code": r.status_code, "error": r.text[:200]}
            info = r.json()
            health_r = await client.get(f"{ELASTICSEARCH_URL}/_cluster/health", auth=_auth())
            health = health_r.json() if health_r.status_code < 400 else {}
            stats_r = await client.get(f"{ELASTICSEARCH_URL}/_cat/indices?format=json&bytes=b", auth=_auth())
            indices = stats_r.json() if stats_r.status_code < 400 else []
            user_idx = [i for i in indices if isinstance(i, dict) and not str(i.get("index", "")).startswith(".")]
            total_docs = sum(int(i.get("docs.count") or 0) for i in user_idx)
            atc_docs = sum(int(i.get("docs.count") or 0) for i in user_idx if str(i.get("index", "")).startswith(INDEX_PREFIX))
            return {
                "ok": True,
                "url": ELASTICSEARCH_URL,
                "cluster_name": info.get("cluster_name"),
                "version": info.get("version", {}).get("number"),
                "health": health.get("status", "unknown"),
                "nodes": health.get("number_of_nodes"),
                "indices_count": len(user_idx),
                "total_docs": total_docs,
                "atc_docs": atc_docs,
                "indices": [
                    {
                        "name": i.get("index"),
                        "docs": i.get("docs.count"),
                        "size": i.get("store.size"),
                        "health": i.get("health"),
                    }
                    for i in user_idx[:80]
                ],
            }
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "url": ELASTICSEARCH_URL, "error": str(exc)}


async def _probe_kibana() -> dict[str, Any]:
    try:
        async with httpx.AsyncClient(timeout=6.0, verify=False) as client:
            r = await client.get(f"{KIBANA_URL}/api/status")
            if r.status_code >= 400:
                return {"ok": False, "url": KIBANA_URL, "ui_url": KIBANA_URL, "status_code": r.status_code}
            data = r.json()
            overall = data.get("status", {}).get("overall", {})
            return {
                "ok": overall.get("level", "available") in ("available", "green", "yellow"),
                "url": KIBANA_URL,
                "ui_url": KIBANA_URL,
                "level": overall.get("level", "available"),
                "version": data.get("version", {}).get("number"),
            }
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "url": KIBANA_URL, "ui_url": KIBANA_URL, "error": str(exc)}


@router.get("/health")
async def search_health():
    es, kb = await _probe_es(), await _probe_kibana()
    return {"ok": es.get("ok") or kb.get("ok"), "elasticsearch": es, "kibana": kb}


@router.get("/elasticsearch")
async def get_elasticsearch():
    result = await _probe_es()
    if not result.get("ok"):
        return JSONResponse(result, status_code=502)
    return result


@router.get("/kibana")
async def get_kibana():
    return await _probe_kibana()


# ──────────────────────────────────────────────────────────────────────────────
# Indices + mappings
# ──────────────────────────────────────────────────────────────────────────────
@router.get("/indices")
async def list_indices(include_system: bool = False):
    try:
        async with httpx.AsyncClient(timeout=10.0, verify=False) as client:
            r = await client.get(f"{ELASTICSEARCH_URL}/_cat/indices?format=json&bytes=b&s=index", auth=_auth())
            if r.status_code >= 400:
                return JSONResponse({"ok": False, "error": r.text[:300]}, status_code=502)
            rows = r.json()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)

    out = []
    for i in rows:
        name = str(i.get("index", ""))
        if not include_system and name.startswith("."):
            continue
        out.append({
            "name": name,
            "docs": int(i.get("docs.count") or 0),
            "size_bytes": int(i.get("store.size") or 0),
            "health": i.get("health"),
            "status": i.get("status"),
            "atc": name.startswith(INDEX_PREFIX),
        })
    out.sort(key=lambda x: (not x["atc"], -x["docs"]))
    return {"ok": True, "count": len(out), "indices": out}


def _flatten_props(props: dict, prefix: str = "") -> list[dict]:
    fields = []
    for name, spec in (props or {}).items():
        full = f"{prefix}{name}"
        ftype = spec.get("type")
        if ftype:
            aggregatable = ftype in ("keyword", "long", "integer", "double", "float", "date", "boolean", "ip")
            fields.append({"name": full, "type": ftype, "aggregatable": aggregatable})
        sub = spec.get("fields") or {}
        for sub_name, sub_spec in sub.items():
            st = sub_spec.get("type")
            fields.append({"name": f"{full}.{sub_name}", "type": st, "aggregatable": st in ("keyword", "ip")})
        if spec.get("properties"):
            fields.extend(_flatten_props(spec["properties"], prefix=f"{full}."))
    return fields


@router.get("/mapping")
async def get_mapping(index: str = Query(...)):
    try:
        async with httpx.AsyncClient(timeout=10.0, verify=False) as client:
            r = await client.get(f"{ELASTICSEARCH_URL}/{index}/_mapping", auth=_auth())
            if r.status_code >= 400:
                return JSONResponse({"ok": False, "error": r.text[:300]}, status_code=502)
            data = r.json()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)

    fields: dict[str, dict] = {}
    for _idx, body in data.items():
        props = body.get("mappings", {}).get("properties", {})
        for f in _flatten_props(props):
            fields[f["name"]] = f
    field_list = sorted(fields.values(), key=lambda x: x["name"])
    return {"ok": True, "index": index, "fields": field_list}


# ──────────────────────────────────────────────────────────────────────────────
# Search + aggregations
# ──────────────────────────────────────────────────────────────────────────────
@router.post("/query")
async def search_query(body: dict = Body(default={})):
    q = (body.get("q") or "").strip()
    indices = body.get("indices") or []
    filters = body.get("filters") or []
    from_ = int(body.get("from") or 0)
    size = min(int(body.get("size") or 20), 100)
    sort = body.get("sort")

    target = ",".join(indices) if indices else f"{INDEX_PREFIX}*"

    must: list[dict] = []
    if q:
        must.append({"query_string": {"query": q, "lenient": True, "default_operator": "AND"}})
    for f in filters:
        field, value = f.get("field"), f.get("value")
        if field and value is not None:
            must.append({"match_phrase": {field: value}})
    es_body: dict[str, Any] = {
        "query": {"bool": {"must": must or [{"match_all": {}}]}},
        "from": from_,
        "size": size,
        "track_total_hits": True,
    }
    if sort and sort.get("field"):
        es_body["sort"] = [{sort["field"]: {"order": sort.get("order", "desc")}}]

    try:
        async with httpx.AsyncClient(timeout=20.0, verify=False) as client:
            r = await client.post(f"{ELASTICSEARCH_URL}/{target}/_search", auth=_auth(), json=es_body)
            if r.status_code >= 400:
                return JSONResponse({"ok": False, "error": r.text[:400]}, status_code=r.status_code)
            data = r.json()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)

    hits = data.get("hits", {})
    return {
        "ok": True,
        "took": data.get("took"),
        "total": hits.get("total", {}).get("value", 0),
        "hits": [
            {"index": h.get("_index"), "id": h.get("_id"), "score": h.get("_score"), "source": h.get("_source")}
            for h in hits.get("hits", [])
        ],
    }


@router.get("/aggs")
async def aggregations(index: str = Query(...), field: str = Query(...), size: int = Query(15, le=50), q: str = Query("")):
    query = {"query_string": {"query": q, "lenient": True}} if q.strip() else {"match_all": {}}
    es_body = {"size": 0, "query": query, "aggs": {"facet": {"terms": {"field": field, "size": size}}}}
    try:
        async with httpx.AsyncClient(timeout=15.0, verify=False) as client:
            r = await client.post(f"{ELASTICSEARCH_URL}/{index}/_search", auth=_auth(), json=es_body)
            if r.status_code >= 400:
                return JSONResponse({"ok": False, "error": r.text[:300]}, status_code=r.status_code)
            data = r.json()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)
    buckets = data.get("aggregations", {}).get("facet", {}).get("buckets", [])
    return {"ok": True, "field": field, "buckets": [{"key": b.get("key"), "count": b.get("doc_count")} for b in buckets]}


# back-compat simple query
@router.get("/elasticsearch/query")
async def es_search(q: str = Query(..., min_length=1), size: int = Query(10, le=50)):
    return await search_query({"q": q, "size": size})


# ──────────────────────────────────────────────────────────────────────────────
# Business data overview (data abstraction — charts & filters over the real data)
# ──────────────────────────────────────────────────────────────────────────────
_biz_cache: dict[str, Any] = {}
_BIZ_TTL = 30.0
_totals_cache: dict[str, Any] = {"at": 0.0, "data": {}}
_TOTALS_TTL = 600.0


def _buckets(agg: dict, key: str, metric: str | None = None) -> list[dict]:
    out = []
    for b in (agg.get(key, {}) or {}).get("buckets", []):
        item = {"key": b.get("key_as_string", b.get("key")), "count": b.get("doc_count", 0)}
        if metric and metric in b:
            item["value"] = round(b[metric].get("value") or 0, 2)
        out.append(item)
    return out


async def _es_aggs(client: httpx.AsyncClient, index: str, filt: dict, aggs: dict) -> dict:
    try:
        r = await client.post(
            f"{ELASTICSEARCH_URL}/{index}/_search",
            params={"ignore_unavailable": "true", "allow_no_indices": "true"},
            auth=_auth(),
            json={"size": 0, "query": filt, "aggs": aggs, "track_total_hits": True},
        )
        if r.status_code >= 400:
            return {}
        return r.json()
    except Exception:  # noqa: BLE001
        return {}


def _source_totals() -> dict:
    """True row counts in the source systems via cheap planner estimates (cached)."""
    if time.time() - _totals_cache["at"] < _TOTALS_TTL and _totals_cache["data"]:
        return _totals_cache["data"]
    out: dict[str, Any] = {}
    try:
        import sql_console as s
        out["orders"] = s._table_row_count("postgres", "public.sales_orders")
        out["hr_events"] = s._table_row_count("mysql", "hr.employee_events")
        out["supply_events"] = s._table_row_count("mongodb", "supplychain.events")
    except Exception:  # noqa: BLE001
        pass
    _totals_cache.update({"at": time.time(), "data": out})
    return out


@router.get("/business")
async def business_overview(q: str = Query(""), region: str = Query("")):
    cache_key = f"{q}|{region}"
    cached = _biz_cache.get(cache_key)
    if cached and time.time() - cached["at"] < _BIZ_TTL:
        return cached["data"]

    must: list[dict] = []
    if q.strip():
        must.append({"query_string": {"query": q, "lenient": True, "default_operator": "AND"}})
    region_must = must + ([{"term": {"region.keyword": region}}] if region.strip() else [])
    filt = {"bool": {"must": must or [{"match_all": {}}]}}
    rfilt = {"bool": {"must": region_must or [{"match_all": {}}]}}

    async with httpx.AsyncClient(timeout=20.0, verify=False) as client:
        orders = await _es_aggs(client, IDX_ORDERS, rfilt, {
            "by_region": {"terms": {"field": "region.keyword", "size": 12}, "aggs": {"rev": {"sum": {"field": "amount"}}}},
            "by_channel": {"terms": {"field": "sales_channel.keyword", "size": 12}, "aggs": {"rev": {"sum": {"field": "amount"}}}},
            "by_status": {"terms": {"field": "order_status.keyword", "size": 12}},
            "top_customers": {"terms": {"field": "customer_name.keyword", "size": 10}, "aggs": {"spend": {"sum": {"field": "amount"}}}},
            "top_products": {"terms": {"field": "product_id", "size": 10}, "aggs": {"rev": {"sum": {"field": "amount"}}}},
            "over_time": {"date_histogram": {"field": "order_ts", "calendar_interval": "month", "min_doc_count": 0}, "aggs": {"rev": {"sum": {"field": "amount"}}}},
            "revenue": {"sum": {"field": "amount"}},
            "avg_order": {"avg": {"field": "amount"}},
        })
        hr = await _es_aggs(client, IDX_HR, rfilt, {
            "by_department": {"terms": {"field": "department.keyword", "size": 15}},
            "by_role": {"terms": {"field": "role_name.keyword", "size": 12}},
            "by_event": {"terms": {"field": "event_type.keyword", "size": 12}},
        })
        supply = await _es_aggs(client, IDX_SUPPLY, rfilt, {
            "by_type": {"terms": {"field": "type.keyword", "size": 15}},
            "by_region": {"terms": {"field": "region.keyword", "size": 12}},
        })
        telemetry = await _es_aggs(client, IDX_TELEMETRY, filt, {
            "by_metric": {"terms": {"field": "metric_type.keyword", "size": 12}, "aggs": {"avg": {"avg": {"field": "metric_value"}}}},
        })

        async def _count(index: str, body_filt: dict) -> int:
            try:
                r = await client.post(
                    f"{ELASTICSEARCH_URL}/{index}/_count",
                    params={"ignore_unavailable": "true", "allow_no_indices": "true"},
                    auth=_auth(), json={"query": body_filt},
                )
                return int(r.json().get("count", 0)) if r.status_code < 400 else 0
            except Exception:  # noqa: BLE001
                return 0

        customers = await _count(IDX_CUSTOMERS, rfilt)
        employees = await _count(IDX_EMPLOYEES, rfilt)
        products = await _count(IDX_PRODUCTS, rfilt)

    o_agg = orders.get("aggregations", {})
    hr_agg = hr.get("aggregations", {})
    sup_agg = supply.get("aggregations", {})
    tel_agg = telemetry.get("aggregations", {})
    orders_indexed = orders.get("hits", {}).get("total", {}).get("value", 0)

    data = {
        "ok": True,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "filters": {"q": q, "region": region},
        "kpis": {
            "customers": customers,
            "employees": employees,
            "products": products,
            "orders_indexed": orders_indexed,
            "revenue_indexed": round((o_agg.get("revenue", {}) or {}).get("value") or 0, 2),
            "avg_order": round((o_agg.get("avg_order", {}) or {}).get("value") or 0, 2),
            "telemetry_indexed": telemetry.get("hits", {}).get("total", {}).get("value", 0),
            "supply_indexed": supply.get("hits", {}).get("total", {}).get("value", 0),
            "source_totals": _source_totals(),
        },
        "orders": {
            "by_region": _buckets(o_agg, "by_region", "rev"),
            "by_channel": _buckets(o_agg, "by_channel", "rev"),
            "by_status": _buckets(o_agg, "by_status"),
            "top_customers": _buckets(o_agg, "top_customers", "spend"),
            "top_products": _buckets(o_agg, "top_products", "rev"),
            "over_time": _buckets(o_agg, "over_time", "rev"),
        },
        "hr": {
            "by_department": _buckets(hr_agg, "by_department"),
            "by_role": _buckets(hr_agg, "by_role"),
            "by_event": _buckets(hr_agg, "by_event"),
        },
        "supply": {
            "by_type": _buckets(sup_agg, "by_type"),
            "by_region": _buckets(sup_agg, "by_region"),
        },
        "telemetry": {"by_metric": _buckets(tel_agg, "by_metric", "avg")},
        "regions": [b["key"] for b in _buckets(o_agg, "by_region")],
    }
    _biz_cache[cache_key] = {"at": time.time(), "data": data}
    return data


# ──────────────────────────────────────────────────────────────────────────────
# Indexer (Trino-federated → ES atc-* indices)
# ──────────────────────────────────────────────────────────────────────────────
_reindex_state: dict[str, Any] = {
    "running": False,
    "started_at": None,
    "finished_at": None,
    "current": None,
    "indices": {},
    "total_docs": 0,
    "errors": [],
    "log": [],
}
_reindex_lock = threading.Lock()


def _bulk(client: httpx.Client, index: str, docs: list[tuple[str, dict]]) -> int:
    if not docs:
        return 0
    lines = []
    for _id, doc in docs:
        lines.append(json.dumps({"index": {"_index": index, "_id": _id}}))
        lines.append(json.dumps(doc, default=str))
    body = "\n".join(lines) + "\n"
    r = client.post(
        f"{ELASTICSEARCH_URL}/_bulk",
        content=body.encode("utf-8"),
        headers={"Content-Type": "application/x-ndjson"},
    )
    if r.status_code >= 400:
        raise RuntimeError(f"bulk failed {r.status_code}: {r.text[:200]}")
    return len(docs)


def _meta(catalog: str, schema: str, table: str, source: str) -> dict:
    return {
        "catalog": catalog,
        "schema": schema,
        "table": table,
        "fqn": f"{catalog}.{schema}.{table}" if schema else f"{catalog}.{table}",
        "source": source,
        "indexed_at": datetime.now(timezone.utc).isoformat(),
    }


def _index_trino_table(sqlmod, client: httpx.Client, catalog: str, schema: str, table: str, limit: int = ROWS_PER_TABLE, where: str | None = None) -> int:
    where_sql = f" WHERE {where}" if where else ""
    res = sqlmod._run_trino(f'SELECT * FROM "{catalog}"."{schema}"."{table}"{where_sql} LIMIT {limit}', limit)
    if not res.get("ok"):
        raise RuntimeError(res.get("error", "query failed")[:200])
    cols = res.get("columns", [])
    rows = res.get("rows", [])
    index = _sanitize_index(catalog, schema, table)
    docs = []
    for n, row in enumerate(rows):
        doc = {cols[i]: row[i] for i in range(min(len(cols), len(row)))}
        doc["meta"] = _meta(catalog, schema, table, catalog)
        docs.append((f"{catalog}.{schema}.{table}:{n}", doc))
    written = 0
    for i in range(0, len(docs), 500):
        written += _bulk(client, index, docs[i:i + 500])
    return written


def _index_entities(sqlmod, client: httpx.Client) -> int:
    """Build de-duplicated entity indices (customers, employees, products) so users
    can search by name / id directly instead of scanning raw transaction rows."""
    written = 0
    for index, catalog_label, schema_label, sql, id_field in ENTITY_QUERIES:
        _reindex_state["current"] = f"entities · {catalog_label}"
        try:
            res = sqlmod._run_trino(sql, 50000)
            if not res.get("ok"):
                _reindex_state["errors"].append(f"{index}: {res.get('error', 'failed')[:120]}")
                continue
            cols = res.get("columns", [])
            rows = res.get("rows", [])
            docs = []
            for row in rows:
                doc = {cols[i]: row[i] for i in range(min(len(cols), len(row)))}
                ident = doc.get(id_field)
                doc["meta"] = _meta(catalog_label, schema_label, id_field, catalog_label)
                docs.append((f"{schema_label}:{ident}", doc))
            n = 0
            for i in range(0, len(docs), 500):
                n += _bulk(client, index, docs[i:i + 500])
            _reindex_state["indices"][index] = n
            _reindex_state["total_docs"] += n
            _reindex_state["log"].append(f"{index} → {n} entities")
            written += n
        except Exception as exc:  # noqa: BLE001
            _reindex_state["errors"].append(f"{index}: {str(exc)[:120]}")
    return written


def _index_neo4j(sqlmod, client: httpx.Client) -> int:
    written = 0
    try:
        driver = sqlmod._neo4j_driver()
    except Exception:  # noqa: BLE001
        return 0
    try:
        with driver.session() as session:
            labels = [r["label"] for r in session.run("CALL db.labels() YIELD label RETURN label")]
            for label in labels:
                index = _sanitize_index("neo4j", "graph", label)
                recs = session.run(f"MATCH (n:`{label}`) RETURN properties(n) AS props LIMIT {ROWS_PER_TABLE}")
                docs = []
                for n, rec in enumerate(recs):
                    props = dict(rec["props"] or {})
                    props["meta"] = _meta("neo4j", "graph", label, "neo4j")
                    docs.append((f"neo4j.{label}:{n}", props))
                for i in range(0, len(docs), 500):
                    written += _bulk(client, index, docs[i:i + 500])
                _reindex_state["indices"][index] = len(docs)
                _reindex_state["log"].append(f"neo4j:{label} → {len(docs)} docs")
    except Exception as exc:  # noqa: BLE001
        _reindex_state["errors"].append(f"neo4j: {exc}")
    finally:
        driver.close()
    return written


def _index_catalog(client: httpx.Client) -> int:
    try:
        from database_inventory import collect_database_inventory_sync
        inv = collect_database_inventory_sync()
    except Exception as exc:  # noqa: BLE001
        _reindex_state["errors"].append(f"catalog: {exc}")
        return 0
    docs = []
    for engine, info in (inv.get("engines") or {}).items():
        for obj in (info.get("objects") or info.get("tables") or []):
            if not isinstance(obj, dict):
                continue
            doc = dict(obj)
            doc["engine"] = engine
            doc["meta"] = _meta("catalog", engine, str(obj.get("name") or obj.get("fqn") or "obj"), "inventory")
            docs.append((f"catalog.{engine}.{obj.get('fqn') or obj.get('name')}", doc))
    written = 0
    for i in range(0, len(docs), 500):
        written += _bulk(client, "atc-catalog", docs[i:i + 500])
    if docs:
        _reindex_state["indices"]["atc-catalog"] = len(docs)
        _reindex_state["log"].append(f"catalog → {len(docs)} docs")
    return written


def _ensure_template(client: httpx.Client) -> None:
    try:
        client.put(f"{ELASTICSEARCH_URL}/_index_template/atc-data", json=ATC_TEMPLATE)
    except Exception as exc:  # noqa: BLE001
        _reindex_state["errors"].append(f"template: {str(exc)[:120]}")


def _drop_atc_indices(client: httpx.Client) -> None:
    """Delete existing atc-* indices so they are re-created with the new
    numeric/date mappings from the template (mappings can't be changed in place)."""
    try:
        r = client.get(f"{ELASTICSEARCH_URL}/_cat/indices/{INDEX_PREFIX}*?format=json&h=index")
        for it in (r.json() if r.status_code < 400 else []):
            nm = it.get("index")
            if nm:
                try:
                    client.delete(f"{ELASTICSEARCH_URL}/{nm}")
                except Exception:  # noqa: BLE001
                    pass
    except Exception as exc:  # noqa: BLE001
        _reindex_state["errors"].append(f"cleanup: {str(exc)[:120]}")


def _reindex_worker():
    import sql_console as sqlmod
    try:
        client = _es_client(timeout=90.0)
        with client:
            # numeric/date mappings + fresh indices
            _reindex_state["current"] = "preparing mappings"
            _ensure_template(client)
            _drop_atc_indices(client)

            # catalog/metadata
            _reindex_state["current"] = "catalog inventory"
            _reindex_state["total_docs"] += _index_catalog(client)

            # business entities (customers / employees / products) — the headline searchables
            _index_entities(sqlmod, client)

            # Trino-federated source data
            cat_res = sqlmod._run_trino("SHOW CATALOGS", 100)
            catalogs = [r[0] for r in cat_res.get("rows", [])] if cat_res.get("ok") else list(INDEX_CATALOGS)
            for catalog in catalogs:
                if catalog not in INDEX_CATALOGS:
                    continue
                try:
                    sch_res = sqlmod._run_trino(f'SHOW SCHEMAS FROM "{catalog}"', 200)
                    schemas = [r[0] for r in sch_res.get("rows", [])] if sch_res.get("ok") else []
                except Exception as exc:  # noqa: BLE001
                    _reindex_state["errors"].append(f"{catalog}: {exc}")
                    continue
                for schema in schemas:
                    if schema in SKIP_SCHEMAS or schema.startswith("system") or schema == "information_schema":
                        continue
                    try:
                        tbl_res = sqlmod._run_trino(f'SHOW TABLES FROM "{catalog}"."{schema}"', 500)
                        tables = [r[0] for r in tbl_res.get("rows", [])] if tbl_res.get("ok") else []
                    except Exception as exc:  # noqa: BLE001
                        _reindex_state["errors"].append(f"{catalog}.{schema}: {exc}")
                        continue
                    for table in tables:
                        _reindex_state["current"] = f"{catalog}.{schema}.{table}"
                        try:
                            cap = ROWS_BUSINESS if (catalog, schema, table) in BUSINESS_TABLES or catalog == "iceberg" else ROWS_PER_TABLE
                            where = BUSINESS_WHERE.get((catalog, schema, table))
                            n = _index_trino_table(sqlmod, client, catalog, schema, table, cap, where)
                            idx = _sanitize_index(catalog, schema, table)
                            _reindex_state["indices"][idx] = n
                            _reindex_state["total_docs"] += n
                            _reindex_state["log"].append(f"{catalog}.{schema}.{table} → {n} docs")
                        except Exception as exc:  # noqa: BLE001
                            _reindex_state["errors"].append(f"{catalog}.{schema}.{table}: {str(exc)[:150]}")

            # Neo4j (not federated by Trino)
            _reindex_state["current"] = "neo4j graph"
            _reindex_state["total_docs"] += _index_neo4j(sqlmod, client)

            # make freshly created indices searchable immediately
            try:
                client.post(f"{ELASTICSEARCH_URL}/{INDEX_PREFIX}*/_refresh")
            except Exception:  # noqa: BLE001
                pass
    except Exception as exc:  # noqa: BLE001
        _reindex_state["errors"].append(f"fatal: {exc}")
    finally:
        _reindex_state["current"] = None
        _reindex_state["running"] = False
        _reindex_state["finished_at"] = datetime.now(timezone.utc).isoformat()
        # keep log bounded
        _reindex_state["log"] = _reindex_state["log"][-200:]


@router.post("/reindex")
async def start_reindex():
    with _reindex_lock:
        if _reindex_state["running"]:
            return {"ok": True, "already_running": True, "state": _reindex_summary()}
        _reindex_state.update({
            "running": True,
            "started_at": datetime.now(timezone.utc).isoformat(),
            "finished_at": None,
            "current": "starting…",
            "indices": {},
            "total_docs": 0,
            "errors": [],
            "log": [],
        })
    threading.Thread(target=_reindex_worker, daemon=True).start()
    return {"ok": True, "started": True}


def _reindex_summary() -> dict[str, Any]:
    s = _reindex_state
    return {
        "running": s["running"],
        "started_at": s["started_at"],
        "finished_at": s["finished_at"],
        "current": s["current"],
        "total_docs": s["total_docs"],
        "index_count": len(s["indices"]),
        "indices": s["indices"],
        "errors": s["errors"][-20:],
        "log": s["log"][-40:],
    }


@router.get("/reindex/status")
async def reindex_status():
    return {"ok": True, **_reindex_summary()}


# ──────────────────────────────────────────────────────────────────────────────
# Kibana provisioning
# ──────────────────────────────────────────────────────────────────────────────
def _kbn_headers():
    return {"kbn-xsrf": "true", "Content-Type": "application/json"}


@router.post("/kibana/setup")
async def kibana_setup():
    """Create data views + an overview dashboard in Kibana (idempotent)."""
    created: list[str] = []
    errors: list[str] = []
    dataview_id = "atc-all-data"
    objects = [
        {
            "type": "index-pattern",
            "id": dataview_id,
            "attributes": {"title": f"{INDEX_PREFIX}*", "name": "ATC — all data", "timeFieldName": "meta.indexed_at"},
        },
        {
            "type": "index-pattern",
            "id": "atc-catalog-dv",
            "attributes": {"title": "atc-catalog", "name": "ATC — data catalog"},
        },
        {
            "type": "visualization",
            "id": "atc-docs-by-source",
            "attributes": {
                "title": "ATC — Documents by source",
                "visState": json.dumps({
                    "title": "ATC — Documents by source",
                    "type": "histogram",
                    "aggs": [
                        {"id": "1", "enabled": True, "type": "count", "schema": "metric", "params": {}},
                        {"id": "2", "enabled": True, "type": "terms", "schema": "segment",
                         "params": {"field": "meta.catalog.keyword", "size": 25, "order": "desc", "orderBy": "1"}},
                    ],
                    "params": {"addLegend": True, "addTooltip": True, "type": "histogram"},
                }),
                "uiStateJSON": "{}",
                "description": "",
                "kibanaSavedObjectMeta": {"searchSourceJSON": json.dumps({
                    "query": {"query": "", "language": "kuery"}, "filter": [],
                    "indexRefName": "kibanaSavedObjectMeta.searchSourceJSON.index",
                })},
            },
            "references": [{"name": "kibanaSavedObjectMeta.searchSourceJSON.index", "type": "index-pattern", "id": dataview_id}],
        },
        {
            "type": "visualization",
            "id": "atc-docs-by-table",
            "attributes": {
                "title": "ATC — Top tables",
                "visState": json.dumps({
                    "title": "ATC — Top tables",
                    "type": "table",
                    "aggs": [
                        {"id": "1", "enabled": True, "type": "count", "schema": "metric", "params": {}},
                        {"id": "2", "enabled": True, "type": "terms", "schema": "bucket",
                         "params": {"field": "meta.table.keyword", "size": 30, "order": "desc", "orderBy": "1"}},
                    ],
                    "params": {"perPage": 15, "showPartialRows": False, "showTotal": True},
                }),
                "uiStateJSON": "{}",
                "description": "",
                "kibanaSavedObjectMeta": {"searchSourceJSON": json.dumps({
                    "query": {"query": "", "language": "kuery"}, "filter": [],
                    "indexRefName": "kibanaSavedObjectMeta.searchSourceJSON.index",
                })},
            },
            "references": [{"name": "kibanaSavedObjectMeta.searchSourceJSON.index", "type": "index-pattern", "id": dataview_id}],
        },
        {
            "type": "dashboard",
            "id": "atc-data-overview",
            "attributes": {
                "title": "ATC Data Overview",
                "description": "Auto-provisioned overview of all data indexed from the ATC platform.",
                "panelsJSON": json.dumps([
                    {"version": "8.19.0", "type": "visualization", "panelIndex": "1", "panelRefName": "panel_1",
                     "gridData": {"x": 0, "y": 0, "w": 24, "h": 15, "i": "1"}, "embeddableConfig": {}},
                    {"version": "8.19.0", "type": "visualization", "panelIndex": "2", "panelRefName": "panel_2",
                     "gridData": {"x": 24, "y": 0, "w": 24, "h": 15, "i": "2"}, "embeddableConfig": {}},
                ]),
                "optionsJSON": json.dumps({"useMargins": True, "hidePanelTitles": False}),
                "timeRestore": False,
                "kibanaSavedObjectMeta": {"searchSourceJSON": json.dumps({"query": {"query": "", "language": "kuery"}, "filter": []})},
            },
            "references": [
                {"name": "panel_1", "type": "visualization", "id": "atc-docs-by-source"},
                {"name": "panel_2", "type": "visualization", "id": "atc-docs-by-table"},
            ],
        },
    ]

    try:
        async with httpx.AsyncClient(timeout=20.0, verify=False, auth=_auth()) as client:
            r = await client.post(
                f"{KIBANA_URL}/api/saved_objects/_bulk_create?overwrite=true",
                headers=_kbn_headers(),
                json=objects,
            )
            if r.status_code >= 400:
                return JSONResponse({"ok": False, "error": r.text[:400]}, status_code=r.status_code)
            data = r.json()
            for obj in data.get("saved_objects", []):
                if obj.get("error"):
                    errors.append(f"{obj.get('type')}/{obj.get('id')}: {obj['error'].get('message')}")
                else:
                    created.append(f"{obj.get('type')}/{obj.get('id')}")
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)

    return {
        "ok": True,
        "created": created,
        "errors": errors,
        "dashboard_url": f"{KIBANA_URL}/app/dashboards#/view/atc-data-overview",
        "discover_url": f"{KIBANA_URL}/app/discover",
    }


@router.get("/kibana/links")
async def kibana_links(index: str = Query("")):
    base = KIBANA_URL
    return {
        "ok": True,
        "kibana": base,
        "discover": f"{base}/app/discover",
        "dashboard": f"{base}/app/dashboards#/view/atc-data-overview",
        "index_management": f"{base}/app/management/data/index_management/indices",
    }