from functools import partial
from typing import Optional

import fastapi
from fastapi import Request
from pydantic import BaseModel, Field
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import Response

app = fastapi.FastAPI()


TRANSLATIONS = {
    "zh_CN": {"操作成功": "操作成功"},
    "en_US": {"操作成功": "Successful operation."},
}

translations = TRANSLATIONS.get("zh_CN")


def set_locale(locale: str):
    global translations
    translations = TRANSLATIONS.get(locale) or TRANSLATIONS.get("zh_CN")


def _(msg: str):
    return translations.get(msg)


class InternationalizationMiddleware(BaseHTTPMiddleware):

    async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        lang_code: Optional[str] = request.headers.get("Accept-Language", None)
        set_locale(lang_code)
        response: _StreamingResponse = await call_next(request)  # type: ignore
        return response


app.add_middleware(InternationalizationMiddleware)


class ResponseModel(BaseModel):
    code: int = 200
    message: str = Field(default_factory=partial(_, msg="操作成功"))


@app.get('/not_work', response_model=ResponseModel)
async def func():
    return {"code": 200}


@app.get('/work', response_model=ResponseModel)
async def func():
    return ResponseModel()

if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, port=8089)