from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List
import asyncio

app = FastAPI()

# Store active WebSocket connections
active_connections: List[WebSocket] = []

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # Accept the incoming WebSocket connection
    await websocket.accept()
    active_connections.append(websocket)  # Add new connection to active list
    print(f"New client connected! Total clients: {len(active_connections)}")

    try:
        while True:
            # Wait for a message from the client
            data = await websocket.receive_text()
            print(f"Received message: {data}")

            # Broadcast the message to all connected clients
            await broadcast_message(f"Broadcast: {data}")
            
            # Optionally echo back the message to the client who sent it
            await websocket.send_text(f"Echo: {data}")
    except WebSocketDisconnect:
        # Handle disconnection of the client
        active_connections.remove(websocket)
        print(f"Client disconnected! Total clients: {len(active_connections)}")

async def broadcast_message(message: str):
    """Helper function to broadcast a message to all active WebSockets."""
    # Send the message to each client in the active connections list
    for connection in active_connections:
        try:
            await connection.send_text(message)
        except Exception as e:
            # Handle connection errors (e.g., client closed connection)
            print(f"Error sending message: {e}")
            active_connections.remove(connection)