from collections.abc import Callable, Generator, Iterable, Iterator
from contextlib import closing
from enum import Enum
from functools import wraps
from typing import Any, Protocol, overload, runtime_checkable
import sqlite3 as sql3

from coolqlite._common import _DONT_INIT_YOURSELF_OR_THINGS_WILL_BREAK, CoolqliteError
from coolqlite._util import _identity_fn


class _RowCount(Enum):
    "The number of rows expected or received."

    Zero = "zero"
    One = "one"
    Any = "any"
    OneOrNone = "one or none"
    MoreThanOne = "more than one"

    def __str__(self):
        return self.value


class Row(sql3.Row):
    """
    Extends a :class:`sqlite3.Row` with the following niceities:

    - Easily allows accessing through a namespace
    - Requesting a missing column by name raises a :exc:`KeyError` instead
      of an :exc:`IndexError`
    - The ``repr`` is more readable
    """

    # if TYPE_CHECKING:

    #     def __new__(cls, cursor: sql3.Cursor, data: tuple[Any, ...]) -> Self: ...
    #     def __init__(self, cursor: sql3.Cursor, data: tuple[Any, ...]): ...

    @overload
    def __getitem__(self, key: str | int) -> Any: ...
    @overload
    def __getitem__(self, key: slice[Any, Any, Any]) -> tuple[Any, ...]: ...
    def __getitem__(self, key: str | int | slice[Any, Any, Any]):
        try:
            return super().__getitem__(key)
        except IndexError:
            if isinstance(key, str):
                raise KeyError(key) from None
            else:
                raise

    # TODO: they can do this themselves!
    # def namespace(self) -> types.SimpleNamespace:
    #     "Access columns as if they were attributes."
    #     return types.SimpleNamespace(self)

    def __repr__(self):  # pragma: no cover
        keys = self.keys()
        if not keys:
            return "Row()"
        key_exprs = [f"{key}={self[key]!r}" for key in keys]
        return f"<row with columns {', '.join(key_exprs)}>"


# @overload
# def mapped_row[T](mapper: Callable[[Row], T]) -> RowFactory[T]: ...
# @overload
# def mapped_row[T, R](
#     mapper: Callable[[R], T], row_factory: RowFactory[R]
# ) -> RowFactory[T]: ...


# # TODO: none should be tuple, no?
# def mapped_row(mapper: Callable, row_factory: RowFactory | None = None) -> RowFactory:
#     """
#     Compose a conversion over a :type:`RowFactory`, making it a :type:`RowFactory` itself.

#     Defaults to :class:`Row` if the wrapped ``row_factory`` is not specified.

#     For something more substantial, see :mod:`coolqlite.cattrs_` .
#     """
#     fac = row_factory if row_factory is not None else Row
#     # TODO: wrapper to update __name__ or whatsits?
#     return lambda cursor, data: mapper(fac(cursor, data))


# TODO: attach rows?
class ResultCountError(CoolqliteError):
    "A different number of rows than expected were returned."

    def __init__(self, *, expected: _RowCount, actual: _RowCount):
        super().__init__(f"Expected {expected} rows but got {actual}")
        self.expected = expected
        self.actual = actual


@runtime_checkable
class CloseableIterator[T](Iterator[T], Protocol):
    def close(self): ...


class _CursorIterator[RowType](Iterator[RowType]):
    def __init__(self, sql_cursor: sql3.Cursor):
        self._inner = sql_cursor

    def __next__(self) -> RowType:
        next_value = self._inner.fetchone()

        if next_value is None:
            # TODO: is this done for us?
            self._inner.close()
            raise StopIteration

        return next_value

    def close(self):
        self._inner.close()


def first_col[T](x: tuple[T, ...]) -> T:
    """Get the first item in a tuple.

    This is only here to help with type checkers complaining
    about being passed a union type for
    :meth:`~coolqlite.Connection.colquery` :

    You use :meth:`~coolqlite.Connection.query` instead,
    wrap the union in a tuple,
    and then give this to :meth:`Results.map` .

    .. code-block:: python

        result = db.query(t"select 1", tuple[int | None]).map(first_col)
        typing.assert_type(result, QueryResults[int | None])
    """
    return x[0]


# TODO: frozen?
class Results[RowType](Iterator[RowType]):
    """An iterator of rows that offers easy ways to
    validate the number of returned rows,
    and is parameterized over the type of the row
    (from the row factory).
    """

    @overload
    def __init__(
        self,
        inner: Iterator[RowType] | Iterable[RowType],
    ): ...
    @overload
    def __init__[Inner](
        self,
        inner: Iterator[Inner] | Iterable[Inner],
        mapping_fn: Callable[[Inner], RowType],
    ): ...

    def __init__(
        self,
        inner: Iterator | Iterable,
        mapping_fn: Callable = _identity_fn,
    ):
        self._close: Callable[[], None]
        match inner:
            case Generator():
                self._inner = inner
                self._close = inner.close
            case Iterator():
                self._inner = inner
                if hasattr(inner, "close") and isinstance(
                    (close := getattr(inner, "close")), Callable
                ):
                    self._close = close
                else:
                    self._close = lambda: None
            case Iterable():
                inner = iter(inner)
                # oh my god this feels dirty
                self.__init__(inner, mapping_fn)
                # close is not assigned because this'll get to
                # the Iterator case and do it there.

        self._closed = False
        self._mapping_fn = mapping_fn

    def __enter__(self):
        return self

    def __exit__(self, *_exc_deets):
        self.close()
        return False

    # TODO: what if it's wrapping a generator,
    # it exhausts itself, and then we call close on it?
    # we have to handle that case, right?

    def __next__(self) -> RowType:
        return self._mapping_fn(next(self._inner))

    def next(self) -> RowType | None:
        "Fetch the next row, or None if no rows remain."
        return next(self, None)

    def one(self) -> RowType:
        """Return exactly one row,
        raising a :exc:`QueryResultCountError`
        if there's zero or more than one.

        If you want *at least one*,
        then use ``next(cursor.any())``.
        """
        with closing(self):
            row = self.next()
            if row is None:
                raise ResultCountError(expected=_RowCount.One, actual=_RowCount.Zero)
            _next = self.next()
            if _next is not None:
                raise ResultCountError(
                    expected=_RowCount.One, actual=_RowCount.MoreThanOne
                )
            return row

    def one_or_none(self) -> RowType | None:
        """Return exactly one or zero rows,
        raising a :exc:`QueryResultCountError`
        if there's two or more.

        Useful for checking for the presence or absence of something.
        """
        with closing(self):
            row = self.next()
            if row is None:
                return None
            _next = self.next()
            if _next is not None:
                raise ResultCountError(
                    expected=_RowCount.OneOrNone, actual=_RowCount.MoreThanOne
                )
            return row

    def at_least_one(self) -> Iterator[RowType]:
        """Get an iterator that is guaranteed to yield at least one row.

        Raises a :exc:`QueryResultCountError` if there are none.
        """
        with closing(self):
            row = self.next()
            if row is None:
                self.close()
                raise ResultCountError(expected=_RowCount.Any, actual=_RowCount.Zero)
            yield row
            yield from self  # closes for us but doubled, uh?

    # IF SUBCLASSING...
    # what's the _right_ way to do this?
    # I guess it inherently needs to be customized, since
    # type(self) could always have a different signature?
    def map[R](self, map: Callable[[RowType], R], /) -> Results[R]:
        """Return a new :class:`Results` that applies the given function
        to every row.
        """
        return Results(self, map)

    # TODO: we can remove this, right?
    # just do a `select count(*)`?

    # def zero(self) -> None:
    #     "If there are any rows at all, raise a :exc:`QueryResultCountError` ."
    #     with closing(self):
    #         if self.next() is not None:
    #             raise QueryResultCountError(
    #                 expected=_RowCount.Zero, actual=_RowCount.Any
    #             )

    def close(self):
        """If the underlying iterator is closeable
        (has a ``close`` method), and hasn't been closed yet,
        close it.

        See :meth:`sqlite3.Cursor.close` if this was based off
        of a sqlite cursor..
        """
        # TODO: do we need to track? or is closing again a no-op?
        if not self._closed:
            self._close()
            self._closed = True


class QueryResults[RowType = Row](Results[RowType]):
    """An extension of :class:`Results` that allows
    access to cursor-specific information,
    like row count and ``lastrowid``.

    Note that the same semantics as in the ``sqlite3`` module apply!
    """

    def __init__(
        self,
        private,
        cursor: sql3.Cursor,
        mapping_fn: Callable[..., RowType] = _identity_fn,
    ):
        assert private is _DONT_INIT_YOURSELF_OR_THINGS_WILL_BREAK
        super().__init__(_CursorIterator(cursor), mapping_fn)
        self._cursor = cursor

    @property
    def rowcount(self):
        "The underlying cursor's ``rowcount``."
        return self._cursor.rowcount

    @property
    def lastrowid(self):
        "The underlying cursor's ``lastrowid``."
        return self._cursor.lastrowid

    def map[R](self, map: Callable[[RowType], R], /) -> QueryResults[R]:
        """Return a new :class:`QueryResults` that applies the given function
        to every row.
        """
        qr = QueryResults(
            _DONT_INIT_YOURSELF_OR_THINGS_WILL_BREAK,
            self._cursor,
            lambda x: map(self._mapping_fn(x)),
        )
        # TODO: we'd need this on the ctor if we make it frozen eventually
        return qr


def as_results[T, **P](fn: Callable[P, Iterable[T]]) -> Callable[P, Results[T]]:
    """Wrap a generator function so it returns a :class:`Results` instead.

    Useful for when you want to do some in-python processing over query results,
    but still want the caller to have control over how many rows they care about.

    .. code-block:: python

        from coolqlite import as_results, Connection

        @as_results
        def advanced_query(db: Connection):
            # imagine this is something more complex
            # than could easily be done in SQL.
            for name in db.colquery(t"select name from Table", str):
                if name.startswith("Alex"):
                    yield name[::-1]
                else:
                    yield name

        only_one = advanced_query(db).one()
    """

    @wraps(fn)
    def wrapped(*args: P.args, **kwargs: P.kwargs):
        return Results(fn(*args, **kwargs))

    return wrapped
