from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from websockets import connect as websockets_connect
from typing import Dict
app = FastAPI()
clients: Dict[str, WebSocket] = {}
@app.get("/clients")
def get_clients():
    return list(clients.keys())
@app.get("/send/{name}/{id}/{message}")
async def sendmsg(name,id,message):
    async with websockets_connect("ws://localhost:7860"+name) as websocket:
        await websocket.send(f"{id}:{message}")
        return await websocket.recv()

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()
    clients[client_id] = websocket
    print(f"Client {client_id} connected.")
    try:
        while True:
            data = await websocket.receive_text()
            message_parts = data.split(":", 1)
            if len(message_parts) == 2:
                target_client_id = message_parts[0].strip()
                msg_content = message_parts[1].strip()
                if target_client_id in clients:
                    target_websocket = clients[target_client_id]
                    await target_websocket.send_text(f"{client_id}: {msg_content}")
                else:
                    await websocket.send_text(f"Error: Client {target_client_id} not found.")
            else:
                await websocket.send_text("Error: Invalid message format. Use 'target_id:message'.")
    except WebSocketDisconnect:
        del clients[client_id]
        print(f"Client disconnected: {client_id}")
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)