from typing import Dict
from typing import List
from typing import NoReturn
from typing import Optional


from aioredis import Redis
from aioredis import from_url
from fastapi import FastAPI
from starlette.requests import Request
from tortoise.contrib.fastapi import register_tortoise

from common.config import settings


def register_orm(app: FastAPI, modules: Optional[Dict[str, List[str]]] = None) -> NoReturn:
    """注册orm"""
    if modules is None:
        modules = {"models": ["models"]}
    register_tortoise(
        app,
        db_url=settings.MYSQL_DATABASE_URL,
        modules=modules
    )


def redis_pool() -> Dict[int, Redis]:
    """redis 16个db 每个创建池"""
    return {db: from_url(settings.REDIS_URL, db=db, max_connections=10, encoding="utf-8", decode_responses=True)
            for db in range(16)
            }


def close_redis_pool(pool: Dict[int, Redis]):
    for v in pool.values():
        v.close()


def get_redis_db(request: Request, db: int = 0) -> Redis:
    """获取到FastAPI中的缓存redis 字典获取 指定 db"""
    return request.app.state.redis_pool.get(db)


def register_redis(app: FastAPI):
    """redis 注册"""

    @app.on_event("startup")
    async def create_redis():
        app.state.redis_pool = redis_pool()

    @app.on_event("shutdown")
    async def close_redis():
        close_redis_pool(app.state.redis_pool)


# orm 迁移生成表
# https://tortoise.github.io/migration.html
TORTOISE_ORM = {
    "connections": {"default": settings.MYSQL_DATABASE_URL},
    "apps": {
        "models": {
            "models": ["models", "aerich.models"],
            "default_connection": "default",
        },
    }
}