"""WebSocket adapter — bidirectional communication at /ws.

Supports tool invocation, heartbeat, and agent orchestration over
persistent WebSocket connections. Reference: src/console_ws/server.py.
"""

from __future__ import annotations

import json
import logging
from typing import Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from ..context import RequestContext
from ..registry import ToolRegistry

logger = logging.getLogger(__name__)

router = APIRouter(tags=["WebSocket"])

_registry: ToolRegistry | None = None
_tenant_registry: Any = None

# Track connected agents for observability
_connected_agents: dict[str, WebSocket] = {}


def init_ws(registry: ToolRegistry, tenant_registry: Any = None) -> None:
    """Wire the shared ToolRegistry into this adapter."""
    global _registry, _tenant_registry
    _registry = registry
    _tenant_registry = tenant_registry


@router.websocket("/ws")
async def ws_handler(websocket: WebSocket) -> None:
    """Handle WebSocket connections for tool invocation.

    Protocol:
    1. Client connects → server accepts
    2. Client sends auth message: {"type": "auth", "token": "...", "tenant_id": "..."}
    3. Server responds: {"type": "auth_ok", "agent_id": "..."}
    4. Client sends tool calls: {"type": "tool_call", "id": "...", "tool": "...", "params": {...}}
    5. Server responds: {"type": "tool_result", "id": "...", "result": {...}}
    6. Heartbeat: {"type": "heartbeat"} → {"type": "heartbeat_ack"}
    """
    await websocket.accept()
    agent_id: str | None = None

    try:
        # Authenticate via first message
        from ..middleware.auth import authenticate_ws
        context = await authenticate_ws(websocket, _tenant_registry)
        agent_id = context.request_id

        _connected_agents[agent_id] = websocket
        await websocket.send_json({
            "type": "auth_ok",
            "agent_id": agent_id,
            "tools": _registry.dynamic.total_count if _registry else 0,
        })

        logger.info("WebSocket agent connected: %s (tenant=%s)", agent_id, context.tenant_id)

        # Message loop
        async for raw in websocket.iter_text():
            try:
                message = json.loads(raw)
            except json.JSONDecodeError:
                await websocket.send_json({
                    "type": "error",
                    "message": "Invalid JSON",
                })
                continue

            msg_type = message.get("type", "")

            if msg_type == "heartbeat":
                await websocket.send_json({"type": "heartbeat_ack"})

            elif msg_type == "tool_call":
                msg_id = message.get("id", "")
                tool = message.get("tool", "")
                params = message.get("params", {})

                if not tool:
                    await websocket.send_json({
                        "type": "error",
                        "id": msg_id,
                        "message": "Missing 'tool' field",
                    })
                    continue

                result = await _registry.invoke(tool, params, context)
                await websocket.send_json({
                    "type": "tool_result",
                    "id": msg_id,
                    "result": result,
                })

            elif msg_type == "discover":
                domain = message.get("domain", "")
                query = message.get("query", "")
                if query:
                    tools = _registry.search_tools(query)
                elif domain:
                    tools = _registry.get_tools_by_domain(domain)
                else:
                    tools = _registry.get_all_domains()
                await websocket.send_json({
                    "type": "discover_result",
                    "id": message.get("id", ""),
                    "data": tools,
                })

            else:
                await websocket.send_json({
                    "type": "error",
                    "message": f"Unknown message type: {msg_type}",
                })

    except WebSocketDisconnect:
        logger.info("WebSocket agent disconnected: %s", agent_id)
    except Exception as exc:
        logger.exception("WebSocket error for agent %s: %s", agent_id, exc)
    finally:
        if agent_id and agent_id in _connected_agents:
            del _connected_agents[agent_id]


@router.get("/ws/agents", tags=["WebSocket"])
async def list_connected_agents():
    """List currently connected WebSocket agents."""
    return {
        "connected": len(_connected_agents),
        "agent_ids": list(_connected_agents.keys()),
    }
