"""Build UI workload + topology payload from lab snapshot."""

from __future__ import annotations

from typing import Any

from node_registry import NODE_AGENT, NODE_REGISTRY
from topology_views import build_all_topologies

OBJECTSCALE_HOST = "10.0.20.111"
OBJECTSCALE_PORT = "9020"
OBJECTSCALE_BUCKET = "data"


def _level(running: int, total: int) -> str:
    if total == 0:
        return "unknown"
    ratio = running / total
    if ratio >= 0.9:
        return "ok"
    if ratio >= 0.5:
        return "warn"
    return "down"


def _app_row(c: dict[str, Any]) -> dict[str, Any]:
    img = c.get("image") or ""
    short_img = img.split("/")[-1].split(":")[0][:20]
    return {
        "name": c.get("name", "?"),
        "state": c.get("state", "unknown"),
        "image": short_img,
        "ports": c.get("ports") or [],
        "host": c.get("host") or "",
    }


def _find_container(containers: list[dict], *patterns: str) -> dict | None:
    for c in containers:
        hay = f"{c.get('name', '')} {c.get('image', '')}".lower()
        if any(p.lower() in hay for p in patterns):
            return c
    return None


def _s3_level(lakehouse: dict[str, Any], objectscale_ok: bool) -> str:
    containers = lakehouse.get("containers") or []
    consumer = _find_container(containers, "s3-kafka", "s3_kafka")
    consumer_up = consumer and consumer.get("state") == "running"
    if objectscale_ok and consumer_up:
        return "ok"
    if objectscale_ok or consumer_up:
        return "warn"
    return "down"


def build_workload_payload(snap: dict[str, Any]) -> dict[str, Any]:
    docker = snap.get("docker", {})
    databases = snap.get("databases", {})
    lakehouse = snap.get("lakehouse", {})
    etl = snap.get("etl", {})
    hadoop = snap.get("hadoop", {})
    gpu = snap.get("gpu", {})
    objectscale = snap.get("objectscale", {})
    command = snap.get("command_center", {})
    try:
        from gpu_config import resolve_gpu_identity
        gpu_id = resolve_gpu_identity(gpu)
    except Exception:
        host = gpu.get("ip") or gpu.get("host") or "10.0.10.106"
        gpu_id = {
            "vm": "atc-gpu-prod",
            "ip": host,
            "ui_url": gpu.get("ui_url") or f"http://{host}:9000",
            "llm_url": gpu.get("vllm_url") or f"http://{host}:8001/v1",
            "vmid": 306,
        }

    docker_apps = [_app_row(c) for c in docker.get("containers", [])]
    db_apps = [_app_row(c) for c in databases.get("containers", [])]
    lake_apps = [_app_row(c) for c in lakehouse.get("containers", [])]

    lake_containers = lakehouse.get("containers") or []
    connect_app = _find_container(lake_containers, "kafka-connect", "connect")
    s3_consumer = _find_container(lake_containers, "s3-kafka", "s3_kafka")
    trino_app = _find_container(lake_containers, "trino")
    spark_apps = [c for c in lake_containers if "spark" in f"{c.get('name', '')} {c.get('image', '')}".lower()]

    hdfs_ok = hadoop.get("reachable", False)
    etl_ok = etl.get("airflow_healthy") and etl.get("kafka_ui_ok")
    objectscale_ok = objectscale.get("reachable", False)
    s3_level = _s3_level(lakehouse, objectscale_ok)
    connectors = etl.get("connectors") or []

    etl_apps = [
        {"name": "Airflow", "state": "running" if etl.get("airflow_healthy") else "down", "image": "airflow", "ports": ["8080"], "host": "10.0.21.55"},
        {"name": "Kafka", "state": "running" if etl.get("kafka_ui_ok") else "down", "image": "kafka", "ports": ["9092"], "host": "10.0.21.36"},
        {"name": "Kafka UI", "state": "running" if etl.get("kafka_ui_ok") else "down", "image": "kafka-ui", "ports": ["9000"], "host": "10.0.21.36"},
        *[
            {"name": c, "state": "running", "image": "connect", "ports": ["8083"], "host": lakehouse.get("host", "10.0.21.50")}
            for c in connectors
        ],
    ]

    s3_apps = [
        {"name": "ObjectScale", "state": "running" if objectscale_ok else "down", "image": "objectscale", "ports": [OBJECTSCALE_PORT], "host": OBJECTSCALE_HOST},
        {"name": f"bucket/{OBJECTSCALE_BUCKET}", "state": "running" if objectscale_ok else "down", "image": "s3", "ports": [], "host": OBJECTSCALE_HOST},
    ]
    if s3_consumer:
        s3_apps.insert(0, _app_row(s3_consumer))

    zones = [
        {
            "id": "docker",
            "label": "DOCKER RACK",
            "x": 8,
            "color": "#b366ff",
            "level": _level(docker.get("running", 0), docker.get("total", 1) or 1),
            "running": docker.get("running", 0),
            "total": docker.get("total", 0),
            "apps": docker_apps,
            "vm": "atc-docker01",
            "ip": "10.0.21.45",
        },
        {
            "id": "db",
            "label": "DB VAULT",
            "x": 22,
            "color": "#ffaa00",
            "level": _level(databases.get("running", 0), databases.get("total", 1) or 1),
            "running": databases.get("running", 0),
            "total": databases.get("total", 0),
            "apps": db_apps,
            "vm": "atc-db02",
            "ip": "10.0.21.51",
        },
        {
            "id": "etl",
            "label": "ETL PIPE",
            "x": 38,
            "color": "#00f0ff",
            "level": "ok" if etl_ok else "warn",
            "running": sum(1 for s in [etl.get("airflow_healthy"), etl.get("kafka_ui_ok"), etl.get("spark_ui_ok")] if s),
            "total": 3,
            "apps": etl_apps,
            "vm": "airflow + kafka",
            "ip": "10.0.21.55 / .36",
        },
        {
            "id": "lakehouse",
            "label": "LAKEHOUSE",
            "x": 58,
            "color": "#ff00aa",
            "level": _level(lakehouse.get("running", 0), lakehouse.get("total", 1) or 1),
            "running": lakehouse.get("running", 0),
            "total": lakehouse.get("total", 0),
            "apps": lake_apps,
            "trino_ok": lakehouse.get("trino_ok"),
            "vm": "atc-lake01",
            "ip": lakehouse.get("host", "10.0.21.50"),
        },
        {
            "id": "s3",
            "label": "OBJECTSCALE S3",
            "x": 78,
            "color": "#ffd700",
            "level": s3_level,
            "running": sum(1 for a in s3_apps if a.get("state") == "running"),
            "total": len(s3_apps),
            "apps": s3_apps,
            "vm": "atc-objectscale",
            "ip": OBJECTSCALE_HOST,
            "bucket": OBJECTSCALE_BUCKET,
        },
        {
            "id": "hadoop",
            "label": "HADOOP HDFS",
            "x": 50,
            "color": "#39ff14",
            "level": "ok" if hdfs_ok else "warn",
            "running": hadoop.get("live_datanodes", 0),
            "total": (hadoop.get("live_datanodes") or 0) + (hadoop.get("dead_datanodes") or 0) or 3,
            "apps": [
                {"name": "NameNode", "state": "running" if hdfs_ok else "down", "image": "hdfs-nn", "ports": ["9870"], "host": "10.0.21.61"},
                *[
                    {"name": dn.get("host", "?").split(".")[0], "state": "running", "image": "datanode", "ports": ["9866"], "host": dn.get("host", "")}
                    for dn in hadoop.get("datanodes", [])
                ],
            ],
            "hdfs_used_gb": hadoop.get("capacity_used_gb"),
            "hdfs_total_gb": hadoop.get("capacity_total_gb"),
            "vm": "hadoop cluster",
            "ip": "10.0.21.61–70",
        },
    ]

    def _node(
        nid: str,
        label: str,
        vm: str,
        ip: str,
        x: float,
        y: float,
        color: str,
        level: str,
        role: str,
        apps: list[dict],
        running: int,
        total: int,
        extra: dict | None = None,
    ) -> dict[str, Any]:
        reg = NODE_REGISTRY.get(nid, {})
        row: dict[str, Any] = {
            "id": nid,
            "label": label,
            "vm": vm,
            "ip": ip,
            "x": x,
            "y": y,
            "color": reg.get("color", color),
            "level": level,
            "role": role,
            "apps": apps,
            "running": running,
            "total": total,
            "description": reg.get("description", ""),
            "agent_id": NODE_AGENT.get(nid),
            "links": reg.get("links", []),
            "endpoints": reg.get("endpoints", []),
            "commands": reg.get("commands", []),
            "vmid": reg.get("vmid"),
            "pve": reg.get("pve"),
        }
        if extra:
            row.update(extra)
        if nid == "gpu":
            row["vm"] = gpu_id["vm"]
            row["ip"] = gpu_id["ip"]
            row["vmid"] = gpu_id.get("vmid", row.get("vmid"))
            row["links"] = [
                {"label": "GPU Lab UI", "url": gpu_id["ui_url"]},
                {"label": "vLLM API", "url": gpu_id["llm_url"]},
            ]
            row["endpoints"] = [
                {"name": "gpu-lab", "host": gpu_id["ip"], "port": "9000", "proto": "http"},
                {"name": "vllm", "host": gpu_id["ip"], "port": "8001", "proto": "http"},
            ]
        return row

    connect_running = 1 if connect_app and connect_app.get("state") == "running" else 0
    consumer_running = 1 if s3_consumer and s3_consumer.get("state") == "running" else 0

    topology_nodes = [
        _node("airflow", "Airflow", "atc-airflow01", "10.0.21.55", 6, 18, "#00f0ff", "ok" if etl.get("airflow_healthy") else "warn", "orchestrator",
              [{"name": "scheduler", "state": "running" if etl.get("airflow_healthy") else "down", "image": "airflow", "ports": ["8080"]}], int(etl.get("airflow_healthy", False)), 1),
        _node("db", "DB Vault", "atc-db02", "10.0.21.51", 22, 18, "#ffaa00", _level(databases.get("running", 0), databases.get("total", 1) or 1), "sources",
              db_apps, databases.get("running", 0), databases.get("total", 0)),
        _node("debezium", "Debezium", "atc-lake01", "10.0.21.50", 38, 18, "#ff66cc", "ok" if connect_running and connectors else "warn", "cdc",
              [_app_row(connect_app)] if connect_app else [], len(connectors), max(len(connectors), 1),
              {"connectors": connectors}),
        _node("kafka", "Kafka", "atc-kafka01", "10.0.21.36", 54, 18, "#00f0ff", "ok" if etl.get("kafka_ui_ok") else "warn", "bus",
              [{"name": "broker", "state": "running" if etl.get("kafka_ui_ok") else "down", "image": "kafka", "ports": ["9092"]}], int(etl.get("kafka_ui_ok", False)), 1),
        _node("lakehouse", "Lakehouse", "atc-lake01", "10.0.21.50", 70, 18, "#ff00aa", _level(lakehouse.get("running", 0), lakehouse.get("total", 1) or 1), "compute",
              lake_apps, lakehouse.get("running", 0), lakehouse.get("total", 0),
              {"trino_ok": lakehouse.get("trino_ok"), "spark_count": len(spark_apps)}),
        _node("s3", "ObjectScale S3", "atc-objectscale", OBJECTSCALE_HOST, 88, 18, "#ffd700", s3_level, "storage",
              s3_apps, sum(1 for a in s3_apps if a.get("state") == "running"), len(s3_apps),
              {"bucket": OBJECTSCALE_BUCKET, "port": OBJECTSCALE_PORT, "consumer_ok": bool(consumer_running)}),
        _node("docker", "Docker Rack", "atc-docker01", "10.0.21.45", 10, 52, "#b366ff", _level(docker.get("running", 0), docker.get("total", 1) or 1), "infra",
              docker_apps, docker.get("running", 0), docker.get("total", 0)),
        _node("hadoop", "Hadoop HDFS", "atc-hadoop-m01", "10.0.21.61", 50, 52, "#39ff14", "ok" if hdfs_ok else "warn", "parallel",
              zones[-1]["apps"], hadoop.get("live_datanodes", 0), zones[-1]["total"],
              {"hdfs_used_gb": hadoop.get("capacity_used_gb"), "hdfs_total_gb": hadoop.get("capacity_total_gb")}),
        _node("gpu", "GPU Lab", gpu_id["vm"], gpu_id["ip"], 88, 52, "#76b900", "ok" if gpu.get("ok") else "down", "inference",
              [{"name": gpu.get("active_model") or "vLLM", "state": "running" if gpu.get("inference_active") else "down", "image": "vllm", "ports": ["8001", "9000"]}],
              gpu.get("gpu_count", 0), gpu.get("gpu_count", 0) or 4,
              {"model": gpu.get("active_model"), "util": round(sum(g.get("util_gpu", 0) for g in gpu.get("gpus", [])) / max(len(gpu.get("gpus", [])), 1), 1),
               "ui_url": gpu_id["ui_url"], "vllm_url": gpu_id["llm_url"]}),
        _node("command", "Command Center", "MCP · VM304", "10.0.21.33", 50, 78, "#00f0ff",
              _level(command.get("running", 0), command.get("total", 1) or 1), "hub",
              command.get("containers") and [_app_row(c) for c in command.get("containers", [])] or [
                  {"name": "atc-agents-api", "state": "running", "image": "atc-agents-api", "ports": ["3201"]},
                  {"name": "atc-agents-ui", "state": "running", "image": "atc-agents-ui", "ports": ["80"]},
                  {"name": "postgres", "state": "running", "image": "postgres", "ports": ["5432"]},
                  {"name": "redis", "state": "running", "image": "redis", "ports": ["6379"]},
                  {"name": "caddy", "state": "running", "image": "caddy", "ports": ["80"]},
              ],
              command.get("running", 5), command.get("total", 5) or 5),
    ]

    def _edge(eid: str, src: str, dst: str, label: str, kind: str, active: bool = True) -> dict[str, Any]:
        return {"id": eid, "from": src, "to": dst, "label": label, "kind": kind, "active": active}

    pipeline_ok = etl.get("airflow_healthy") and len(connectors) > 0 and etl.get("kafka_ui_ok")
    s3_flow_ok = pipeline_ok and consumer_running and objectscale_ok

    topology_edges = [
        _edge("e-seed", "airflow", "db", "seed data", "pipeline", bool(etl.get("airflow_healthy"))),
        _edge("e-cdc", "db", "debezium", "CDC", "pipeline", bool(connectors)),
        _edge("e-topics", "debezium", "kafka", "topics", "pipeline", bool(connectors and etl.get("kafka_ui_ok"))),
        _edge("e-stream", "kafka", "lakehouse", "stream", "pipeline", bool(etl.get("kafka_ui_ok") and lakehouse.get("trino_ok"))),
        _edge("e-s3", "lakehouse", "s3", "s3-kafka-consumer", "pipeline", bool(s3_flow_ok)),
        _edge("e-iceberg", "lakehouse", "s3", "Trino Iceberg", "query", bool(lakehouse.get("trino_ok") and objectscale_ok)),
        _edge("e-trino-db", "lakehouse", "db", "federated SQL", "query", bool(lakehouse.get("trino_ok"))),
        _edge("e-hdfs", "lakehouse", "hadoop", "parallel layer", "parallel", bool(hdfs_ok)),
        _edge("e-monitor-docker", "command", "docker", "monitor", "infra", True),
        _edge("e-monitor-gpu", "command", "gpu", "LLM", "infra", bool(gpu.get("ok"))),
    ]

    topologies = build_all_topologies(
        topology_nodes,
        topology_edges,
        snap,
        pipeline_active=s3_flow_ok,
        connectors=connectors,
    )

    return {
        "ts": snap.get("ts"),
        "zones": zones,
        "topology": topologies["architecture"],
        "topologies": topologies,
        "gpu": {
            "level": "ok" if gpu.get("ok") and gpu.get("inference_active") else ("warn" if gpu.get("ok") else "down"),
            "model": gpu.get("active_model"),
            "inference_active": gpu.get("inference_active"),
            "gpu_count": gpu.get("gpu_count", 0),
            "avg_util": round(
                sum(g.get("util_gpu", 0) for g in gpu.get("gpus", [])) / max(len(gpu.get("gpus", [])), 1),
                1,
            ),
            "gpus": gpu.get("gpus", []),
        },
        "totals": {
            "apps_running": sum(z["running"] for z in zones if z["id"] not in ("hadoop",)) + (hadoop.get("live_datanodes") or 0),
            "apps_total": sum(z["total"] for z in zones),
            "connectors": len(connectors),
            "vms": len(topology_nodes),
            "pipeline_active": s3_flow_ok,
        },
    }