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":
        try:
            print("Going with mysql")
            global mysql_pool
            mysql_pool = await aiomysql.create_pool(host=mysql_host,user=mysql_user,password=mysql_pwd,db=mysql_db_name, minsize=1, maxsize=10)
            yield
        finally: 
            mysql_pool.close()
            await mysql_pool.wait_closed()
    elif mode == "sqlite":
        try:
            print("Going with sqlite")
            yield
        finally: pass # no action required now
    else:
        raise EnvironmentError("DB_MODE must either be mysql or sqlite")