"""
FastAPI WebSocket router for exporting audit report rows.
"""

import asyncio

from box import Box
from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from ... import logger
from ...constants import RequestParams
from ...logic.reporting.spreadsheets import report_rows
from ...users.manage import verify_auth0_websocket

logger = logger.new_logger(__name__)


ROUTE: str = "/audits/ws/export"


def create_router(state: Box, *args, **kwargs):
    """Create FastAPI router."""

    app = APIRouter()

    @app.websocket("/audit_report_rows")
    async def audit_report_rows(websocket: WebSocket):
        """Generate a spreadsheet with the audit report rows. This is a
        spreadsheet that contains the rows with issues for sending it to the
        audited label together with the audit report.

        This endpoint uses websockets to allow for long-running operations, as
        audits with hundreds of thousands of rows can take some minutes to
        generate the report.
        """
        # Ensure the user is authenticated.
        json, _ = await verify_auth0_websocket(websocket)
        audit_group_id = int(json[RequestParams.AUDIT_GROUP_ID])

        async def keep_alive():
            start = asyncio.get_event_loop().time()
            while True:
                # Stop keep-alive after 15 minutes (900 seconds)
                if asyncio.get_event_loop().time() - start > 900:
                    break
                # pylint: disable=broad-except
                # noinspection PyBroadException
                try:
                    await websocket.send_text("(websocket keep-alive ping)")
                except Exception:
                    break
                await asyncio.sleep(10)

        # Start the keep-alive task to prevent the WebSocket connection from
        # timing out due to inactivity.
        keep_alive_task = asyncio.create_task(keep_alive())

        try:
            await websocket.send_text("Generating report...")

            report_bytes = await report_rows.run(state.db, audit_group_id)

            await websocket.send_bytes(report_bytes)
            await websocket.send_text("Report generation completed.")

        except WebSocketDisconnect as ex:
            logger.info(f"WebSocket closed with code: {ex.code}")

        finally:
            keep_alive_task.cancel()

    return app