atc-agents/api/lab_context.py

"""Live lab metrics for all ATC domains — fed to vLLM as context."""

from __future__ import annotations

import asyncio
import json
import os
import time
from datetime import datetime, timezone
from typing import Any

import httpx

from agent_terminal import TerminalLogFn
from dockhand_envs import DOCKHAND_ENV_COMMAND_CENTER, DOCKHAND_ENVS
from database_inventory import collect_database_inventory
from node_registry import NODE_REGISTRY

DOCKHAND_URL = os.getenv("DOCKHAND_URL", "http://10.0.21.45:8082")
DOCKHAND_API_TOKEN = os.getenv("DOCKHAND_API_TOKEN", "")


def _dockhand_headers() -> dict[str, str]:
    if DOCKHAND_API_TOKEN:
        return {"Authorization": f"Bearer {DOCKHAND_API_TOKEN}"}
    return {}

HDFS_NN_URL = os.getenv("HDFS_NN_URL", "http://10.0.21.61:9870")
LAKEHOUSE_HOST = os.getenv("LAKEHOUSE_HOST", "10.0.21.50")
AIRFLOW_URL = os.getenv("AIRFLOW_URL", "http://10.0.21.55:8080")
KAFKA_UI_URL = os.getenv("KAFKA_UI_URL", "http://10.0.21.36:9000")
KAFKA_CONNECT_URL = os.getenv("KAFKA_CONNECT_URL", f"http://{LAKEHOUSE_HOST}:8083")
TRINO_URL = os.getenv("TRINO_URL", f"http://{LAKEHOUSE_HOST}:8089")
SPARK_UI_URL = os.getenv("SPARK_UI_URL", f"http://{LAKEHOUSE_HOST}:8080")
try:
    from gpu_config import get_gpu_urls as _get_gpu_urls
except Exception:
    _get_gpu_urls = None  # type: ignore

GPU_URL = os.getenv("GPU_URL", "http://10.0.10.106:9000")
OBJECTSCALE_URL = os.getenv("OBJECTSCALE_URL", "http://10.0.20.111:9020")
YARN_URL = os.getenv("YARN_URL", "http://10.0.21.62:8088")

AGENT_PRIMARY_DOMAIN = {
    "infra-sentinel": "docker",
    "data-custodian": "databases",
    "lakehouse-ops": "lakehouse",
    "hadoop-ranger": "hadoop",
    "etl-guardian": "etl",
}


async def _log(log: TerminalLogFn | None, level: str, phase: str, text: str) -> None:
    if log:
        await log(level, phase, text)


async def _get_json(
    client: httpx.AsyncClient,
    url: str,
    log: TerminalLogFn | None = None,
    label: str = "",
    timeout: float = 6.0,
) -> Any | None:
    name = label or url
    t0 = time.monotonic()
    await _log(log, "cmd", "fetch", f"$ GET {url}")
    try:
        r = await client.get(url, timeout=timeout)
        ms = int((time.monotonic() - t0) * 1000)
        if r.status_code < 400:
            await _log(log, "ok", "fetch", f"← {r.status_code} {name} ({ms}ms)")
            return r.json()
        await _log(log, "warn", "fetch", f"← {r.status_code} {name} ({ms}ms)")
    except Exception as exc:
        ms = int((time.monotonic() - t0) * 1000)
        await _log(log, "err", "fetch", f"✗ {name}: {exc} ({ms}ms)")
    return None


async def _probe_ok(
    client: httpx.AsyncClient,
    url: str,
    log: TerminalLogFn | None = None,
    label: str = "",
) -> bool:
    name = label or url
    t0 = time.monotonic()
    await _log(log, "cmd", "probe", f"$ GET {url}")
    try:
        r = await client.get(url, timeout=4.0)
        ms = int((time.monotonic() - t0) * 1000)
        ok = r.status_code < 500
        await _log(log, "ok" if ok else "warn", "probe", f"← {r.status_code} {name} ({'UP' if ok else 'DOWN'}, {ms}ms)")
        return ok
    except Exception as exc:
        ms = int((time.monotonic() - t0) * 1000)
        await _log(log, "err", "probe", f"✗ {name}: {exc} ({ms}ms)")
        return False


def _container_rows(containers: list[dict], host: str = "") -> list[dict[str, Any]]:
    rows = []
    for c in containers:
        ports = sorted({str(p.get("PublicPort")) for p in c.get("ports", []) if p.get("PublicPort")})
        rows.append({
            "name": c.get("name"),
            "state": c.get("state"),
            "image": c.get("image"),
            "status": c.get("status"),
            "ports": ports,
            "host": host,
        })
    return rows


async def dockhand_containers(
    client: httpx.AsyncClient,
    env_id: int,
    log: TerminalLogFn | None = None,
) -> list[dict]:
    url = f"{DOCKHAND_URL}/api/containers?env={env_id}"
    await _log(log, "cmd", "fetch", f"$ GET {url}")
    t0 = time.monotonic()
    try:
        r = await client.get(f"{DOCKHAND_URL}/api/containers", params={"env": env_id}, headers=_dockhand_headers(), timeout=8.0)
        ms = int((time.monotonic() - t0) * 1000)
        r.raise_for_status()
        data = r.json()
        await _log(log, "ok", "fetch", f"← Dockhand env {env_id}: {len(data)} containers ({ms}ms)")
        return data
    except Exception as exc:
        ms = int((time.monotonic() - t0) * 1000)
        await _log(log, "err", "fetch", f"✗ Dockhand env {env_id}: {exc} ({ms}ms)")
        return []


async def collect_hdfs(client: httpx.AsyncClient, log: TerminalLogFn | None = None) -> dict[str, Any]:
    ctx: dict[str, Any] = {"reachable": False, "namenode": HDFS_NN_URL}
    await _log(log, "info", "fetch", "▸ HDFS NameNode JMX metrics")
    try:
        fs_url = f"{HDFS_NN_URL}/jmx?qry=Hadoop:service=NameNode,name=FSNamesystem"
        nn_url = f"{HDFS_NN_URL}/jmx?qry=Hadoop:service=NameNode,name=NameNodeInfo"
        t0 = time.monotonic()
        await _log(log, "cmd", "fetch", f"$ GET {fs_url}")
        await _log(log, "cmd", "fetch", f"$ GET {nn_url}")
        fs_r, nn_r = await asyncio.gather(
            client.get(fs_url),
            client.get(nn_url),
            return_exceptions=True,
        )
        ms = int((time.monotonic() - t0) * 1000)
        if isinstance(fs_r, httpx.Response) and fs_r.status_code == 200:
            beans = fs_r.json().get("beans", [])
            if beans:
                b = beans[0]
                ctx.update({
                    "reachable": True,
                    "hostname": b.get("tag.Hostname"),
                    "ha_state": b.get("tag.HAState"),
                    "capacity_total_gb": b.get("CapacityTotalGB"),
                    "capacity_used_gb": b.get("CapacityUsedGB"),
                    "capacity_remaining_gb": b.get("CapacityRemainingGB"),
                    "files_total": b.get("FilesTotal"),
                    "blocks_total": b.get("BlocksTotal"),
                    "live_datanodes": b.get("NumLiveDataNodes"),
                    "dead_datanodes": b.get("NumDeadDataNodes"),
                    "missing_blocks": b.get("MissingBlocks"),
                    "under_replicated_blocks": b.get("UnderReplicatedBlocks"),
                    "corrupt_blocks": b.get("CorruptBlocks"),
                    "default_replication_factor": 3,
                })
                await _log(
                    log, "ok", "fetch",
                    f"← HDFS: {b.get('CapacityUsedGB')}GB used, {b.get('FilesTotal')} files, "
                    f"{b.get('NumLiveDataNodes')} datanodes ({ms}ms)",
                )
        else:
            await _log(log, "warn", "fetch", f"← FSNamesystem JMX failed ({ms}ms)")

        if isinstance(nn_r, httpx.Response) and nn_r.status_code == 200:
            beans = nn_r.json().get("beans", [])
            if beans:
                b = beans[0]
                live = json.loads(b.get("LiveNodes") or "{}")
                ctx["hdfs_version"] = b.get("Version")
                ctx["safemode"] = b.get("Safemode") or "off"
                ctx["percent_used"] = round(float(b.get("PercentUsed", 0)) * 100, 4)
                ctx["datanodes"] = [
                    {
                        "host": host.split(":")[0],
                        "capacity_gb": round(node.get("capacity", 0) / (1024**3), 1),
                        "used_gb": round(node.get("used", 0) / (1024**3), 4),
                        "blocks": node.get("numBlocks", 0),
                        "state": node.get("adminState"),
                    }
                    for host, node in live.items()
                ]

        yarn_url = f"{YARN_URL}/ws/v1/cluster/info"
        await _log(log, "cmd", "fetch", f"$ GET {yarn_url}")
        try:
            yr = await client.get(yarn_url, timeout=4.0)
            if yr.status_code == 200:
                yinfo = yr.json().get("clusterInfo", {})
                ctx["yarn_ok"] = True
                ctx["yarn_state"] = yinfo.get("state", "UNKNOWN")
                ctx["yarn_rm"] = YARN_URL
                await _log(log, "ok", "fetch", f"← YARN RM: {yinfo.get('state', '?')}")
            else:
                ctx["yarn_ok"] = False
        except Exception:
            ctx["yarn_ok"] = False

    except Exception as exc:
        ctx["error"] = str(exc)
        await _log(log, "err", "fetch", f"✗ HDFS: {exc}")
    return ctx


async def collect_etl(client: httpx.AsyncClient, log: TerminalLogFn | None = None) -> dict[str, Any]:
    await _log(log, "info", "fetch", "▸ ETL stack (Airflow, Kafka, Spark)")
    health, kafka_ok, spark_ok = await asyncio.gather(
        _get_json(client, f"{AIRFLOW_URL}/api/v2/monitor/health", log, "Airflow health"),
        _probe_ok(client, KAFKA_UI_URL, log, "Kafka UI"),
        _probe_ok(client, SPARK_UI_URL, log, "Spark UI"),
    )
    connectors: list[str] = []
    await _log(log, "cmd", "fetch", f"$ GET {KAFKA_CONNECT_URL}/connectors")
    t0 = time.monotonic()
    try:
        r = await client.get(f"{KAFKA_CONNECT_URL}/connectors", timeout=5.0)
        ms = int((time.monotonic() - t0) * 1000)
        if r.status_code == 200:
            connectors = r.json() if isinstance(r.json(), list) else []
            await _log(log, "ok", "fetch", f"← Kafka Connect: {len(connectors)} connectors ({ms}ms)")
            for c in connectors:
                await _log(log, "info", "fetch", f"  · {c}")
        else:
            await _log(log, "warn", "fetch", f"← Kafka Connect {r.status_code} ({ms}ms)")
    except Exception as exc:
        await _log(log, "err", "fetch", f"✗ Kafka Connect: {exc}")

    airflow_detail: dict[str, str] = {}
    if isinstance(health, dict):
        for comp, info in health.items():
            if isinstance(info, dict) and "status" in info:
                airflow_detail[comp] = info["status"]
                await _log(log, "info", "fetch", f"  Airflow {comp}: {info['status']}")

    connector_status: list[dict[str, Any]] = []
    for name in connectors:
        status_url = f"{KAFKA_CONNECT_URL}/connectors/{name}/status"
        await _log(log, "cmd", "fetch", f"$ GET {status_url}")
        try:
            sr = await client.get(status_url, timeout=5.0)
            if sr.status_code == 200:
                st = sr.json()
                conn = st.get("connector", {})
                tasks = st.get("tasks", [])
                state = conn.get("state", "UNKNOWN")
                task_states = [t.get("state", "?") for t in tasks]
                connector_status.append({
                    "name": name,
                    "state": state,
                    "tasks": task_states,
                })
                await _log(log, "info", "fetch", f"  · {name}: {state} tasks={task_states}")
        except Exception as exc:
            connector_status.append({"name": name, "state": "ERROR", "error": str(exc)})

    return {
        "airflow_url": AIRFLOW_URL,
        "airflow_healthy": airflow_detail.get("scheduler") == "healthy",
        "airflow_components": airflow_detail,
        "kafka_ui_url": KAFKA_UI_URL,
        "kafka_ui_ok": kafka_ok,
        "kafka_connect_url": KAFKA_CONNECT_URL,
        "connectors": connectors,
        "connector_status": connector_status,
        "spark_ui_url": SPARK_UI_URL,
        "spark_ui_ok": spark_ok,
    }


async def collect_lakehouse(
    client: httpx.AsyncClient,
    containers: list[dict],
    log: TerminalLogFn | None = None,
) -> dict[str, Any]:
    await _log(log, "info", "fetch", "▸ Lakehouse (Trino, Spark, Kafka Connect)")
    trino_info = await _get_json(client, f"{TRINO_URL}/v1/info", log, "Trino /v1/info")
    running = sum(1 for c in containers if c.get("state") == "running")
    for c in containers:
        ports = ",".join(str(p.get("PublicPort")) for p in c.get("ports", []) if p.get("PublicPort")) or "internal"
        await _log(log, "info", "fetch", f"  · {c.get('name')}: {c.get('state')} ports={ports}")
    return {
        "host": LAKEHOUSE_HOST,
        "trino_url": TRINO_URL,
        "trino_ok": trino_info is not None,
        "trino_version": (trino_info or {}).get("nodeVersion", {}).get("version"),
        "trino_uptime": (trino_info or {}).get("uptime"),
        "trino_coordinator": (trino_info or {}).get("coordinator"),
        "spark_ui_url": SPARK_UI_URL,
        "kafka_connect_url": KAFKA_CONNECT_URL,
        "containers": _container_rows(containers, LAKEHOUSE_HOST),
        "running": running,
        "total": len(containers),
    }


async def collect_databases(
    client: httpx.AsyncClient,
    containers: list[dict],
    log: TerminalLogFn | None = None,
) -> dict[str, Any]:
    await _log(log, "info", "fetch", "▸ Database vault (Dockhand env 5)")
    running = sum(1 for c in containers if c.get("state") == "running")
    rows = _container_rows(containers)
    by_engine: dict[str, list[str]] = {}
    for r in rows:
        img = (r.get("image") or "").lower()
        name = (r.get("name") or "").lower()
        if "postgres" in img or "postgres" in name:
            engine = "PostgreSQL"
        elif "mysql" in img or "mysql" in name:
            engine = "MySQL"
        elif "mongo" in img or "mongo" in name:
            engine = "MongoDB"
        elif "cassandra" in img or "cassandra" in name:
            engine = "Cassandra"
        elif "neo4j" in img or "neo4j" in name:
            engine = "Neo4j"
        else:
            engine = "Other"
        port_str = ",".join(r["ports"]) or "internal"
        by_engine.setdefault(engine, []).append(f"{r['name']} ({r['state']}, ports {port_str})")
        await _log(log, "info", "fetch", f"  · {r['name']}: {r['state']} [{engine}] ports={port_str}")

    return {
        "dockhand_env": 5,
        "running": running,
        "total": len(containers),
        "containers": rows,
        "by_engine": by_engine,
    }


async def collect_command_center(
    client: httpx.AsyncClient,
    containers: list[dict],
    log: TerminalLogFn | None = None,
) -> dict[str, Any]:
    await _log(log, "info", "fetch", f"▸ Command Center VM304 (Dockhand env {DOCKHAND_ENV_COMMAND_CENTER})")
    running = sum(1 for c in containers if c.get("state") == "running")
    rows = _container_rows(containers, "10.0.21.33")
    for r in rows:
        lvl = "info" if r.get("state") == "running" else "warn"
        await _log(log, lvl, "fetch", f"  · {r.get('name')}: {r.get('state')}")
    return {
        "dockhand_env": DOCKHAND_ENV_COMMAND_CENTER,
        "dockhand_stack": "atc-agents-vm304",
        "host": "10.0.21.33",
        "vmid": 304,
        "url": "http://10.0.21.33/",
        "running": running,
        "total": len(containers),
        "containers": rows,
    }


async def collect_docker_rack(
    client: httpx.AsyncClient,
    containers: list[dict],
    log: TerminalLogFn | None = None,
) -> dict[str, Any]:
    await _log(log, "info", "fetch", "▸ Docker rack (Dockhand env 1)")
    running = sum(1 for c in containers if c.get("state") == "running")
    not_running = [c["name"] for c in containers if c.get("state") != "running"]
    for c in containers:
        ports = ",".join(str(p.get("PublicPort")) for p in c.get("ports", []) if p.get("PublicPort")) or "internal"
        lvl = "info" if c.get("state") == "running" else "warn"
        await _log(log, lvl, "fetch", f"  · {c.get('name')}: {c.get('state')} ports={ports}")
    return {
        "dockhand_url": DOCKHAND_URL,
        "dockhand_env": 1,
        "running": running,
        "total": len(containers),
        "not_running": not_running,
        "containers": _container_rows(containers, "10.0.21.45"),
    }


async def collect_objectscale(client: httpx.AsyncClient, log: TerminalLogFn | None = None) -> dict[str, Any]:
    await _log(log, "info", "fetch", "▸ ObjectScale S3 storage")
    ctx: dict[str, Any] = {
        "host": "10.0.20.111",
        "url": OBJECTSCALE_URL,
        "port": "9020",
        "bucket": "data",
        "reachable": False,
    }
    t0 = time.monotonic()
    await _log(log, "cmd", "probe", f"$ GET {OBJECTSCALE_URL}")
    try:
        r = await client.get(OBJECTSCALE_URL, timeout=4.0)
        ms = int((time.monotonic() - t0) * 1000)
        # 403/401 means API is up but unauthenticated
        ctx["reachable"] = r.status_code in (200, 401, 403, 405)
        ctx["status_code"] = r.status_code
        await _log(
            log, "ok" if ctx["reachable"] else "warn", "probe",
            f"← ObjectScale {r.status_code} ({'UP' if ctx['reachable'] else 'DOWN'}, {ms}ms)",
        )
    except Exception as exc:
        ms = int((time.monotonic() - t0) * 1000)
        ctx["error"] = str(exc)
        await _log(log, "err", "probe", f"✗ ObjectScale: {exc} ({ms}ms)")
    return ctx


async def collect_gpu_metrics(client: httpx.AsyncClient, log: TerminalLogFn | None = None) -> dict[str, Any]:
    gpu_url = GPU_URL
    host = GPU_URL
    if _get_gpu_urls is not None:
        try:
            u = _get_gpu_urls()
            gpu_url = u["gpu_url"]
            host = u["host"]
        except Exception:
            pass
    await _log(log, "info", "fetch", f"▸ GPU Lab metrics @ {host}")
    base = {"ok": False, "host": host, "ui_url": gpu_url}
    try:
        metrics_url = f"{gpu_url}/api/gpu/metrics"
        model_url = f"{gpu_url}/api/active-model"
        await _log(log, "cmd", "fetch", f"$ GET {metrics_url}")
        await _log(log, "cmd", "fetch", f"$ GET {model_url}")
        t0 = time.monotonic()
        metrics_r, model_r = await asyncio.gather(
            client.get(metrics_url),
            client.get(model_url),
            return_exceptions=True,
        )
        ms = int((time.monotonic() - t0) * 1000)
        gpus: list[dict[str, Any]] = []
        if isinstance(metrics_r, httpx.Response) and metrics_r.status_code == 200:
            current = metrics_r.json().get("current", {})
            gpus = [
                {
                    "index": g["index"],
                    "name": g["name"],
                    "util_gpu": g.get("util_gpu", 0),
                    "memory_used_mib": g.get("memory_used_mib", 0),
                    "memory_total_mib": g.get("memory_total_mib", 0),
                    "temperature_c": g.get("temperature_c", 0),
                    "power_w": g.get("power_w", 0),
                }
                for g in current.get("gpus", [])
            ]
            await _log(log, "ok", "fetch", f"← GPU metrics: {len(gpus)} devices ({ms}ms)")
            for g in gpus:
                await _log(
                    log, "info", "fetch",
                    f"  GPU{g['index']}: util {g['util_gpu']:.0f}% VRAM "
                    f"{g['memory_used_mib']:.0f}/{g['memory_total_mib']:.0f} MiB",
                )

        active_model = None
        inference_active = False
        vllm_url = None
        if isinstance(model_r, httpx.Response) and model_r.status_code == 200:
            model_data = model_r.json()
            active_model = model_data.get("name")
            inference_active = bool(model_data.get("inference_active"))
            vllm_url = model_data.get("base_url")
            await _log(log, "ok", "fetch", f"← Active model: {active_model} inference={'ON' if inference_active else 'OFF'}")

        return {
            **base,
            "ok": len(gpus) > 0 or inference_active,
            "inference_active": inference_active,
            "active_model": active_model,
            "vllm_url": vllm_url,
            "gpu_count": len(gpus),
            "gpus": gpus,
        }
    except Exception as exc:
        await _log(log, "err", "fetch", f"✗ GPU Lab: {exc}")
        return {**base, "error": str(exc)}


def _section_docker(d: dict[str, Any]) -> list[str]:
    lines = [
        f"Docker rack (Dockhand env 1): {d['running']}/{d['total']} running",
        f"Dockhand: {d['dockhand_url']}",
    ]
    if d.get("not_running"):
        lines.append(f"Not running: {', '.join(d['not_running'])}")
    for c in d.get("containers", []):
        port_str = ",".join(c["ports"]) if c["ports"] else "internal"
        lines.append(f"  - {c['name']}: {c['state']} | {c['image']} | ports {port_str}")
    return lines


def _fmt_count(n: Any) -> str:
    if n is None:
        return "?"
    try:
        return f"{int(n):,}"
    except (TypeError, ValueError):
        return str(n)


def _section_databases(d: dict[str, Any]) -> list[str]:
    lines = [f"Databases (Dockhand env {d['dockhand_env']}): {d['running']}/{d['total']} running"]
    for engine, items in d.get("by_engine", {}).items():
        lines.append(f"  {engine}:")
        for item in items:
            lines.append(f"    - {item}")
    inv = d.get("inventory") or {}
    if inv:
        lines.append(
            f"  Live data inventory @ {inv.get('host', '?')}: "
            f"{inv.get('engines_ok', 0)}/{inv.get('engines_total', 0)} engines queried"
        )
        for key, eng in (inv.get("engines") or {}).items():
            if not eng.get("ok"):
                err = str(eng.get("error", "unknown"))[:100]
                lines.append(f"  {eng.get('engine', key)}: ERROR — {err}")
                continue
            label = eng.get("engine", key)
            if eng.get("size_human"):
                lines.append(f"  {label} ({eng.get('database', '')}): {eng['size_human']}")
            for tbl in eng.get("tables") or []:
                rows = tbl.get("rows")
                cols = ", ".join((tbl.get("columns") or [])[:8])
                extra = ""
                if tbl.get("top_regions"):
                    extra = f" | regions: {tbl['top_regions']}"
                elif tbl.get("top_event_types"):
                    extra = f" | event_types: {tbl['top_event_types']}"
                lines.append(f"    · {tbl['name']}: {_fmt_count(rows)} rows | cols: {cols}{extra}")
            for coll in eng.get("collections") or []:
                docs = coll.get("documents")
                fields = ", ".join(coll.get("fields") or [])
                extra = f" | types: {coll['top_types']}" if coll.get("top_types") else ""
                lines.append(f"    · {coll['name']}: {_fmt_count(docs)} docs | fields: {fields}{extra}")
            for node in eng.get("nodes") or []:
                lines.append(f"    · {node['label']} nodes: {_fmt_count(node.get('count'))}")
            if eng.get("relationships"):
                rels = ", ".join(f"{r['type']}={_fmt_count(r.get('count'))}" for r in eng["relationships"][:5])
                lines.append(f"    · relationships: {rels or 'none'}")
    return lines


def _section_lakehouse(d: dict[str, Any]) -> list[str]:
    lines = [
        f"Lakehouse host: {d['host']} — {d['running']}/{d['total']} containers running",
        f"Trino: {d['trino_url']} — {'UP' if d['trino_ok'] else 'DOWN'}"
        + (f" (v{d['trino_version']}, uptime {d.get('trino_uptime')})" if d.get("trino_ok") else ""),
        f"Spark UI: {d['spark_ui_url']}",
        f"Kafka Connect: {d['kafka_connect_url']}",
    ]
    for c in d.get("containers", []):
        port_str = ",".join(c["ports"]) if c["ports"] else "internal"
        lines.append(f"  - {c['name']}: {c['state']} | {c['image']} | ports {port_str}")
    return lines


def _section_etl(d: dict[str, Any]) -> list[str]:
    lines = [
        f"Airflow ({d['airflow_url']}): {'HEALTHY' if d['airflow_healthy'] else 'DEGRADED'}",
    ]
    for comp, st in d.get("airflow_components", {}).items():
        lines.append(f"  - {comp}: {st}")
    lines.append(f"Kafka UI ({d['kafka_ui_url']}): {'UP' if d['kafka_ui_ok'] else 'DOWN'}")
    lines.append(f"Kafka Connect ({d['kafka_connect_url']}): connectors {d.get('connectors') or 'none listed'}")
    if d.get("connectors"):
        lines.append("  Registered connector names (exact): " + ", ".join(d["connectors"]))
    for cs in d.get("connector_status") or []:
        tasks = cs.get("tasks") or []
        lines.append(f"  Connector {cs['name']}: {cs.get('state', '?')}" + (f" tasks={tasks}" if tasks else ""))
    lines.append(f"Spark UI ({d['spark_ui_url']}): {'UP' if d['spark_ui_ok'] else 'DOWN'}")
    return lines


def _section_hadoop(h: dict[str, Any]) -> list[str]:
    lines = ["HDFS / Hadoop:"]
    if not h.get("reachable"):
        lines.append(f"  UNREACHABLE: {h.get('error', 'NameNode probe failed')}")
        return lines
    lines.extend([
        f"  NameNode: {h['namenode']} ({h.get('hostname')}, HA {h.get('ha_state')})",
        f"  Version: {h.get('hdfs_version')}, safemode: {h.get('safemode')}",
        f"  Capacity: {h.get('capacity_used_gb')} GB used / {h.get('capacity_total_gb')} GB total "
        f"({h.get('capacity_remaining_gb')} GB free, {h.get('percent_used', 0)}% used)",
        f"  Files: {h.get('files_total')}, Blocks: {h.get('blocks_total')}",
        f"  DataNodes: {h.get('live_datanodes')} live, {h.get('dead_datanodes')} dead",
        f"  Replication factor (dfs.replication): {h.get('default_replication_factor')}",
        f"  Block health: missing={h.get('missing_blocks')}, under-replicated={h.get('under_replicated_blocks')}, corrupt={h.get('corrupt_blocks')}",
    ])
    for dn in h.get("datanodes", []):
        lines.append(
            f"    - {dn['host']}: {dn['used_gb']} GB / {dn['capacity_gb']} GB, {dn['blocks']} blocks, {dn['state']}"
        )
    if (h.get("capacity_used_gb") or 0) < 0.01 and (h.get("files_total") or 0) > 0:
        lines.append("  Note: metadata/small files only — almost no user data stored yet.")
    return lines


def _section_gpu(g: dict[str, Any]) -> list[str]:
    lines = ["GPU Lab / vLLM inference:"]
    if not g.get("ok"):
        lines.append(f"  OFFLINE: {g.get('error', 'unreachable')}")
        return lines
    lines.extend([
        f"  Manager: {g.get('ui_url')}",
        f"  Model: {g.get('active_model')} (inference {'ON' if g.get('inference_active') else 'OFF'})",
        f"  vLLM endpoint: {g.get('vllm_url')}",
        f"  GPUs: {g.get('gpu_count')}x V100",
    ])
    for gpu in g.get("gpus", []):
        lines.append(
            f"    GPU{gpu['index']}: util {gpu['util_gpu']:.0f}%, "
            f"VRAM {gpu['memory_used_mib']:.0f}/{gpu['memory_total_mib']:.0f} MiB, "
            f"{gpu['temperature_c']}°C, {gpu['power_w']:.0f}W"
        )
    return lines




def _section_objectscale(o: dict[str, Any]) -> list[str]:
    lines = [
        f"ObjectScale S3 ({o.get('host')}:{o.get('port')}): {'REACHABLE' if o.get('reachable') else 'DOWN'}",
        f"  API: {o.get('url')} (HTTP {o.get('status_code', '?')})",
        f"  Bucket: {o.get('bucket', 'data')} — landing zone for s3-kafka-consumer & Iceberg",
    ]
    if o.get("error"):
        lines.append(f"  Error: {o['error']}")
    return lines


def _section_command_center(c: dict[str, Any]) -> list[str]:
    lines = [
        f"Command Center VM304: {c.get('host')} — {c.get('running', 0)}/{c.get('total', 0)} containers",
        f"  URL: {c.get('url')}",
        f"  Dockhand env: {c.get('dockhand_env')}",
    ]
    for row in c.get("containers", []):
        port_str = ",".join(row["ports"]) if row.get("ports") else "internal"
        lines.append(f"  - {row['name']}: {row['state']} | ports {port_str}")
    return lines


def collect_governance(log: TerminalLogFn | None = None) -> dict[str, Any]:
    """In-process governance snapshot: CDC, movements, PII catalog, lineage.

    Reads directly from the live modules (no HTTP) so the LLM always has the
    current data-platform governance picture.
    """
    out: dict[str, Any] = {"openmetadata_url": os.getenv("OPENMETADATA_URL", "http://10.0.21.47:8585")}
    try:
        from cdc_consumer import snapshot as _cdc_snapshot
        out["cdc"] = _cdc_snapshot(15)
    except Exception as exc:
        out["cdc"] = {"error": str(exc)}
    try:
        from movements import MOVEMENTS, last_runs
        runs = last_runs()
        out["movements"] = [
            {"id": m["id"], "label": m["label"], "kind": m["kind"], "from": m["from"], "to": m["to"],
             "last": runs.get(m["id"])}
            for m in MOVEMENTS
        ]
    except Exception as exc:
        out["movements"] = {"error": str(exc)}
    try:
        from pii_catalog import get_pii
        pii = get_pii()
        out["pii_summary"] = pii.get("summary", {})
        out["pii_source"] = pii.get("source")
        out["pii_datasets"] = [
            {"label": d["label"], "pii_count": d["pii_count"], "all_masked": d["all_masked"],
             "columns": [f"{c['name']}:{c['category']}{'(masked)' if c['masked'] else ''}" for c in d["pii_columns"]]}
            for d in pii.get("datasets", [])
        ]
    except Exception as exc:
        out["pii_summary"] = {"error": str(exc)}
    out["lineage"] = [
        "postgres_sales.public.sales_orders -> iceberg.curated_masked.sales_orders_masked (PII masked)",
        "mysql_hr.hr.employee_events -> iceberg.curated_masked.employee_events_masked (PII masked)",
        "hdfs:/data/historical/sales_orders -> iceberg.hadoop.historical_sales_hdfs",
    ]
    try:
        from dq_monitor import summary_for_llm as _dq
        out["data_quality"] = _dq()
    except Exception as exc:
        out["data_quality"] = {"error": str(exc)}
    try:
        from observability import summary_for_llm as _obs
        out["observability"] = _obs()
    except Exception as exc:
        out["observability"] = {"error": str(exc)}
    try:
        from catalog_governance import summary_for_llm as _own
        out["ownership"] = _own()
    except Exception as exc:
        out["ownership"] = {"error": str(exc)}
    return out


def _section_governance(g: dict[str, Any]) -> list[str]:
    lines = ["Data governance (CDC · movements · PII · lineage):",
             f"  OpenMetadata catalog: {g.get('openmetadata_url')} (catalog, lineage, PII auto-classification)"]
    cdc = g.get("cdc") or {}
    if "error" not in cdc:
        by_src = ", ".join(f"{k}={v}" for k, v in (cdc.get("by_source") or {}).items()) or "none"
        lines.append(f"  CDC stream: connected={cdc.get('connected')} consumed={cdc.get('consumed')} "
                     f"changes/15m={cdc.get('window_total')} ({by_src})")
    movements = g.get("movements")
    if isinstance(movements, list):
        lines.append("  Data movements (trigger via Command Center / ETL agent):")
        for m in movements:
            last = m.get("last") or {}
            st = last.get("state", "never run")
            extra = f" rows={last.get('rows')}" if last.get("rows") is not None else ""
            lines.append(f"    - {m['id']}: {m['from']}→{m['to']} [{m['kind']}] last={st}{extra}")
    ps = g.get("pii_summary") or {}
    if "error" not in ps:
        lines.append(f"  PII catalog ({g.get('pii_source')}): {ps.get('pii_columns', 0)} PII cols, "
                     f"{ps.get('masked_columns', 0)} masked / {ps.get('unmasked_columns', 0)} unmasked")
        for d in g.get("pii_datasets") or []:
            tag = "all masked" if d["all_masked"] else "UNMASKED"
            lines.append(f"    - {d['label']}: {d['pii_count']} PII [{tag}] {', '.join(d['columns'][:8])}")
    lines.append("  Lineage:")
    for ln in g.get("lineage") or []:
        lines.append(f"    - {ln}")
    dq = g.get("data_quality") or {}
    if isinstance(dq, dict) and "error" not in dq:
        lines.append(f"  Data quality (continuous, live tables): platform score {dq.get('platform_dq_score')}")
        for d in dq.get("datasets") or []:
            iss = f" issues: {', '.join(d['issues'][:3])}" if d.get("issues") else ""
            lines.append(f"    - {d['key']}: score {d.get('score')}{iss}")
    own = g.get("ownership") or {}
    if isinstance(own, dict) and "error" not in own:
        orph = own.get("orphan_datasets") or []
        lines.append(f"  Ownership: {len(own.get('owners', {}))} assigned"
                     + (f", orphans (no owner): {', '.join(orph)}" if orph else ", no orphans"))
        for k, v in (own.get("owners") or {}).items():
            if v.get("owner"):
                lines.append(f"    - {k}: owner={v.get('owner')} steward={v.get('steward') or '—'} tier={v.get('tier') or '—'}")
    obs = g.get("observability") or {}
    if isinstance(obs, dict) and "error" not in obs:
        ac = obs.get("active_alerts") or {}
        lines.append(f"  Observability alerts: {ac.get('total', 0)} active "
                     f"(critical={ac.get('critical', 0)}, warning={ac.get('warning', 0)})")
        for a in (obs.get("alerts") or [])[:5]:
            lines.append(f"    - [{a['severity']}] {a['dataset']}: {a['message']}")
    return lines


def _section_cluster_registry(_: dict[str, Any]) -> list[str]:
    """Static cluster map — always available even when probes fail."""
    lines = ["Cluster infrastructure map (Proxmox VMs & roles):"]
    for nid, node in NODE_REGISTRY.items():
        if nid in ("mo-commander", "bart-commander", "network-watcher", "mcp-coordinator"):
            continue
        vmid = node.get("vmid", "?")
        lines.append(
            f"  - {node['label']}: {node.get('vm')} VMID {vmid} @ {node.get('ip')} — {node.get('role')}"
        )
        desc = node.get("description") or ""
        if desc:
            lines.append(f"      {desc[:140]}")
    lines.append("")
    lines.append("Supervisors & control plane:")
    for nid in ("mo-commander", "bart-commander", "mcp-coordinator", "network-watcher"):
        node = NODE_REGISTRY[nid]
        lines.append(f"  - {node['label']}: {(node.get('description') or '')[:120]}")
    return lines

SECTION_BUILDERS = {
    "docker": _section_docker,
    "databases": _section_databases,
    "lakehouse": _section_lakehouse,
    "etl": _section_etl,
    "hadoop": _section_hadoop,
    "gpu": _section_gpu,
    "objectscale": _section_objectscale,
    "command_center": _section_command_center,
    "governance": _section_governance,
    "cluster_registry": _section_cluster_registry,
}

DOMAIN_ORDER = ["docker", "databases", "lakehouse", "etl", "hadoop", "gpu", "objectscale", "command_center", "governance", "cluster_registry"]


async def collect_full_lab_context(
    gpu_data: dict[str, Any] | None = None,
    log: TerminalLogFn | None = None,
    include_inventory: bool = True,
) -> dict[str, Any]:
    """Gather all lab domains in parallel with optional live terminal logging."""
    await _log(log, "info", "fetch", "═══ Lab snapshot collection started ═══")
    async with httpx.AsyncClient(timeout=10.0, verify=False) as client:
        if gpu_data is None:
            gpu_data = await collect_gpu_metrics(client, log)

        docker_raw, db_raw, lake_raw, cc_raw, gpu_raw, hdfs, etl, objectscale = await asyncio.gather(
            dockhand_containers(client, DOCKHAND_ENVS["docker01"], log),
            dockhand_containers(client, DOCKHAND_ENVS["db02"], log),
            dockhand_containers(client, DOCKHAND_ENVS["lakehouse"], log),
            dockhand_containers(client, DOCKHAND_ENV_COMMAND_CENTER, log),
            dockhand_containers(client, DOCKHAND_ENVS["gpu_dev"], log),
            collect_hdfs(client, log),
            collect_etl(client, log),
            collect_objectscale(client, log),
        )
        docker = await collect_docker_rack(client, docker_raw, log)
        databases = await collect_databases(client, db_raw, log)
        if include_inventory:
            try:
                databases["inventory"] = await collect_database_inventory()
                inv_ok = databases["inventory"].get("engines_ok", 0)
                await _log(log, "ok", "fetch", f"← Database inventory: {inv_ok} engines")
            except Exception as exc:
                await _log(log, "warn", "fetch", f"✗ Database inventory: {exc}")
                databases["inventory"] = {"error": str(exc)}
        lakehouse = await collect_lakehouse(client, lake_raw, log)
        command_center = await collect_command_center(client, cc_raw, log)
        if gpu_data is not None and gpu_raw:
            gpu_running = sum(1 for c in gpu_raw if c.get("state") == "running")
            gpu_data = {
                **gpu_data,
                "dockhand_env": DOCKHAND_ENVS["gpu_dev"],
                "dockhand_containers": _container_rows(gpu_raw, GPU_URL.replace("http://", "").split(":")[0]),
                "dockhand_running": gpu_running,
                "dockhand_total": len(gpu_raw),
            }

    try:
        governance = collect_governance(log)
        await _log(log, "ok", "fetch", "← Governance: CDC/movements/PII/lineage")
    except Exception as exc:
        await _log(log, "warn", "fetch", f"✗ Governance: {exc}")
        governance = {"error": str(exc)}

    await _log(log, "ok", "fetch", "═══ Lab snapshot complete ═══")
    return {
        "ts": datetime.now(timezone.utc).isoformat(),
        "docker": docker,
        "databases": databases,
        "lakehouse": lakehouse,
        "etl": etl,
        "hadoop": hdfs,
        "gpu": gpu_data,
        "objectscale": objectscale,
        "command_center": command_center,
        "governance": governance,
    }


def format_context_for_agent(agent_id: str, snapshot: dict[str, Any]) -> str:
    """Format full lab snapshot for LLM; primary domain first."""
    primary = AGENT_PRIMARY_DOMAIN.get(agent_id, "docker")
    lines = [
        f"ATC Lab live snapshot — {snapshot.get('ts')}",
        f"Your primary domain: {primary.upper()}",
    ]
    if snapshot.get("domains_summary"):
        lines.append(f"Health summary: {json.dumps(snapshot['domains_summary'], default=str)}")
    lines.extend(["", f"=== PRIMARY: {primary.upper()} ==="])

    if primary in snapshot and primary in SECTION_BUILDERS:
        lines.extend(SECTION_BUILDERS[primary](snapshot[primary]))
    lines.append("")
    lines.append("=== FULL LAB (all domains) ===")

    if "cluster_registry" not in snapshot:
        snapshot = {**snapshot, "cluster_registry": {}}

    for domain in DOMAIN_ORDER:
        if domain == primary:
            continue
        if domain not in snapshot or domain not in SECTION_BUILDERS:
            continue
        lines.append("")
        lines.append(f"--- {domain.upper()} ---")
        lines.extend(SECTION_BUILDERS[domain](snapshot[domain]))

    return "\n".join(lines)