"""Gate business traffic during deployment and database recovery, including sockets."""
import asyncio
from pathlib import Path
from starlette.types import ASGIApp, Receive, Scope, Send
from starlette.responses import JSONResponse

MAINTENANCE_FILE = Path('/var/lib/replypier-updater/maintenance')


class SystemMaintenanceMiddleware:
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        path = scope.get('path', '')
        exempt = path == '/health' or (path == '/api/v1/system-updates' and scope.get('method') == 'GET')
        if scope['type'] in {'http', 'websocket'} and MAINTENANCE_FILE.exists() and not exempt:
            if scope['type'] == 'websocket':
                await send({'type': 'websocket.close', 'code': 1013})
            else:
                response = JSONResponse({'detail': '系统更新中，接待将在更新完成后恢复。'},
                                        status_code=503, headers={'Retry-After': '5'})
                await response(scope, receive, send)
            return
        await self.app(scope, receive, send)


async def wait_for_deployment() -> None:
    while MAINTENANCE_FILE.exists():
        await asyncio.sleep(1)