"""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

DOCKHAND_URL = os.getenv("DOCKHAND_URL", "http://10.0.21.45:8082")
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")
GPU_URL = os.getenv("GPU_URL", "http://10.0.20.106:9000")

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}, 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()
                ]
    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']}")

    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,
        "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_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_gpu_metrics(client: httpx.AsyncClient, log: TerminalLogFn | None = None) -> dict[str, Any]:
    await _log(log, "info", "fetch", "▸ GPU Lab metrics")
    base = {"ok": False, "host": GPU_URL, "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 _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}")
    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"]))
    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


SECTION_BUILDERS = {
    "docker": _section_docker,
    "databases": _section_databases,
    "lakehouse": _section_lakehouse,
    "etl": _section_etl,
    "hadoop": _section_hadoop,
    "gpu": _section_gpu,
}

DOMAIN_ORDER = ["docker", "databases", "lakehouse", "etl", "hadoop", "gpu"]


async def collect_full_lab_context(
    gpu_data: dict[str, Any] | None = None,
    log: TerminalLogFn | None = None,
) -> 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, hdfs, etl = await asyncio.gather(
            dockhand_containers(client, 1, log),
            dockhand_containers(client, 5, log),
            dockhand_containers(client, 9, log),
            collect_hdfs(client, log),
            collect_etl(client, log),
        )
        docker = await collect_docker_rack(client, docker_raw, log)
        databases = await collect_databases(client, db_raw, log)
        lakehouse = await collect_lakehouse(client, lake_raw, log)

    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,
    }


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) ===")

    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)