from fastapi import FastAPI, WebSocket
import asyncio
import logging

app = FastAPI()

# Store active WebSocket connections
connected_clients = []

@app.websocket("/ws/logs")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    connected_clients.append(websocket)
    try:
        while True:
            await websocket.receive_text()
    except:
        connected_clients.remove(websocket)

# Custom logging handler to send logs via WebSocket
class WebSocketHandler(logging.Handler):
    def emit(self, record):
        log_message = self.format(record)
        # Send to all connected clients
        asyncio.create_task(broadcast_log(log_message))

async def broadcast_log(message):
    for client in connected_clients[:]:  # Copy list to avoid modification during iteration
        try:
            await client.send_json({
                "type": "log",
                "message": message,
                "timestamp": asyncio.get_event_loop().time()
            })
        except:
            connected_clients.remove(client)

# Set up logging
logger = logging.getLogger("app")
logger.setLevel(logging.INFO)
ws_handler = WebSocketHandler()
logger.addHandler(ws_handler)

@app.get("/test")
async def test_endpoint():
    logger.info("This is a test log message from backend!")
    return {"message": "Check your browser console"}