from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from services.upload_file.upload_ws import job_state
from api.routers.ws.manager import manager

router = APIRouter()

@router.websocket("/ws/test")
async def ws_test(ws: WebSocket):
    await ws.accept()
    await ws.send_text("OK")


@router.websocket("/ws/job/{job_id}")
async def ws_job(job_id: str, ws: WebSocket):
    await manager.connect(job_id, ws)

    # kirim progress terakhir (jika reconnect)
    if job_id in job_state:
        await ws.send_json(job_state[job_id])

    try:
        while True:
            await ws.receive_text()  # keep alive
    except WebSocketDisconnect:
        manager.disconnect(job_id, ws)