from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import Response
from typing import Dict
import asyncio
import pickle

from paddle_run.core import TaskResult

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[WebSocket, asyncio.Queue] = {}
        self.idle_workers: asyncio.Queue = asyncio.Queue()

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[websocket] = asyncio.Queue()
        await self.idle_workers.put(websocket)

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            del self.active_connections[websocket]

    async def get_available_worker(self) -> WebSocket:
        worker = await self.idle_workers.get()
        if worker not in self.active_connections:
            return await self.get_available_worker()
        return worker

    def return_worker(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.idle_workers.put_nowait(websocket)

    async def dispatch_message(self, websocket: WebSocket, message: bytes):
        if websocket in self.active_connections:
            await self.active_connections[websocket].put(message)

    async def wait_for_response(self, websocket: WebSocket) -> bytes:
        queue = self.active_connections.get(websocket)
        if queue:
            return await queue.get()
        raise Exception("Connection lost")


manager = ConnectionManager()

@app.websocket("/ws/register")
async def websocket_endpoint(websocket: WebSocket, password: str):
    if password != "123456789":
        await websocket.close()
        return

    await manager.connect(websocket)

    try:
        while True:
            data = await websocket.receive_bytes()
            await manager.dispatch_message(websocket, data)

    except WebSocketDisconnect:
        manager.disconnect(websocket)
    except Exception as e:
        print(f"Connection error: {e}")
        manager.disconnect(websocket)


from fastapi import Request


@app.post("/run_proxy")
async def run_proxy(request: Request):
    body = await request.body()
    worker = await manager.get_available_worker()

    try:
        await worker.send_bytes(body)
        result_data = await manager.wait_for_response(worker)
        return Response(content=result_data, media_type="application/octet-stream")

    except Exception as e:
        return Response(content=pickle.dumps(TaskResult("err", False, "", f"Proxy Error: {str(e)}")),
                        media_type="application/octet-stream")
    finally:
        manager.return_worker(worker)


@app.get("/status")
def get_status():
    return {
        "total_connections": len(manager.active_connections),
        "idle_workers": manager.idle_workers.qsize()
    }