import time
import typing
import aiosqlite
import aiosqlitepool

# Local imports
from . import udfs


class SQLiteAsyncConnection:
    """Class to create SQLite pooled connection and execute SQL queries asynchronous."""

    # ==-----------------------------------------------------------------------------== #
    # Public methods                                                                    #
    # ==-----------------------------------------------------------------------------== #
    def __init__(self, path: str, pool_size: int = 4, idle_timeout: int | None = 30, queue_timeout: int | None = 60 * 60 * 24, user_defined_functions: list[typing.Callable[..., typing.Any]] = list()):
        """Creates instance of SQLite pooled connection to execute SQL queries asynchronous."""

        # Save initializer arguments as class attributes
        self.path = path
        self.pool_size = pool_size
        self.idle_timeout = idle_timeout
        self.queue_timeout = queue_timeout
        self.user_defined_functions = user_defined_functions

        # Initialize pool
        self.connection_pool = aiosqlitepool.SQLiteConnectionPool(self.__create_connection, self.pool_size, idle_timeout=self.idle_timeout, acquisition_timeout=self.queue_timeout)

    async def execute_single(self, query: str, *args: typing.Any) -> None:
        """Execute single SQL query."""

        # Flag if query successed
        query_successed = True

        # Execute SQL query
        async with self.connection_pool.connection() as connection:

            try:

                # Start time of query execution
                start_time = time.perf_counter()

                # Execute query
                async with connection.execute(query, args) as cursor:
                    columns = [column[0] for column in cursor.description] if cursor.description else None
                    data = await cursor.fetchall() if columns else None

                return {
                    "execution_time_ms": (time.perf_counter() - start_time) * 1000
                } | ({
                    "columns": columns,
                    "data": [{key: value for key, value in zip(columns, item)} for item in data]
                } if columns else {})

            except Exception:

                # Change flag of query success
                query_successed = False

                # Reraise exception
                raise

            finally:

                # If transaction is open
                if connection.in_transaction:

                    # Commit or rollback changes
                    await connection.execute("COMMIT;") if query_successed else await connection.rollback()

    # ==-----------------------------------------------------------------------------== #
    # Private methods                                                                   #
    # ==-----------------------------------------------------------------------------== #
    async def __create_connection(self) -> aiosqlite.Connection:
        """Private method to produce new aiosqlite and give it to connection pool."""

        # Create aiosqlite connection instance
        connection = await aiosqlite.connect(self.path)

        # Use WAL journal mode for more concurrent performance
        await connection.execute("PRAGMA journal_mode = WAL;")

        # Use auto-checkpoint for WAL
        await connection.execute("PRAGMA wal_autocheckpoint = 1000;")

        # Allow to use pragma keys
        await connection.execute("PRAGMA foreign_keys = ON;")

        # Register functions provided by default
        await self.__registrate_functions(connection, [item for item in dict(udfs.__dict__).values() if hasattr(item, "__code__")])

        # Register functions provided by user
        await self.__registrate_functions(connection, self.user_defined_functions)

        # Return connection instance
        return connection

    async def __registrate_functions(self, connection: aiosqlite.Connection, functions: list[typing.Callable[..., typing.Any]]) -> None:
        """Private method to register user-defined functions and make it callable from SQL queries."""

        # Iterate all of the functions
        for function in functions:

            # Delete return value from function annotations, if it exists
            if "return" in (function_annotations := function.__annotations__.copy()):
                del function_annotations["return"]

            # If not all function arguments is annotated
            if function.__code__.co_argcount != len(function_annotations):
                raise aiosqlite.Error(f"All `{function.__name__}` function arguments have to be annotated")

            # Function registration
            await connection.create_function(function.__name__.upper(), -1, function)
