"""
WebSocket endpoints for real-time communication.
"""

import json
import logging
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from ..websocket import connection_manager, handle_websocket_message, websocket_connection


logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    client_id: Optional[str] = Query(None, description="Client identifier")
):
    """Main WebSocket endpoint for real-time data streaming."""
    async with websocket_connection(websocket, client_id):
        try:
            while True:
                # Receive message from client
                data = await websocket.receive_text()
                
                try:
                    message = json.loads(data)
                    await handle_websocket_message(websocket, message)
                except json.JSONDecodeError:
                    await connection_manager._send_message(websocket, {
                        "type": "error",
                        "message": "Invalid JSON message",
                        "timestamp": connection_manager._get_timestamp()
                    })
                except Exception as e:
                    logger.error(f"Error handling WebSocket message: {e}")
                    await connection_manager._send_message(websocket, {
                        "type": "error",
                        "message": "Message processing error",
                        "timestamp": connection_manager._get_timestamp()
                    })
                    
        except WebSocketDisconnect:
            logger.info(f"WebSocket client {client_id or 'anonymous'} disconnected")
        except Exception as e:
            logger.error(f"WebSocket endpoint error: {e}")


@router.get("/ws/stats")
async def websocket_stats():
    """Get WebSocket connection statistics."""
    return connection_manager.get_connection_stats()


@router.post("/ws/broadcast")
async def broadcast_message(message: dict):
    """Broadcast a message to all connected WebSocket clients."""
    await connection_manager.broadcast_to_all(message)
    return {"success": True, "message": "Message broadcasted to all clients"}


@router.post("/ws/broadcast/{channel}")
async def broadcast_to_channel(channel: str, message: dict):
    """Broadcast a message to clients subscribed to a specific channel."""
    await connection_manager.broadcast_to_channel(channel, message)
    return {"success": True, "message": f"Message broadcasted to {channel} channel"}