import aiomysql, aiosqlite, os
from fastapi import FastAPI
from contextlib import asynccontextmanager

# Either "mysql" or "sqlite"
mode = os.environ.get("DB_MODE", "sqlite")
# For SQLite usage
sqlite_file = os.environ.get("SQLITE_FILE","./database.db")

# For MySQL usage
mysql_pool = None

mysql_host = os.environ.get("MYSQL_HOST","127.0.0.1")
mysql_user = os.environ.get("MYSQL_USER","yt2mp3")
mysql_db_name = os.environ.get("MYSQL_DB","yt2mp3")
mysql_pwd = os.environ.get("MYSQL_PASSWORD","")


@asynccontextmanager
async def lifespan(app: FastAPI):
    if mode == "mysql":
        pool = await aiomysql.create_pool(host=mysql_host,user=mysql_user,password=mysql_pwd,db=mysql_db_name, minsize=1, maxsize=10)
        try:
            yield
        finally: 
            pool.close()
            await pool.wait_closed()
    else:
        try:
            yield
        finally: pass # no action required now