atc-agents/api/main.py

"""ATC Command Center API — FastAPI backend."""

from __future__ import annotations

import asyncio
import json
import os
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any

import httpx
import redis.asyncio as aioredis
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from sqlalchemy import Column, DateTime, String, Text, create_engine, select
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker

REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
DOCKHAND_URL = os.getenv("DOCKHAND_URL", "http://10.0.21.45:8082")
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:////data/atc-agents.db")

AGENTS = [
    {
        "id": "etl-guardian",
        "name": "ETL Guardian",
        "color": "#00f0ff",
        "zone": "etl",
        "role": "Airflow, Kafka, Debezium, S3 pipeline",
    },
    {
        "id": "lakehouse-ops",
        "name": "Lakehouse Ops",
        "color": "#ff00aa",
        "zone": "lakehouse",
        "role": "Spark, Trino, Iceberg",
    },
    {
        "id": "data-custodian",
        "name": "Data Custodian",
        "color": "#ffaa00",
        "zone": "db",
        "role": "PostgreSQL, MySQL, Mongo, Cassandra, Neo4j",
    },
    {
        "id": "hadoop-ranger",
        "name": "Hadoop Ranger",
        "color": "#39ff14",
        "zone": "hadoop",
        "role": "HDFS, YARN cluster",
    },
    {
        "id": "infra-sentinel",
        "name": "Infra Sentinel",
        "color": "#b366ff",
        "zone": "docker",
        "role": "Docker, Proxmox, monitoring",
    },
]

ZONES = [
    {"id": "docker", "label": "DOCKER RACK", "x": 8, "color": "#b366ff"},
    {"id": "db", "label": "DB VAULT", "x": 28, "color": "#ffaa00"},
    {"id": "lakehouse", "label": "LAKEHOUSE HUB", "x": 50, "color": "#ff00aa"},
    {"id": "hadoop", "label": "HADOOP CLUSTER", "x": 72, "color": "#39ff14"},
    {"id": "etl", "label": "ETL PIPE", "x": 92, "color": "#00f0ff"},
]

INTENT_KEYWORDS: dict[str, list[str]] = {
    "data-custodian": ["database", "db", "postgres", "mysql", "mongo", "cassandra", "neo4j", "sql"],
    "lakehouse-ops": ["trino", "spark", "lakehouse", "iceberg", "query"],
    "hadoop-ranger": ["hadoop", "hdfs", "yarn", "datanode"],
    "infra-sentinel": ["docker", "container", "vm", "proxmox", "infra", "grafana"],
    "etl-guardian": ["airflow", "dag", "debezium", "kafka", "connector", "etl", "pipeline", "s3"],
}


class Base(DeclarativeBase):
    pass


class FeedEntry(Base):
    __tablename__ = "feed"

    id = Column(String, primary_key=True)
    ts = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    agent_id = Column(String)
    level = Column(String, default="info")
    message = Column(Text)


class Approval(Base):
    __tablename__ = "approvals"

    id = Column(String, primary_key=True)
    ts = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    agent_id = Column(String)
    action = Column(Text)
    reason = Column(Text)
    status = Column(String, default="pending")


engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=engine)
Base.metadata.create_all(engine)

redis_client: aioredis.Redis | None = None
ws_clients: set[WebSocket] = set()


class PromptRequest(BaseModel):
    message: str = Field(min_length=1, max_length=2000)


class ApprovalDecision(BaseModel):
    approved: bool


def route_agent(message: str) -> str:
    lower = message.lower()
    scores = {aid: sum(1 for kw in kws if kw in lower) for aid, kws in INTENT_KEYWORDS.items()}
    best = max(scores, key=scores.get)
    if scores[best] == 0:
        return "infra-sentinel"
    return best


async def publish_event(event: dict[str, Any]) -> None:
    payload = json.dumps(event, default=str)
    if redis_client:
        await redis_client.publish("ops", payload)
    dead = []
    for ws in ws_clients:
        try:
            await ws.send_text(payload)
        except Exception:
            dead.append(ws)
    for ws in dead:
        ws_clients.discard(ws)


def add_feed(agent_id: str, message: str, level: str = "info") -> dict:
    entry_id = str(uuid.uuid4())[:8]
    with SessionLocal() as db:
        row = FeedEntry(id=entry_id, agent_id=agent_id, message=message, level=level)
        db.add(row)
        db.commit()
    return {
        "id": entry_id,
        "ts": datetime.now(timezone.utc).isoformat(),
        "agent_id": agent_id,
        "message": message,
        "level": level,
    }


async def dockhand_env_containers(env_id: int) -> list[dict]:
    try:
        async with httpx.AsyncClient(timeout=8.0) as client:
            r = await client.get(f"{DOCKHAND_URL}/api/containers", params={"env": env_id})
            r.raise_for_status()
            return r.json()
    except Exception:
        return []


async def probe_url(url: str) -> bool:
    try:
        async with httpx.AsyncClient(timeout=4.0, verify=False) as client:
            r = await client.get(url)
            return r.status_code < 500
    except Exception:
        return False


async def collect_status() -> dict[str, Any]:
    db_containers = await dockhand_env_containers(5)
    db_running = sum(1 for c in db_containers if c.get("state") == "running")
    db_total = len(db_containers) or 6

    lake_containers = await dockhand_env_containers(9)
    lake_running = sum(1 for c in lake_containers if c.get("state") == "running")
    lake_total = len(lake_containers) or 6

    docker_containers = await dockhand_env_containers(1)
    docker_running = sum(1 for c in docker_containers if c.get("state") == "running")

    hdfs_ok = await probe_url("http://10.0.21.61:9870")
    kafka_ok = await probe_url("http://10.0.21.36:9000")
    airflow_ok = await probe_url("http://10.0.21.55:8080")

    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"

    return {
        "ts": datetime.now(timezone.utc).isoformat(),
        "domains": {
            "docker": {"level": "ok" if docker_running >= 5 else "warn", "label": f"{docker_running} containers", "running": docker_running},
            "databases": {"level": level(db_running, db_total), "label": f"{db_running}/{db_total} up", "running": db_running, "total": db_total},
            "lakehouse": {"level": level(lake_running, lake_total), "label": f"{lake_running}/{lake_total} up", "running": lake_running, "total": lake_total},
            "hadoop": {"level": "ok" if hdfs_ok else "warn", "label": "NN up" if hdfs_ok else "NN check"},
            "etl": {"level": "ok" if kafka_ok and airflow_ok else "warn", "label": "Kafka+Airflow"},
        },
        "kafka_ok": kafka_ok,
        "airflow_ok": airflow_ok,
        "hdfs_ok": hdfs_ok,
    }


async def run_agent_task(agent_id: str, message: str, prompt_id: str) -> str:
    zone = next(a["zone"] for a in AGENTS if a["id"] == agent_id)
    await publish_event({"type": "agent_dispatch", "agent_id": agent_id, "zone": zone, "prompt_id": prompt_id})
    await asyncio.sleep(0.8)
    await publish_event({"type": "agent_fetch", "agent_id": agent_id, "zone": zone, "prompt_id": prompt_id})

    status = await collect_status()
    answer_parts = [f"**{next(a['name'] for a in AGENTS if a['id'] == agent_id)}** reporting:"]

    if agent_id == "data-custodian":
        db = status["domains"]["databases"]
        containers = await dockhand_env_containers(5)
        names = ", ".join(f"{c['name']}:{c.get('state','?')}" for c in containers[:8])
        answer_parts.append(f"Databases {db['label']}. Containers: {names or 'unreachable'}.")
    elif agent_id == "etl-guardian":
        answer_parts.append(
            f"Kafka UI: {'OK' if status['kafka_ok'] else 'DOWN'}. "
            f"Airflow: {'OK' if status['airflow_ok'] else 'DOWN'}. "
            f"Lakehouse {status['domains']['lakehouse']['label']}."
        )
    elif agent_id == "lakehouse-ops":
        answer_parts.append(f"Lakehouse stack {status['domains']['lakehouse']['label']}. Trino at 10.0.21.50:8089.")
    elif agent_id == "hadoop-ranger":
        answer_parts.append(f"HDFS NameNode: {'reachable' if status['hdfs_ok'] else 'unreachable'} on 10.0.21.61:9870.")
    else:
        answer_parts.append(
            f"Docker {status['domains']['docker']['label']}. "
            f"Overall lab health snapshot collected."
        )

    answer = " ".join(answer_parts)
    await asyncio.sleep(0.6)
    await publish_event({"type": "agent_return", "agent_id": agent_id, "zone": zone, "prompt_id": prompt_id})
    feed = add_feed(agent_id, f"Prompt answered: {message[:80]}", "info")
    await publish_event({"type": "feed", "entry": feed})
    await publish_event({"type": "prompt_result", "prompt_id": prompt_id, "agent_id": agent_id, "answer": answer})
    return answer


async def heartbeat_loop() -> None:
    while True:
        try:
            status = await collect_status()
            await publish_event({"type": "status", "data": status})
            for domain, info in status["domains"].items():
                if info["level"] == "down":
                    agent = "data-custodian" if domain == "databases" else "infra-sentinel"
                    feed = add_feed(agent, f"Alert: {domain} is DOWN ({info['label']})", "warn")
                    await publish_event({"type": "feed", "entry": feed})
        except Exception as exc:
            await publish_event({"type": "error", "message": str(exc)})
        await asyncio.sleep(60)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global redis_client
    redis_client = aioredis.from_url(REDIS_URL, decode_responses=True)
    task = asyncio.create_task(heartbeat_loop())
    add_feed("infra-sentinel", "ATC Command Center API online", "info")
    yield
    task.cancel()
    if redis_client:
        await redis_client.close()


app = FastAPI(title="ATC Command Center API", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/api/health")
async def health():
    return {"ok": True, "ts": datetime.now(timezone.utc).isoformat()}


@app.get("/api/status")
async def get_status():
    return await collect_status()


@app.get("/api/agents")
async def get_agents():
    return {"agents": AGENTS, "zones": ZONES}


@app.get("/api/feed")
async def get_feed(limit: int = 50):
    with SessionLocal() as db:
        rows = db.execute(select(FeedEntry).order_by(FeedEntry.ts.desc()).limit(limit)).scalars().all()
    return {
        "entries": [
            {
                "id": r.id,
                "ts": r.ts.isoformat() if r.ts else None,
                "agent_id": r.agent_id,
                "message": r.message,
                "level": r.level,
            }
            for r in rows
        ]
    }


@app.get("/api/approvals")
async def get_approvals():
    with SessionLocal() as db:
        rows = db.execute(select(Approval).where(Approval.status == "pending")).scalars().all()
    return {
        "approvals": [
            {
                "id": r.id,
                "ts": r.ts.isoformat() if r.ts else None,
                "agent_id": r.agent_id,
                "action": r.action,
                "reason": r.reason,
                "status": r.status,
            }
            for r in rows
        ]
    }


@app.post("/api/approvals/{approval_id}/decide")
async def decide_approval(approval_id: str, body: ApprovalDecision):
    with SessionLocal() as db:
        row = db.get(Approval, approval_id)
        if not row:
            return {"error": "not found"}
        row.status = "approved" if body.approved else "denied"
        db.commit()
        agent_id = row.agent_id
        action = row.action
    msg = f"Approval {'approved' if body.approved else 'denied'}: {action}"
    feed = add_feed(agent_id, msg, "info" if body.approved else "warn")
    await publish_event({"type": "feed", "entry": feed})
    await publish_event({"type": "approval_update", "id": approval_id, "status": row.status})
    return {"ok": True, "status": row.status}


@app.post("/api/prompt")
async def post_prompt(body: PromptRequest):
    prompt_id = str(uuid.uuid4())[:8]
    agent_id = route_agent(body.message)
    add_feed(agent_id, f"Prompt received: {body.message}", "info")
    asyncio.create_task(run_agent_task(agent_id, body.message, prompt_id))
    return {"prompt_id": prompt_id, "agent_id": agent_id, "status": "dispatched"}


@app.websocket("/api/ws/ops")
async def ws_ops(websocket: WebSocket):
    await websocket.accept()
    ws_clients.add(websocket)
    try:
        status = await collect_status()
        await websocket.send_text(json.dumps({"type": "status", "data": status}, default=str))
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        ws_clients.discard(websocket)