from random import randint

from fastapi import FastAPI
import os
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import FastAPI

from tortoise import Tortoise, generate_config
from tortoise.contrib.fastapi import RegisterTortoise

from models import RawData


@asynccontextmanager
async def lifespan_test(app: FastAPI) -> AsyncGenerator[None, None]:
    config = generate_config(
        "sqlite://:memory:",
        app_modules={"models": ["models"]},
        testing=True,
        connection_label="models",
    )
    async with RegisterTortoise(
            app=app,
            config=config,
            generate_schemas=True,
            add_exception_handlers=True,
            _create_db=True,
    ):
        # db connected
        yield
        # app teardown
    # db connections closed
    await Tortoise._drop_databases()


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    if getattr(app.state, "testing", None):
        async with lifespan_test(app) as _:
            yield
    else:
        config = generate_config(
            "sqlite:///db.sqlite3",
            app_modules={"models": ["models"]},
            connection_label="models",
        )
        async with RegisterTortoise(
                app=app,
                config=config,
                generate_schemas=True,
                add_exception_handlers=True,
                _create_db=True,
        ):
            # db connected
            yield
            # app teardown
        # db connections closed
        await Tortoise.close_connections()


app = FastAPI(lifespan=lifespan)


@app.get("/")
async def root():
    cnt = await RawData.all().count()
    # 生成一个范围内的随机整数，随机返回一篇文章
    r = randint(1, cnt)
    record = await RawData.get(id=r)
    return {"title": record.title, "content": record.content}