"""
OPC WebSocket Router - Real-time updates for tasks and agent executions
"""

import logging
from datetime import datetime
from typing import Any

from fastapi import APIRouter, HTTPException, Query, WebSocket, WebSocketDisconnect, status

import auth_service as auth

logger = logging.getLogger(__name__)

router = APIRouter()


def serialize_datetime(obj: Any) -> Any:
    """Recursively serialize datetime objects to ISO format strings"""
    if isinstance(obj, datetime):
        return obj.isoformat()
    elif isinstance(obj, dict):
        return {key: serialize_datetime(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [serialize_datetime(item) for item in obj]
    else:
        return obj


class ConnectionManager:
    """Manages WebSocket connections with per-user tracking"""

    def __init__(self):
        self.active_connections: dict[WebSocket, str] = {}

    async def connect(self, websocket: WebSocket, username: str):
        """Accept and register a new connection"""
        await websocket.accept()
        self.active_connections[websocket] = username
        logger.info(f"WebSocket connected: {username}. Total connections: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        """Remove a connection"""
        username = self.active_connections.pop(websocket, "unknown")
        logger.info(f"WebSocket disconnected: {username}. Total connections: {len(self.active_connections)}")

    async def broadcast(self, message: dict):
        """Broadcast message to all connected clients"""
        disconnected = set()
        for connection in self.active_connections:
            try:
                await connection.send_json(message)
            except Exception as e:
                logger.error(f"Failed to send message: {e}")
                disconnected.add(connection)

        # Clean up disconnected clients
        for conn in disconnected:
            self.disconnect(conn)


manager = ConnectionManager()


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, token: str = Query(None)):
    """WebSocket endpoint for OPC real-time updates. Requires JWT token via ?token= query param."""
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Missing authentication token")
        return

    try:
        user = auth.user_from_access_token(token, include_auth_header=False)
    except HTTPException:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid authentication token")
        return

    await manager.connect(websocket, user.username)

    try:
        while True:
            # Keep connection alive and receive any client messages
            data = await websocket.receive_text()

            # Echo back for ping/pong
            if data == "ping":
                await websocket.send_text("pong")

    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
        manager.disconnect(websocket)


async def broadcast_task_update(task_id: int, action: str, task_data: dict):
    """Broadcast task update to all connected clients"""
    # Serialize all datetime objects recursively
    serialized_data = serialize_datetime(task_data)

    message = {
        "type": "task_update",
        "task_id": task_id,
        "action": action,  # created, updated, moved, deleted
        "data": serialized_data,
        "timestamp": datetime.utcnow().isoformat(),
    }
    await manager.broadcast(message)


async def broadcast_agent_execution(execution_id: int, status: str, execution_data: dict):
    """Broadcast agent execution update"""
    # Serialize all datetime objects recursively
    serialized_data = serialize_datetime(execution_data)

    message = {
        "type": "agent_execution",
        "execution_id": execution_id,
        "status": status,  # pending, running, completed, failed, pending_approval
        "data": serialized_data,
        "timestamp": datetime.utcnow().isoformat(),
    }
    await manager.broadcast(message)


async def broadcast_agent_status(agent_id: str, status: str, message_text: str = ""):
    """Broadcast agent status update"""
    message = {
        "type": "agent_status",
        "agent_id": agent_id,
        "status": status,  # idle, working, error
        "message": message_text,
    }
    await manager.broadcast(message)