import re
from typing import Any, Iterable, Iterator, cast
import aiomysql


class SelectResult[T]:
    rowcount: int

    async def fetchone(self) -> T | None:
        """
        Fetches a single row from the cursor. None indicates that
        no more rows are available.
        """

    async def fetchall(self) -> list[T]:
        """Fetchs all available rows from the cursor."""

    async def fetchmany(self, size: int | None = None) -> list[T]:
        """
        Fetch up to size rows from the cursor. Result set may be smaller
        than size. If size is not defined, cursor.arraysize is used.
        """


class OtherResult:
    rowcount: int


class InsertWithLastRowIdResult:
    rowcount: int
    lastrowid: int


class InsertWithOptLastRowIdResult:
    rowcount: int
    lastrowid: int | None


class UntypedResult:
    """
    Return type of execute, if the mysql-type-plugin is not used by mypy
    """

    rowcount: int
    lastrowid: int | None

    async def fetchone(self) -> Any | None:
        """
        Fetches a single row from the cursor. None indicates that
        no more rows are available.
        """

    async def fetchall(self) -> list[Any]:
        """Fetchs all available rows from the cursor."""

    async def fetchmany(self, size: int | None = None) -> list[Any]:
        """
        Fetch up to size rows from the cursor. Result set may be smaller
        than size. If size is not defined, cursor.arraysize is used.
        """

    def __iter__(self) -> Iterator[Any]: ...


async def execute[CursorType: aiomysql.Cursor | aiomysql.DictCursor](
    c: CursorType, sql: str, *args: Any
) -> UntypedResult:
    """
    Call c.execute, with the given sqlstatement and argumens, and return c.

    The mysql-type-plugin for mysql will special type this function such
    that the number and types of args matches the what is expected for the query
    by analyzing the mysql-schema.sql file in the project root.

    The return type will be changed to either: InsertResult, OtherResult or
    SelectResult depending on the query type.
    For select queries the SelectResult will be typed with either a Tuple with
    arguments with Tuple or TypedDict determined from the sql. Based on wether
    the curser is a dict cursor or not.

    Arguments in the query are expected to be presented as %s or _LIST_
    . For list arguments the corresponding entry in args is assumed to be a list
    and the _LIST_ in the query is replaced by %s,%s,...,%s with where
    the number of $s's is equal to the list length
    """
    if "_LIST_" not in sql:
        await c.execute(sql, args)
    else:
        rargs = list(reversed(args))
        flatargs = []

        def replace_arg(mo) -> str:
            nonlocal flatargs
            while True:
                a = rargs.pop()
                if a is None:
                    raise Exception("Number of _LIST_ arguments do not match")
                if isinstance(a, list):
                    flatargs += a
                    if not a:
                        return "null"
                    return ", ".join(["%s"] * len(a))
                flatargs.append(a)

        sql = re.sub("_LIST_", replace_arg, sql)
        while rargs:
            a = rargs.pop()
            if isinstance(a, list):
                raise Exception("Number of _LIST_ arguments do not match")
            flatargs.append(a)
        await c.execute(sql, flatargs)
    return cast(UntypedResult, c)


async def executemany[CursorType: aiomysql.Cursor | aiomysql.DictCursor](
    c: CursorType, sql: str, args: Iterable[Iterable[Any]]
) -> OtherResult:
    """
    Call c.executemany, with the given sqlstatement and argumens, and return c.

    The mysql-type-plugin for mysql will special type this function such
    that the number and types of args matches the what is expected for the query
    by analyzing the mysql-schema.sql file in the project root.
    """
    await c.executemany(sql, args)
    return cast(OtherResult, c)


async def fetchone[CursorType: aiomysql.Cursor | aiomysql.DictCursor](
    c: CursorType, sql: str, *args: Any
) -> Any | None:
    """
    Shortcut for execute(c, sql, *args).fetchone()
    """
    ret = await execute(c, sql, *args)
    return await ret.fetchone()


async def fetchall[CursorType: aiomysql.Cursor | aiomysql.DictCursor](
    c: CursorType, sql: str, *args: Any
) -> list[Any]:
    """
    Shortcut for execute(c, sql, *args).fetchall()
    """
    ret = await execute(c, sql, *args)
    return await ret.fetchall()
