"""
OPC WebSocket Router - Real-time updates for tasks and agent executions
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from typing import Set, Any, Dict
from datetime import datetime
import json
import asyncio
import logging
from auth_service import get_current_user_ws

logger = logging.getLogger(__name__)

router = APIRouter()

# Active WebSocket connections
active_connections: Set[WebSocket] = set()


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"""

    def __init__(self):
        self.active_connections: Set[WebSocket] = set()

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

    def disconnect(self, websocket: WebSocket):
        """Remove a connection"""
        self.active_connections.discard(websocket)
        logger.info(f"WebSocket disconnected. 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):
    """WebSocket endpoint for OPC real-time updates"""
    await manager.connect(websocket)

    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)