from dataclasses import dataclass
from datetime import datetime
from typing import Callable, Awaitable, Any
from typing import List, TypeVar, Tuple, Generic

from pydantic import BaseModel as PdModel
from tortoise.exceptions import IntegrityError, FieldError
from tortoise.queryset import QuerySet, MODEL

from web_foundation.environment.resources.database.models import ConcurrentEntity
from web_foundation.environment.resources.database.models import GenericDbType
from web_foundation.environment.resources.database.utils import integrity_error_format
from web_foundation.environment.workers.web.ext.request_handler import ProtectIdentity
from web_foundation.errors.app.application import InconsistencyError

EntityId = TypeVar("EntityId", bound=int)


async def retry_dummy(dummy: Any) -> bool:
    return False


@dataclass
class AccessProtectIdentity(Generic[GenericDbType], ProtectIdentity):
    pass


class ModelLoader(Generic[GenericDbType]):

    @classmethod
    async def fetch_related(cls, model: GenericDbType, query: QuerySet[MODEL],
                            fetch_fields: List[str] | None = None) -> QuerySet[MODEL]:
        if fetch_fields:
            for_select = {field for fk_field in model._meta.fk_fields for field in fetch_fields if
                          fk_field == field}
            for_prefetch = set(fetch_fields) - for_select
            if for_select:
                query = query.select_related(*for_select)
            if for_prefetch:
                query = query.prefetch_related(*for_prefetch)
        return query

    @staticmethod
    async def get_entity(model: GenericDbType,
                         entity_id: EntityId) -> GenericDbType:  # TODO TYPED USER
        entity = await model.get_or_none(id=entity_id)
        if not entity:
            raise InconsistencyError(message=f"{model.__name__} not found")  # type: ignore
        return entity

    @staticmethod
    async def create(model: GenericDbType, protection: AccessProtectIdentity, dto: PdModel,
                     **kwargs) -> GenericDbType:
        entity_kwargs = {}
        m2m_fields = {}
        for field, value in dto.dict().items():
            if value is None:
                continue
            if field.endswith('_id') or field.endswith('_ids'):
                raw_field = field[:field.find("_id")]
                if raw_field in model._meta.m2m_fields:
                    m2m_model = model._meta.fields_map[raw_field].model
                    if not isinstance(value, (list, set)):
                        value = [value]
                    m2m_entities = await m2m_model.filter(id__in=value)
                    if len(m2m_entities) != len(value):
                        raise InconsistencyError(message=f"Some {m2m_model.__name__} not found")
                    m2m_fields[raw_field] = m2m_entities
                else:
                    entity_kwargs[field] = value
            else:
                entity_kwargs[field] = value
        try:
            entity = await model.create(**entity_kwargs)
        except IntegrityError as exception:
            raise integrity_error_format(exception)
        except ValueError as exception:
            raise InconsistencyError(exception)

        for m2m_field, m2m_entities in m2m_fields.items():
            await getattr(entity, m2m_field).add(*m2m_entities)
        return entity

    @staticmethod
    async def read(model: GenericDbType, entity_id: EntityId, protection: AccessProtectIdentity,
                   fetch_fields: List[str] | None = None,
                   **kwargs) -> GenericDbType | None:
        query = model.filter(id=entity_id)
        if kwargs:
            query = query.filter(**kwargs)
        query = await ModelLoader.fetch_related(model, query, fetch_fields)
        entity = await query.first()
        return entity

    @staticmethod
    async def read_all(model: GenericDbType, protection: AccessProtectIdentity, limit: int = None, offset: int = None,
                       fetch_fields: List[str] = None,
                       **kwargs) -> Tuple[List[GenericDbType], int]:
        query = model.filter()
        if kwargs:
            if "order_by" in kwargs:
                query = query.order_by(kwargs.pop("order_by"))
            query = query.filter(**kwargs)

        if limit:
            query = query.limit(limit)
        if offset:
            query = query.offset(offset)
        query = await ModelLoader.fetch_related(model, query, fetch_fields)
        try:
            entities = await query
            query._limit = None
            query._offset = None
            total = await query.count()
        except FieldError:
            raise InconsistencyError(message="Incorrect filter")
        return entities, total

    @staticmethod
    async def update(model: GenericDbType, entity_id: EntityId, protection: AccessProtectIdentity, dto: PdModel,
                     **kwargs) -> GenericDbType:
        entity = await ModelLoader.get_entity(model, entity_id)
        m2m_fields = {}
        entity_kwargs = {field: value for field, value in dto.dict().items() if value is not None}
        for field, value in entity_kwargs.items():
            if field.endswith('_id') or field.endswith('_ids'):
                raw_field = field[:field.find("_id")]
                if raw_field in model._meta.m2m_fields:
                    m2m_model = model._meta.fields_map[raw_field].model
                    if not isinstance(value, (list, set)):
                        value = [value]
                    m2m_entities = await m2m_model.filter(id__in=value)
                    if len(m2m_entities) != len(value):
                        raise InconsistencyError(message=f"Some {m2m_model.__name__} not found")
                    m2m_fields[raw_field] = m2m_entities
                    continue
            setattr(entity, field, value)
        try:
            await entity.save(update_fields=list(entity_kwargs.keys()))
        except IntegrityError as exception:
            raise integrity_error_format(exception)
        except ValueError as exception:
            raise InconsistencyError(exception)
        for m2m_field, m2m_entities in m2m_fields.items():
            await getattr(entity, m2m_field).add(*m2m_entities)
        return entity

    @staticmethod
    async def delete(model: GenericDbType, entity_id: EntityId, protection: AccessProtectIdentity,
                     **kwargs) -> GenericDbType:
        assert entity_id is not None
        entity = await ModelLoader.get_entity(model, entity_id)
        await entity.delete()
        return entity

    @classmethod
    async def concurrent_update(cls, query: QuerySet, should_retry: Callable[[Any], Awaitable[bool]] = retry_dummy,
                                **kwargs):
        """
            This function is designed to apply updates to database objects in concurrent manner
            and uses Read-Modify-Write Pattern (RMW).

            It's implemented in three basic steps:
            1. Read object -> acquire/own it;
            2. Modify local copy;
            3. Apply updates only if we still own it (we have it's latest version)

            Step 3 may fail in case when concurrency detected:
            some other database session modified target object
            while we were modifying our local copy, which means
            we do not 'own' the object and must refresh it
            in order to get the latest version.

            It is perceived by comparing local 'updated_at' field
            with it's actual value at step 3.
            So the typical produced query is going to look like:
            update table set value=x where updated_at=<our local updated_at value>

            'query' is required to contain some filtering expression
            'apply' must contain kwargs to be applied in case of success
            'should_retry' functor can be provided as guard who checks retry neccesarity

            NOTE: To keep database consistency regarding
            concurrently accessed objects - every code must MODIFY
            such objects only through this function.

            NOTE: requires connections to use 'read committed' policy
        """
        while True:
            entity: ConcurrentEntity = await query.first()
            if not entity:
                return None

            kwargs.update({'updated_at': datetime.now()})
            done = await query.filter(updated_at=entity.updated_at) \
                .update(**kwargs)

            if done:
                return await query.first()
            elif not await should_retry(query):
                return None
