"""
OPC WebSocket Router - Real-time updates for tasks and agent executions
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from typing import Set
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()


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"""
    message = {
        "type": "task_update",
        "task_id": task_id,
        "action": action,  # created, updated, moved, deleted
        "data": task_data,
        "timestamp": task_data.get("updated_at")
    }
    await manager.broadcast(message)


async def broadcast_agent_execution(execution_id: int, status: str, execution_data: dict):
    """Broadcast agent execution update"""
    message = {
        "type": "agent_execution",
        "execution_id": execution_id,
        "status": status,  # pending, running, completed, failed, pending_approval
        "data": execution_data,
        "timestamp": execution_data.get("started_at") or execution_data.get("completed_at")
    }
    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)