# +----------------------------------------------------------------------
# | WaitAdmin(fastapi)快速开发后台管理系统
# +----------------------------------------------------------------------
# | 欢迎阅读学习程序代码,建议反馈是我们前进的动力
# | 程序完全开源可支持商用,允许去除界面版权信息
# | gitee: https://gitee.com/wafts/waitadmin-python
# | github: https://github.com/topwait/waitadmin-python
# | 官方网站: https://www.waitadmin.cn
# | WaitAdmin团队版权所有并拥有最终解释权
# +----------------------------------------------------------------------
# | Author: WaitAdmin Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import typing
import asyncio
from abc import ABC
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp


def init_middlewares(app: FastAPI):
 # 跨域中间件
 cors_middleware: typing.Type[any] = CORSMiddleware
 app.add_middleware(
 cors_middleware,
 allow_credentials=True,
 allow_origins=["*"],
 allow_methods=["*"],
 allow_headers=["*"]
 )

 # 超时中间件
 timeout_middleware: typing.Type[any] = TimeoutMiddleware
 app.add_middleware(timeout_middleware, timeout=500)


class TimeoutMiddleware(BaseHTTPMiddleware, ABC):
 def __init__(self, app: ASGIApp, timeout: int = 15):
 super().__init__(app)
 self.timeout = timeout

 async def dispatch(self, request, call_next):
 try:
 return await asyncio.wait_for(call_next(request), timeout=self.timeout)
 except asyncio.TimeoutError:
 return JSONResponse({"code": 1, "msg": "Request timeout", "data": []})