from typing import TypeVar, List, Optional

from sqlalchemy import and_, select, update, delete
from sqlalchemy.engine import ScalarResult
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.inspection import inspect
from sqlalchemy.sql import Select


class ModelMixin:
    # TODO: maybe rename "db_commit" to "commit" ?
    # TODO: add bulk methods
    """
    Base model for Sqlalchemy's ORM

    All mutable actions are flushed by default (autoflush=True, autocommit=True).
    After executing all queries, use session.commit() or session.rollback().
    Certain classes (i.e. starlette_web.common.http.base_endpoint.BaseHTTPEndpoint)
    already handle transaction logic for you.
    """

    class Meta:
        order_by = ()

    @classmethod
    def prepare_query(
        cls,
        limit: int = None,
        offset: int = None,
        order_by=(),
        **filter_kwargs,
    ) -> Select:
        _order_by = []

        if order_by:
            _order_by = [o for o in order_by]
        else:
            _order_by = [o for o in cls.Meta.order_by]

        _order_by = [
            getattr(cls, field.lstrip("-")).desc() if field.startswith("-") else getattr(cls, field)
            for field in _order_by
        ]

        query = select(cls).filter(cls._filter_criteria(filter_kwargs)).order_by(*_order_by)

        if limit is not None:
            query = query.limit(limit)

        if offset is not None:
            query = query.offset(offset)

        return query

    @classmethod
    async def async_filter(
        cls, db_session: AsyncSession, limit: int = None, offset: int = None, **filter_kwargs
    ) -> ScalarResult:
        query = cls.prepare_query(limit=limit, offset=offset, **filter_kwargs)
        result = await db_session.execute(query)
        return result.scalars()

    @classmethod
    async def async_get(cls, db_session: AsyncSession, **filter_kwargs) -> "DBModel":
        query = cls.prepare_query(**filter_kwargs)
        result = await db_session.execute(query)
        return result.scalars().first()

    @classmethod
    async def async_update(
        cls,
        db_session: AsyncSession,
        filter_kwargs: dict,
        update_data: dict,
        db_commit=False,
    ):
        query = (
            update(cls)
            .where(cls._filter_criteria(filter_kwargs))
            .values(**update_data)
            .execution_options(synchronize_session="fetch")
        )
        await db_session.execute(query)
        if db_commit:
            await db_session.commit()

    @classmethod
    async def async_delete(cls, db_session: AsyncSession, filter_kwargs: dict, db_commit=False):
        query = (
            delete(cls)
            .where(cls._filter_criteria(filter_kwargs))
            .execution_options(synchronize_session="fetch")
        )
        await db_session.execute(query)
        if db_commit:
            await db_session.commit()

    @classmethod
    async def async_create(cls, db_session: AsyncSession, db_commit=False, **data) -> "DBModel":
        instance = cls(**data)  # noqa
        db_session.add_all([instance])

        # Explicit flush() to populate newly created instance with primary key (if autogenerated).
        await db_session.flush()

        if db_commit:
            await db_session.commit()
        return instance

    @classmethod
    async def async_create_or_update(
        cls,
        db_session: AsyncSession,
        filter_kwargs: dict,
        update_data: dict,
        db_commit=False,
        force_update=False,
        update_to_null=False,
    ) -> "DBModel":
        """
        Allows to create entry, if not exists, or update an existing query.

        :param db_session: active instance of AsyncSession
        :param filter_kwargs: function will search entry by these conditions
        :param update_data: function will update these fields.
            "update_data" has priority over "filter_kwargs" in case of a new entry.
        :param db_commit: bool, whether to execute db_session.commit()
        :param force_update: bool, whether to do update is instance exists and has not been changed
        :param update_to_null: bool, whether to update currently non-null values to new null values
        :returns: instance of DBModel
        """
        # TODO: contrib.postgres.PostgresModelMixin with this method reloaded
        #  (using on_conflict do update)
        if instance := await cls.async_get(db_session, **filter_kwargs):
            if cls._object_needs_update(
                instance.to_dict(),
                update_data,
                force_update=force_update,
                update_to_null=update_to_null,
            ):
                await cls.async_update(
                    db_session,
                    filter_kwargs=filter_kwargs,
                    update_data=update_data,
                    db_commit=db_commit,
                )
        else:
            instance = await cls.async_create(
                db_session,
                **filter_kwargs,
                **update_data,
                db_commit=db_commit,
            )
        return instance

    @classmethod
    async def async_get_or_create(
        cls,
        db_session: AsyncSession,
        filter_kwargs: dict,
        extra_data: Optional[dict] = None,
        db_commit=False,
    ) -> "DBModel":
        """
        Allows to get existing instance or create, if not exists.
        Allows passing additional data upon creation.

        :param db_session: active instance of AsyncSession
        :param filter_kwargs: function will search entry by these conditions
        :param extra_data: extra data to create an entry, if not exists
        :param db_commit: bool, whether to execute db_session.commit()
        :returns: instance of DBModel
        """
        instance = await cls.async_get(db_session, **filter_kwargs)
        if not instance:
            extra_data = extra_data or {}
            instance = await cls.async_create(
                db_session,
                **filter_kwargs,
                **extra_data,
                db_commit=db_commit,
            )
        return instance

    async def update(self, db_session: AsyncSession, db_commit=False, **update_data):
        primary_keys_names = [key.name for key in inspect(self.__class__).primary_key]
        filter_kwargs = {key: getattr(self, key) for key in primary_keys_names}
        await self.async_update(db_session, filter_kwargs, update_data)

        if db_commit:
            await db_session.commit()

    async def delete(self, db_session: AsyncSession, db_commit=False):
        # The delete() method only marks object as deleted,
        # so flush() has to be used explicitly, as per documentation.
        await db_session.delete(self)
        await db_session.flush()

        if db_commit:
            await db_session.commit()

    def to_dict(self, excluded_fields: List[str] = None) -> dict:
        excluded_fields = excluded_fields or []
        res = {}
        for field in self.__dict__:
            if field not in excluded_fields and not field.startswith("_"):
                res[field] = getattr(self, field)

        return res

    @classmethod
    def _filter_criteria(cls, filter_kwargs):
        filters = []
        for filter_name, filter_value in filter_kwargs.items():
            field, _, criteria = filter_name.partition("__")
            if criteria in ("eq", ""):
                filters.append(getattr(cls, field) == filter_value)
            elif criteria == "gt":
                filters.append(getattr(cls, field) > filter_value)
            elif criteria == "lt":
                filters.append(getattr(cls, field) < filter_value)
            elif criteria == "is":
                filters.append(getattr(cls, field).is_(filter_value))
            elif criteria == "in":
                filters.append(getattr(cls, field).in_(filter_value))
            elif criteria == "notin":
                filters.append(getattr(cls, field).notin_(filter_value))
            elif criteria == "inarr":
                filters.append(getattr(cls, field).contains([filter_value]))
            elif criteria == "icontains":
                filters.append(getattr(cls, field).ilike(f"%{filter_value}%"))
            elif criteria == "ne":
                filters.append(getattr(cls, field) != filter_value)
            else:
                raise NotImplementedError(f"Unexpected criteria: {criteria}")

        return and_(True, *filters)

    @staticmethod
    def _object_needs_update(
        dict_original,
        dict_new,
        force_update=False,
        update_to_null=True,
    ) -> bool:
        if force_update:
            return True

        for key in dict_new:
            if dict_new[key] is None and not update_to_null:
                return False
            if dict_original.get(key) != dict_new[key]:
                return True

        return False


DBModel = TypeVar("DBModel", bound=ModelMixin)
