import peewee
from peewee_async import AsyncDatabase

try:
    import aiosqlite
except ImportError:
    aiosqlite = None

__all__ = ["SqliteDatabase"]


class AsyncSqliteConnection:
    def __init__(self, *, database=None, loop=None, timeout=None, **kwargs):
        self._connections = []
        self._used_connections = []
        self.loop = loop
        self.database = database
        self.timeout = timeout
        self.connect_kwargs = kwargs

    async def acquire(self):
        if self._connections:
            conn = self._connections.pop()
        else:
            conn = aiosqlite.connect(database=self.database, **self.connect_kwargs)

        await conn.__aenter__()
        self._used_connections.append(conn)
        return conn

    async def release(self, conn):
        if conn in self._used_connections:
            self._used_connections.remove(conn)
        await conn.__aexit__(None, None, None)
        self._connections.append(conn)

    async def connect(self):
        pass

    async def close(self):
        for conn in self._used_connections:
            await conn.__aexit__(None, None, None)
        self._used_connections = self._connections = None

    async def cursor(self, conn=None, *args, **kwargs):
        in_transaction = conn is not None
        if not conn:
            conn = await self.acquire()
        cursor = await conn.cursor()
        # cursor.release is a coroutine
        cursor.release = self.release_cursor(  # pylint: disable = assignment-from-no-return
            cursor, in_transaction
        )
        return cursor

    async def release_cursor(self, cursor, in_transaction=False):
        conn = cursor._conn
        await cursor.__aexit__(None, None, None)
        if not in_transaction:
            await self.release(conn)


class AsyncSqliteMixin(AsyncDatabase):
    if aiosqlite:
        import sqlite3

        Error = sqlite3.Error

    def init_async(self, conn_class=AsyncSqliteConnection):
        self._async_conn_cls = conn_class

    @property
    def connect_kwargs_async(self):
        return {**self.connect_kwargs}


class SqliteDatabase(AsyncSqliteMixin, peewee.SqliteDatabase):
    def init(self, database, **kwargs):
        if not aiosqlite:
            raise Exception("Error, aiosqlite is not installed!")
        super().init(database, **kwargs)
        self.init_async()

    async def last_insert_id_async(self, cursor, model):
        """Get ID of last inserted row.
        """
        if model._meta.auto_increment:
            return cursor.lastrowid