import asyncio

import redis.asyncio as aioredis
from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.core.config import get_settings

router = APIRouter(tags=["ws"])
settings = get_settings()


@router.websocket("/ws/jobs/{resource_id}")
async def job_progress(websocket: WebSocket, resource_id: str):
    await websocket.accept()
    redis_client = aioredis.from_url(settings.redis_url)
    pubsub = redis_client.pubsub()
    await pubsub.subscribe(f"jobs:{resource_id}")

    async def forward():
        async for message in pubsub.listen():
            if message["type"] == "message":
                await websocket.send_text(message["data"].decode("utf-8"))

    forward_task = asyncio.create_task(forward())
    try:
        while True:
            # Keep the connection alive; client isn't expected to send data.
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        forward_task.cancel()
        await pubsub.unsubscribe(f"jobs:{resource_id}")
        await pubsub.close()
        await redis_client.close()