from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Optional, Union

from ormar.queryset.utils import get_relationship_alias_model_and_str

if TYPE_CHECKING:  # pragma: no cover
    from ormar import Model


@dataclass
class Excludable:
    """
    Class that keeps sets of fields to exclude and include
    """

    include: set = field(default_factory=set)
    exclude: set = field(default_factory=set)

    def get_copy(self) -> "Excludable":
        """
        Return copy of self to avoid in place modifications
        :return: copy of self with copied sets
        :rtype: ormar.models.excludable.Excludable
        """
        _copy = self.__class__()
        _copy.include = {x for x in self.include}
        _copy.exclude = {x for x in self.exclude}
        return _copy

    def set_values(self, value: set, is_exclude: bool) -> None:
        """
        Appends the data to include/exclude sets.

        :param value: set of values to add
        :type value: set
        :param is_exclude: flag if values are to be excluded or included
        :type is_exclude: bool
        """
        prop = "exclude" if is_exclude else "include"
        current_value = getattr(self, prop)
        current_value.update(value)
        setattr(self, prop, current_value)

    def is_included(self, key: str) -> bool:
        """
        Check if field in included (in set or set is {...})
        :param key: key to check
        :type key: str
        :return: result of the check
        :rtype: bool
        """
        return (... in self.include or key in self.include) if self.include else True

    def is_excluded(self, key: str) -> bool:
        """
        Check if field in excluded (in set or set is {...})
        :param key: key to check
        :type key: str
        :return: result of the check
        :rtype: bool
        """
        return (... in self.exclude or key in self.exclude) if self.exclude else False


class ExcludableItems:
    """
    Keeps a dictionary of Excludables by alias + model_name keys
    to allow quick lookup by nested models without need to travers
    deeply nested dictionaries and passing include/exclude around
    """

    def __init__(self) -> None:
        self.items: dict[str, Excludable] = dict()

    @classmethod
    def from_excludable(cls, other: "ExcludableItems") -> "ExcludableItems":
        """
        Copy passed ExcludableItems to avoid inplace modifications.

        :param other: other excludable items to be copied
        :type other: ormar.models.excludable.ExcludableItems
        :return: copy of other
        :rtype: ormar.models.excludable.ExcludableItems
        """
        new_excludable = cls()
        for key, value in other.items.items():
            new_excludable.items[key] = value.get_copy()
        return new_excludable

    def include_entry_count(self) -> int:
        """
        Returns count of include items inside
        """
        count = 0
        for key in self.items.keys():
            count += len(self.items[key].include)
        return count

    def get(self, model_cls: type["Model"], alias: str = "") -> Excludable:
        """
        Return Excludable for given model and alias.

        :param model_cls: target model to check
        :type model_cls: ormar.models.metaclass.ModelMetaclass
        :param alias: table alias from relation manager
        :type alias: str
        :return: Excludable for given model and alias
        :rtype: ormar.models.excludable.Excludable
        """
        key = f"{alias + '_' if alias else ''}{model_cls.get_name(lower=True)}"
        excludable = self.items.get(key)
        if not excludable:
            excludable = Excludable()
            self.items[key] = excludable
        return excludable

    def build(
        self,
        items: Union[list[str], str, tuple[str], set[str], dict],
        model_cls: type["Model"],
        is_exclude: bool = False,
    ) -> None:
        """
        Receives the one of the types of items and parses them as to achieve
        a end situation with one excludable per alias/model in relation.

        Each excludable has two sets of values - one to include, one to exclude.

        :param items: values to be included or excluded
        :type items: Union[list[str], str, tuple[str], set[str], dict]
        :param model_cls: source model from which relations are constructed
        :type model_cls: ormar.models.metaclass.ModelMetaclass
        :param is_exclude: flag if items should be included or excluded
        :type is_exclude: bool
        """
        if isinstance(items, str):
            items = {items}

        if isinstance(items, dict):
            self._traverse_dict(
                values=items,
                source_model=model_cls,
                model_cls=model_cls,
                is_exclude=is_exclude,
            )

        else:
            items = set(items)
            nested_items = set(x for x in items if "__" in x)
            items.difference_update(nested_items)
            self._set_excludes(
                items=items,
                model_name=model_cls.get_name(lower=True),
                is_exclude=is_exclude,
            )
            if nested_items:
                self._traverse_list(
                    values=nested_items, model_cls=model_cls, is_exclude=is_exclude
                )

    def _set_excludes(
        self, items: set, model_name: str, is_exclude: bool, alias: str = ""
    ) -> None:
        """
        Sets set of values to be included or excluded for given key and model.

        :param items: items to include/exclude
        :type items: set
        :param model_name: name of model to construct key
        :type model_name: str
        :param is_exclude: flag if values should be included or excluded
        :type is_exclude: bool
        :param alias:
        :type alias: str
        """
        key = f"{alias + '_' if alias else ''}{model_name}"
        excludable = self.items.get(key)
        if not excludable:
            excludable = Excludable()
        excludable.set_values(value=items, is_exclude=is_exclude)
        self.items[key] = excludable

    def _traverse_dict(  # noqa: CFQ002
        self,
        values: dict,
        source_model: type["Model"],
        model_cls: type["Model"],
        is_exclude: bool,
        related_items: Optional[list] = None,
        alias: str = "",
    ) -> None:
        """
        Goes through dict of nested values and construct/update Excludables.

        :param values: items to include/exclude
        :type values: dict
        :param source_model: source model from which relations are constructed
        :type source_model: ormar.models.metaclass.ModelMetaclass
        :param model_cls: model from which current relation is constructed
        :type model_cls: ormar.models.metaclass.ModelMetaclass
        :param is_exclude: flag if values should be included or excluded
        :type is_exclude: bool
        :param related_items: list of names of related fields chain
        :type related_items: list
        :param alias: alias of relation
        :type alias: str
        """
        self_fields = set()
        related_items = related_items[:] if related_items else []
        for key, value in values.items():
            if value is ...:
                self_fields.add(key)
            elif isinstance(value, set):
                (
                    table_prefix,
                    target_model,
                    _,
                    _,
                ) = get_relationship_alias_model_and_str(
                    source_model=source_model, related_parts=related_items + [key]
                )
                self._set_excludes(
                    items=value,
                    model_name=target_model.get_name(),
                    is_exclude=is_exclude,
                    alias=table_prefix,
                )
            else:
                # dict
                related_items.append(key)
                (
                    table_prefix,
                    target_model,
                    _,
                    _,
                ) = get_relationship_alias_model_and_str(
                    source_model=source_model, related_parts=related_items
                )
                self._traverse_dict(
                    values=value,
                    source_model=source_model,
                    model_cls=target_model,
                    is_exclude=is_exclude,
                    related_items=related_items,
                    alias=table_prefix,
                )
        if self_fields:
            self._set_excludes(
                items=self_fields,
                model_name=model_cls.get_name(),
                is_exclude=is_exclude,
                alias=alias,
            )

    def _traverse_list(
        self, values: set[str], model_cls: type["Model"], is_exclude: bool
    ) -> None:
        """
        Goes through list of values and construct/update Excludables.

        :param values: items to include/exclude
        :type values: set
        :param model_cls: model from which current relation is constructed
        :type model_cls: ormar.models.metaclass.ModelMetaclass
        :param is_exclude: flag if values should be included or excluded
        :type is_exclude: bool
        """
        # here we have only nested related keys
        for key in values:
            key_split = key.split("__")
            related_items, field_name = key_split[:-1], key_split[-1]
            table_prefix, target_model, _, _ = get_relationship_alias_model_and_str(
                source_model=model_cls, related_parts=related_items
            )
            self._set_excludes(
                items={field_name},
                model_name=target_model.get_name(),
                is_exclude=is_exclude,
                alias=table_prefix,
            )
