"""
LogiFlow Command Center — real-time JARVIS-style mission control.

A read-only tap on the live infrastructure:
  - subscribes to `logistics.>`  (every real agent-to-agent message on the bus)
  - subscribes to `telemetry.>`  (agent state + task lifecycle events)
and streams everything to the browser over WebSocket.

Runs as a separate process — it never interferes with the agents:
plain (core) NATS subscriptions only, no durable consumers, no JetStream acks.

Run:
    python main.py --command-center          # http://localhost:8600
    python dashboard/command_center.py       # same, standalone
"""
import asyncio
import json
import sys
import time
from collections import deque
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Set

# Allow running as `python dashboard/command_center.py`
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from dotenv import load_dotenv
load_dotenv()

import nats
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse

from config.system_config import NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD, AGENT_CONFIG
from core.logger import logger

STATIC_DIR = Path(__file__).resolve().parent / "static"
PORT = 8600

KNOWN_AGENTS = [
    "JARVIS", "ORDER_AGENT", "DISPATCH_AGENT", "FLEET_AGENT",
    "HUB_AGENT", "CUSTOMER_AGENT", "EXCEPTION_AGENT", "ROUTE_OPTIMIZER",
]

AGENT_DESCRIPTIONS = {cfg["name"]: cfg["description"] for cfg in AGENT_CONFIG.values()}


class CommandCenterState:
    """Aggregated live view of the agent system, fed by NATS events."""

    def __init__(self):
        self.bus_connected = False
        self.agents: Dict[str, Dict[str, Any]] = {
            a: {"agent_id": a, "status": "offline", "current_task": None,
                "tasks_completed": 0, "tasks_failed": 0, "last_seen": None,
                "msgs_sent": 0, "msgs_received": 0,
                "description": AGENT_DESCRIPTIONS.get(a, "")}
            for a in KNOWN_AGENTS
        }
        self.agent_recent: Dict[str, deque] = {a: deque(maxlen=40) for a in KNOWN_AGENTS}
        self.agent_durations: Dict[str, deque] = {a: deque(maxlen=50) for a in KNOWN_AGENTS}
        self.events: deque = deque(maxlen=300)
        self.totals = {"messages": 0, "tasks_completed": 0, "tasks_failed": 0, "exceptions": 0}
        self.msg_times: deque = deque(maxlen=5000)  # unix ts of bus messages, for rate

    def rate_per_min(self) -> int:
        cutoff = time.time() - 60
        return sum(1 for t in self.msg_times if t >= cutoff)

    def snapshot(self) -> Dict[str, Any]:
        return {
            "type": "snapshot",
            "bus_connected": self.bus_connected,
            "agents": self.agents,
            "events": list(self.events),
            "totals": self.totals,
            "rate_per_min": self.rate_per_min(),
        }


state = CommandCenterState()
websockets: Set[WebSocket] = set()
app = FastAPI(title="LogiFlow Command Center")


async def broadcast(event: Dict[str, Any]):
    """Push one event to every connected browser."""
    dead = []
    for ws in websockets:
        try:
            await ws.send_text(json.dumps(event))
        except Exception:
            dead.append(ws)
    for ws in dead:
        websockets.discard(ws)


def _compact(payload: Any, limit: int = 240) -> str:
    try:
        s = json.dumps(payload, default=str)
    except Exception:
        s = str(payload)
    return s[:limit] + ("…" if len(s) > limit else "")


async def _on_bus_message(msg):
    """Any real agent message on `logistics.>` (direct or broadcast)."""
    now = time.time()
    state.totals["messages"] += 1
    state.msg_times.append(now)
    try:
        data = json.loads(msg.data.decode())
    except Exception:
        data = {"raw": msg.data.decode(errors="replace")}

    msg_type = data.get("message_type", msg.subject.split(".")[-1])
    if msg_type == "EXCEPTION_DETECTED":
        state.totals["exceptions"] += 1

    event = {
        "type": "bus",
        "ts": data.get("timestamp", datetime.now().isoformat()),
        "sender": data.get("sender", "?"),
        "recipient": data.get("recipient", msg.subject.split(".")[-1] if ".direct." in msg.subject else "ALL"),
        "msg_type": msg_type,
        "summary": _compact(data.get("payload", data)),
        "correlation_id": data.get("correlation_id"),
    }
    for role, key in (("sender", "msgs_sent"), ("recipient", "msgs_received")):
        aid = event[role]
        if aid in state.agents:
            state.agents[aid][key] = state.agents[aid].get(key, 0) + 1
            state.agent_recent.setdefault(aid, deque(maxlen=40)).append(event)
    state.events.append(event)
    await broadcast(event)


async def _on_telemetry(msg):
    """Agent state / task lifecycle events on `telemetry.>`."""
    try:
        data = json.loads(msg.data.decode())
    except Exception:
        return

    kind = data.get("kind")
    if kind == "agent":
        agent_id = data.get("agent_id")
        if not agent_id:
            return
        entry = state.agents.setdefault(agent_id, {
            "msgs_sent": 0, "msgs_received": 0,
            "description": AGENT_DESCRIPTIONS.get(agent_id, ""),
        })
        entry.update({
            "agent_id": agent_id,
            "status": data.get("status", "idle"),
            "current_task": data.get("current_task"),
            "tasks_completed": data.get("tasks_completed", 0),
            "tasks_failed": data.get("tasks_failed", 0),
            "last_seen": data.get("ts"),
        })
        await broadcast({"type": "agent", **entry})

    elif kind == "task":
        if data.get("status") == "completed":
            state.totals["tasks_completed"] += 1
        elif data.get("status") == "failed":
            state.totals["tasks_failed"] += 1
        event = {
            "type": "task",
            "ts": data.get("ts", datetime.now().isoformat()),
            "agent_id": data.get("agent_id"),
            "task_id": data.get("task_id"),
            "task_type": data.get("task_type"),
            "status": data.get("status"),
            "error": data.get("error"),
            "duration_ms": data.get("duration_ms"),
        }
        aid = event["agent_id"]
        if aid:
            state.agent_recent.setdefault(aid, deque(maxlen=40)).append(event)
            if event["duration_ms"] is not None:
                state.agent_durations.setdefault(aid, deque(maxlen=50)).append(event["duration_ms"])
        state.events.append(event)
        await broadcast(event)


async def nats_tap():
    """Connect to NATS (with retry forever) and hold plain subscriptions."""
    while True:
        try:
            nc = await nats.connect(
                servers=[f"nats://{NATS_HOST}:{NATS_PORT}"],
                user=NATS_USER,
                password=NATS_PASSWORD,
                max_reconnect_attempts=-1,
                disconnected_cb=_bus_down,
                reconnected_cb=_bus_up,
            )
            await nc.subscribe("logistics.>", cb=_on_bus_message)
            await nc.subscribe("telemetry.>", cb=_on_telemetry)
            await _bus_up()
            logger.info(f"Command Center tapped into NATS at {NATS_HOST}:{NATS_PORT}")
            while nc.is_connected or nc.is_reconnecting:
                await asyncio.sleep(2)
            await _bus_down()
        except Exception as e:
            logger.warning(f"Command Center NATS connect failed: {e} — retrying in 5s")
            await _bus_down()
            await asyncio.sleep(5)


async def _bus_up():
    state.bus_connected = True
    await broadcast({"type": "sys", "bus_connected": True})


async def _bus_down():
    if state.bus_connected:
        logger.warning("Command Center lost NATS connection")
    state.bus_connected = False
    await broadcast({"type": "sys", "bus_connected": False})


async def rate_ticker():
    """Push the message rate + mark stale agents offline every 2s."""
    while True:
        await asyncio.sleep(2)
        now = datetime.now()
        for entry in state.agents.values():
            last = entry.get("last_seen")
            if last and entry.get("status") != "offline":
                try:
                    age = (now - datetime.fromisoformat(last)).total_seconds()
                    if age > 15:
                        entry["status"] = "offline"
                        await broadcast({"type": "agent", **entry})
                except ValueError:
                    pass
        if websockets:
            await broadcast({
                "type": "rate",
                "rate_per_min": state.rate_per_min(),
                "totals": state.totals,
            })


@app.on_event("startup")
async def startup():
    asyncio.create_task(nats_tap())
    asyncio.create_task(rate_ticker())


@app.get("/")
async def index():
    return FileResponse(STATIC_DIR / "command_center.html")


@app.get("/api/state")
async def api_state():
    return state.snapshot()


@app.get("/api/agent/{agent_id}")
async def api_agent(agent_id: str):
    """Full dossier for one agent — used by the UI detail panel."""
    entry = state.agents.get(agent_id)
    if not entry:
        return {"error": "unknown agent"}
    durations = list(state.agent_durations.get(agent_id, []))
    return {
        **entry,
        "avg_task_ms": round(sum(durations) / len(durations)) if durations else None,
        "recent": list(state.agent_recent.get(agent_id, []))[-30:],
    }


@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
    await ws.accept()
    websockets.add(ws)
    await ws.send_text(json.dumps(state.snapshot()))
    try:
        while True:
            await ws.receive_text()  # keepalive pings from the browser
    except WebSocketDisconnect:
        pass
    finally:
        websockets.discard(ws)


def run(port: int = PORT):
    logger.info(f"LogiFlow Command Center → http://localhost:{port}")
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="warning")


if __name__ == "__main__":
    run()