from __future__ import annotations

import types
from collections import defaultdict
from collections.abc import AsyncIterator, Callable, Collection, Generator, Iterable
from copy import copy
from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, TypeVar, cast, overload

from pypika_tortoise import JoinType, Order, Table
from pypika_tortoise.analytics import Count
from pypika_tortoise.functions import Cast
from pypika_tortoise.queries import QueryBuilder
from pypika_tortoise.terms import Case, Field, Star, Term, ValueWrapper

from tortoise.backends.base.client import BaseDBAsyncClient, Capabilities
from tortoise.exceptions import (
    DoesNotExist,
    FieldError,
    IntegrityError,
    MultipleObjectsReturned,
    ParamsError,
)
from tortoise.expressions import Expression, Q, RawSQL, ResolveContext, ResolveResult
from tortoise.fields.relational import (
    ForeignKeyFieldInstance,
    OneToOneFieldInstance,
    RelationalField,
)
from tortoise.filters import FilterInfoDict
from tortoise.query_utils import (
    Prefetch,
    QueryModifier,
    TableCriterionTuple,
    expand_lookup_expression,
    get_joins_for_related_field,
)
from tortoise.router import router
from tortoise.utils import chunk

# Empty placeholder - Should never be edited.

QUERY: QueryBuilder = QueryBuilder()

if TYPE_CHECKING:  # pragma: nocoverage
    from tortoise.models import Model

MODEL = TypeVar("MODEL", bound="Model")
PRIMARY_KEY = TypeVar("PRIMARY_KEY")
T_co = TypeVar("T_co", covariant=True)
SINGLE = TypeVar("SINGLE", bound=bool)


class QuerySetSingle(Protocol[T_co]):
    """
    Awaiting on this will resolve a single instance of the Model object, and not a sequence.
    """

    # pylint: disable=W0104
    def __await__(self) -> Generator[Any, None, T_co]: ...  # pragma: nocoverage

    def prefetch_related(
        self, *args: str | Prefetch
    ) -> QuerySetSingle[T_co]: ...  # pragma: nocoverage

    def select_related(self, *args: str) -> QuerySetSingle[T_co]: ...  # pragma: nocoverage

    def annotate(
        self, **kwargs: Expression | Term
    ) -> QuerySetSingle[T_co]: ...  # pragma: nocoverage

    def only(self, *fields_for_select: str) -> QuerySetSingle[T_co]: ...  # pragma: nocoverage

    def values_list(
        self, *fields_: str, flat: bool = False
    ) -> ValuesListQuery[Literal[True]]: ...  # pragma: nocoverage

    def values(
        self, *args: str, **kwargs: str
    ) -> ValuesQuery[Literal[True]]: ...  # pragma: nocoverage


class AwaitableQuery(Generic[MODEL]):
    __slots__ = (
        "query",
        "model",
        "_joined_tables",
        "_db",
        "_capabilities",
        "_annotations",
        "_custom_filters",
        "_q_objects",
    )

    def __init__(self, model: type[MODEL]) -> None:
        self._joined_tables: list[Table] = []
        self.model: type[MODEL] = model
        self.query: QueryBuilder = QUERY
        self._db: BaseDBAsyncClient = None  # type: ignore
        self._capabilities: Capabilities | None = None
        self._annotations: dict[str, Expression | Term] = {}
        self._custom_filters: dict[str, FilterInfoDict] = {}
        self._q_objects: list[Q] = []

    @property
    def capabilities(self) -> Capabilities:
        if self._capabilities is None:
            self._capabilities = self.model._meta.db.capabilities
        return self._capabilities

    @capabilities.setter
    def capabilities(self, value: Capabilities) -> None:
        self._capabilities = value

    def _choose_db(self, for_write: bool = False) -> BaseDBAsyncClient:
        """
        Return the connection that will be used if this query is executed now.

        :return: BaseDBAsyncClient:
        """
        if self._db:
            return self._db
        if for_write:
            db = router.db_for_write(self.model)
        else:
            db = router.db_for_read(self.model)
        return db or self.model._meta.db

    def _choose_db_if_not_chosen(self, for_write: bool = False) -> None:
        if self._db is None:
            self._db = self._choose_db(for_write)  # type: ignore

    def resolve_filters(self, fields_for_select: Collection[str] | None = None) -> None:
        """Builds the common filters for a QuerySet."""
        has_aggregate = self._resolve_annotate(fields_for_select)

        modifier = QueryModifier()
        for node in self._q_objects:
            modifier &= node.resolve(
                ResolveContext(
                    model=self.model,
                    table=self.model._meta.basetable,
                    annotations=self._annotations,
                    custom_filters=self._custom_filters,
                )
            )

        for join in modifier.joins:
            if join[0] not in self._joined_tables:
                self.query = self.query.join(join[0], how=JoinType.left_outer).on(join[1])
                self._joined_tables.append(join[0])

        self.query._havings = modifier.having_criterion
        self.query._wheres = modifier.where_criterion

        if has_aggregate and (self._joined_tables or self.query._havings or self.query._orderbys):
            self.query = self.query.groupby(
                *[self.model._meta.basetable[field] for field in self.model._meta.db_fields]
            )

    def _join_table_by_field(
        self, table: Table, related_field_name: str, related_field: RelationalField
    ) -> Table:
        joins = get_joins_for_related_field(table, related_field, related_field_name)
        for join in joins:
            self._join_table(join)
        return joins[-1][0]

    def _join_table(self, table_criterio_tuple: TableCriterionTuple) -> None:
        if table_criterio_tuple[0] not in self._joined_tables:
            self.query = self.query.join(table_criterio_tuple[0], how=JoinType.left_outer).on(
                table_criterio_tuple[1]
            )
            self._joined_tables.append(table_criterio_tuple[0])

    @staticmethod
    def _resolve_ordering_string(ordering: str, reverse: bool = False) -> tuple[str, Order]:
        order_type = Order.asc
        if ordering[0] == "-":
            field_name = ordering[1:]
            order_type = Order.desc
        else:
            field_name = ordering

        if reverse:
            order_type = Order.desc if order_type == Order.asc else Order.asc

        return field_name, order_type

    def resolve_ordering(
        self,
        model: type[Model],
        table: Table,
        orderings: Iterable[tuple[str, str | Order]],
        annotations: dict[str, Term | Expression],
        fields_for_select: Collection[str] | None = None,
    ) -> None:
        """
        Applies standard ordering to QuerySet.

        :param model: The Model this queryset is based on.
        :param table: ``pypika_tortoise.Table`` to keep track of the virtual SQL table
            (to allow self referential joins)
        :param orderings: What columns/order to order by
        :param annotations:  Annotations that may be ordered on
        :param fields_for_select: Contains fields that are selected in the SELECT clause if
            .only(), .values() or .values_list() are used.

        :raises FieldError: If a field provided does not exist in model.
        """
        # Do not apply default ordering for annotated queries to not mess them up
        if not orderings and self.model._meta.ordering and not annotations:
            orderings = self.model._meta.ordering

        for ordering in orderings:
            field_name = ordering[0]
            if field_name in model._meta.fetch_fields:
                raise FieldError(
                    "Filtering by relation is not possible. Filter by nested field of related model"
                )

            related_field_name, __, forwarded = field_name.partition("__")
            if related_field_name in model._meta.fetch_fields:
                related_field = cast(RelationalField, model._meta.fields_map[related_field_name])
                related_table = self._join_table_by_field(table, related_field_name, related_field)
                self.resolve_ordering(
                    related_field.related_model,
                    related_table,
                    [(forwarded, ordering[1])],
                    {},
                )
            elif field_name in annotations:
                term: Term
                if not fields_for_select or field_name in fields_for_select:
                    # The annotation is SELECTed, we can just reference it in the following cases:
                    # - Empty fields_for_select means that all columns and annotations are selected,
                    #   hence we can reference the annotation.
                    # - The annotation is in fields_for_select, hence we can reference it.
                    term = Field(field_name)
                else:
                    # The annotation is not in SELECT, resolve it
                    annotation = annotations[field_name]
                    if isinstance(annotation, Term):
                        term = annotation
                    else:
                        term = annotation.resolve(
                            ResolveContext(
                                model=self.model,
                                table=table,
                                annotations=annotations,
                                custom_filters={},
                            )
                        ).term
                self.query = self.query.orderby(term, order=ordering[1])
            else:
                field_object = model._meta.fields_map.get(field_name)

                if not field_object:
                    raise FieldError(f"Unknown field {field_name} for model {model.__name__}")
                field_name = field_object.source_field or field_name
                field = table[field_name]

                func = field_object.get_for_dialect(
                    model._meta.db.capabilities.dialect, "function_cast"
                )
                if func:
                    field = func(field_object, field)

                self.query = self.query.orderby(field, order=ordering[1])

    def _resolve_annotate(self, fields_for_select: Collection[str] | None = None) -> bool:
        if not self._annotations:
            return False

        annotation_info: dict[str, ResolveResult] = {}
        for key, annotation in self._annotations.items():
            if isinstance(annotation, Term):
                annotation_info[key] = ResolveResult(term=annotation)
            else:
                annotation_info[key] = annotation.resolve(
                    ResolveContext(
                        model=self.model,
                        table=self.model._meta.basetable,
                        annotations=self._annotations,
                        custom_filters=self._custom_filters,
                    )
                )

        for key, info in annotation_info.items():
            for join in info.joins:
                self._join_table(join)
            if key in self._annotations and (fields_for_select is None or key in fields_for_select):
                self.query._select_other(info.term.as_(key))  # type:ignore[arg-type]

        return any(info.term.is_aggregate for info in annotation_info.values())

    def sql(self, params_inline=False) -> str:
        """
        Returns the SQL query that will be executed. By default, it will return the query with
        placeholders, but if you set `params_inline=True`, it will inline the parameters.

        :param params_inline: Whether to inline the parameters
        """
        self._choose_db_if_not_chosen()

        self._make_query()
        if params_inline:
            sql = self.query.get_sql()
        else:
            sql, _ = self.query.get_parameterized_sql()
        return sql

    def _make_query(self) -> None:
        raise NotImplementedError()  # pragma: nocoverage

    async def _execute(self) -> Any:
        raise NotImplementedError()  # pragma: nocoverage


class QuerySet(AwaitableQuery[MODEL]):
    __slots__ = (
        "fields",
        "_prefetch_map",
        "_prefetch_queries",
        "_single",
        "_raise_does_not_exist",
        "_db",
        "_limit",
        "_offset",
        "_fields_for_select",
        "_filter_kwargs",
        "_orderings",
        "_distinct",
        "_having",
        "_group_bys",
        "_select_for_update",
        "_select_for_update_nowait",
        "_select_for_update_skip_locked",
        "_select_for_update_of",
        "_select_for_update_no_key",
        "_select_related",
        "_select_related_idx",
        "_use_indexes",
        "_force_indexes",
    )

    def __init__(self, model: type[MODEL]) -> None:
        super().__init__(model)
        self.fields: set[str] = model._meta.db_fields
        self._prefetch_map: dict[str, set[str | Prefetch]] = {}
        self._prefetch_queries: dict[str, list[tuple[str | None, QuerySet]]] = {}
        self._single: bool = False
        self._raise_does_not_exist: bool = False
        self._limit: int | None = None
        self._offset: int | None = None
        self._filter_kwargs: dict[str, Any] = {}
        self._orderings: list[tuple[str, Any]] = []
        self._distinct: bool = False
        self._having: dict[str, Any] = {}
        self._fields_for_select: tuple[str, ...] = ()
        self._group_bys: tuple[str, ...] = ()
        self._select_for_update: bool = False
        self._select_for_update_nowait: bool = False
        self._select_for_update_skip_locked: bool = False
        self._select_for_update_of: set[str] = set()
        self._select_for_update_no_key: bool = False
        self._select_related: set[str] = set()
        self._select_related_idx: list[
            tuple[type[Model], int, Table | str, type[Model], Iterable[str | None]]
        ] = []  # format with: model,idx,model_name,parent_model
        self._force_indexes: set[str] = set()
        self._use_indexes: set[str] = set()

    def _clone(self) -> QuerySet[MODEL]:
        queryset = self.__class__.__new__(self.__class__)
        queryset.fields = self.fields
        queryset.model = self.model
        queryset.query = self.query
        queryset.capabilities = self.capabilities
        queryset._prefetch_map = copy(self._prefetch_map)
        queryset._prefetch_queries = copy(self._prefetch_queries)
        queryset._single = self._single
        queryset._raise_does_not_exist = self._raise_does_not_exist
        queryset._db = self._db
        queryset._limit = self._limit
        queryset._offset = self._offset
        queryset._fields_for_select = self._fields_for_select
        queryset._filter_kwargs = copy(self._filter_kwargs)
        queryset._orderings = copy(self._orderings)
        queryset._joined_tables = copy(self._joined_tables)
        queryset._q_objects = copy(self._q_objects)
        queryset._distinct = self._distinct
        queryset._annotations = copy(self._annotations)
        queryset._having = copy(self._having)
        queryset._custom_filters = copy(self._custom_filters)
        queryset._group_bys = copy(self._group_bys)
        queryset._select_for_update = self._select_for_update
        queryset._select_for_update_nowait = self._select_for_update_nowait
        queryset._select_for_update_skip_locked = self._select_for_update_skip_locked
        queryset._select_for_update_of = self._select_for_update_of
        queryset._select_for_update_no_key = self._select_for_update_no_key
        queryset._select_related = self._select_related
        queryset._select_related_idx = self._select_related_idx
        queryset._force_indexes = self._force_indexes
        queryset._use_indexes = self._use_indexes
        return queryset

    def _filter_or_exclude(self, *args: Q, negate: bool, **kwargs: Any) -> QuerySet[MODEL]:
        queryset = self._clone()
        for arg in args:
            if not isinstance(arg, Q):
                raise TypeError("expected Q objects as args")
            if negate:
                queryset._q_objects.append(~arg)
            else:
                queryset._q_objects.append(arg)

        for key, value in kwargs.items():
            if negate:
                queryset._q_objects.append(~Q(**{key: value}))
            else:
                queryset._q_objects.append(Q(**{key: value}))

        return queryset

    def filter(self, *args: Q, **kwargs: Any) -> QuerySet[MODEL]:
        """
        Filters QuerySet by given kwargs. You can filter by related objects like this:

        .. code-block:: python3

            Team.filter(events__tournament__name='Test')

        You can also pass Q objects to filters as args.
        """
        return self._filter_or_exclude(negate=False, *args, **kwargs)

    def exclude(self, *args: Q, **kwargs: Any) -> QuerySet[MODEL]:
        """
        Same as .filter(), but with appends all args with NOT
        """
        return self._filter_or_exclude(negate=True, *args, **kwargs)

    def _parse_orderings(
        self, orderings: tuple[str, ...], reverse=False
    ) -> list[tuple[str, Order]]:
        """
        Convert ordering from strings to standard items for queryset.

        :param orderings: What columns/order to order by
        :param reverse:  Whether reverse order
        :return: standard ordering for QuerySet.
        """
        new_ordering = []
        for ordering in orderings:
            field_name, order_type = self._resolve_ordering_string(ordering, reverse=reverse)

            if not (
                field_name.split("__")[0] in self.model._meta.fields
                or field_name in self._annotations
            ):
                raise FieldError(f"Unknown field {field_name} for model {self.model.__name__}")
            new_ordering.append((field_name, order_type))
        return new_ordering

    def order_by(self, *orderings: str) -> QuerySet[MODEL]:
        """
        Accept args to filter by in format like this:

        .. code-block:: python3

            .order_by('name', '-tournament__name')

        Supports ordering by related models too.
        A '-' before the name will result in descending sort order, default is ascending.

        :raises FieldError: If unknown field has been provided.
        """
        queryset = self._clone()
        queryset._orderings = self._parse_orderings(orderings)
        return queryset

    def _as_single(self) -> QuerySetSingle[MODEL | None]:
        self._single = True
        self._limit = 1
        return cast(QuerySetSingle[MODEL | None], self)

    def latest(self, *orderings: str) -> QuerySetSingle[MODEL | None]:
        """
        Returns the most recent object by ordering descending on the providers fields.

        :params orderings: Fields to order by.

        :raises FieldError: If unknown or no fields has been provided.
        """
        if not orderings:
            raise FieldError("No fields passed")
        queryset = self._clone()
        queryset._orderings = self._parse_orderings(orderings, reverse=True)
        return queryset._as_single()

    def earliest(self, *orderings: str) -> QuerySetSingle[MODEL | None]:
        """
        Returns the earliest object by ordering ascending on the specified field.

        :params orderings: Fields to order by.

        :raises FieldError: If unknown or no fields has been provided.
        """
        if not orderings:
            raise FieldError("No fields passed")
        queryset = self._clone()
        queryset._orderings = self._parse_orderings(orderings)
        return queryset._as_single()

    def limit(self, limit: int) -> QuerySet[MODEL]:
        """
        Limits QuerySet to given length.

        :raises ParamsError: Limit should be non-negative number.
        """
        if limit < 0:
            raise ParamsError("Limit should be non-negative number")

        queryset = self._clone()
        queryset._limit = limit
        return queryset

    def offset(self, offset: int) -> QuerySet[MODEL]:
        """
        Query offset for QuerySet.

        :raises ParamsError: Offset should be non-negative number.
        """
        if offset < 0:
            raise ParamsError("Offset should be non-negative number")

        queryset = self._clone()
        queryset._offset = offset
        if self.capabilities.requires_limit and queryset._limit is None:
            queryset._limit = 1000000
        return queryset

    def __getitem__(self, key: slice) -> QuerySet[MODEL]:
        """
        Query offset and limit for Queryset.

        :raises ParamsError: QuerySet indices must be slices.

        :raises ParamsError: Slice steps should be 1 or None.

        :raises ParamsError: Slice start should be non-negative number or None.

        :raises ParamsError: Slice stop should be non-negative number greater that slice start,
        or None.
        """
        if not isinstance(key, slice):
            raise ParamsError("QuerySet indices must be slices.")

        if not (key.step is None or (isinstance(key.step, int) and key.step == 1)):
            raise ParamsError("Slice steps should be 1 or None.")

        start = key.start if key.start is not None else 0

        if not isinstance(start, int) or start < 0:
            raise ParamsError("Slice start should be non-negative number or None.")
        if key.stop is not None and (not isinstance(key.stop, int) or key.stop <= start):
            raise ParamsError(
                "Slice stop should be non-negative number greater that slice start, or None.",
            )

        queryset = self.offset(start)
        if key.stop:
            queryset = queryset.limit(key.stop - start)
        return queryset

    def distinct(self) -> QuerySet[MODEL]:
        """
        Make QuerySet distinct.

        Only makes sense in combination with a ``.values()`` or ``.values_list()`` as it
        precedes all the fetched fields with a distinct.
        """
        queryset = self._clone()
        queryset._distinct = True
        return queryset

    def select_for_update(
        self,
        nowait: bool = False,
        skip_locked: bool = False,
        of: tuple[str, ...] = (),
        no_key: bool = False,
    ) -> QuerySet[MODEL]:
        """
        Make QuerySet select for update.

        Returns a queryset that will lock rows until the end of the transaction,
        generating a SELECT ... FOR UPDATE SQL statement on supported databases.

        :param nowait:
            If `True`, raise an error if the lock cannot be obtained immediately.
        :param skip_locked:
            If `True`, skip rows that are already locked by other transactions instead of waiting.
        :param of:
            Specify the tables to lock when dealing with multiple related tables, e.g. when using `select_related`.
            Provide a tuple of table names to indicate which tables' rows should be locked. By default, all fetched
            rows are locked.
        :param no_key:
            If `True`, use the lower SELECT ... FOR NO KEY UPDATE lock strength on PostgreSQL to allow creating or
            deleting rows in other tables that reference the locked rows via foreign keys. The parameter is ignored
            on other backends.
        """
        if self.capabilities.support_for_update:
            queryset = self._clone()
            queryset._select_for_update = True
            queryset._select_for_update_nowait = nowait
            queryset._select_for_update_skip_locked = skip_locked
            queryset._select_for_update_of = set(of)
            queryset._select_for_update_no_key = (
                no_key and self.capabilities.support_for_no_key_update
            )
            return queryset
        return self

    def annotate(self, **kwargs: Expression | Term) -> QuerySet[MODEL]:
        """
        Annotate result with aggregation or function result.

        :raises TypeError: Value of kwarg is expected to be a ``Function`` instance.
        """
        from tortoise.models import get_filters_for_field

        queryset = self._clone()
        for key, annotation in kwargs.items():
            # if not isinstance(annotation, (Function, Term)):
            #     raise TypeError("value is expected to be Function/Term instance")
            queryset._annotations[key] = annotation
            queryset._custom_filters.update(get_filters_for_field(key, None, key))
        return queryset

    def group_by(self, *fields: str) -> QuerySet[MODEL]:
        """
        Make QuerySet returns list of dict or tuple with group by.

        Must call before .values() or .values_list()
        """
        queryset = self._clone()
        queryset._group_bys = fields
        return queryset

    def values_list(self, *fields_: str, flat: bool = False) -> ValuesListQuery[Literal[False]]:
        """
        Make QuerySet returns list of tuples for given args instead of objects.

        If call after `.get()`, `.get_or_none()` or `.first()` return tuples for given args instead of object.

        If ```flat=True`` and only one arg is passed can return flat list or just scalar.

        If no arguments are passed it will default to a tuple containing all fields
        in order of declaration.
        """
        if self._fields_for_select:
            raise ValueError(".values_list() cannot be used with .only()")

        fields_for_select_list = fields_ or [
            field for field in self.model._meta.fields_map if field in self.model._meta.db_fields
        ] + list(self._annotations.keys())
        return ValuesListQuery(
            db=self._db,
            model=self.model,
            q_objects=self._q_objects,
            single=self._single,
            raise_does_not_exist=self._raise_does_not_exist,
            flat=flat,
            fields_for_select_list=fields_for_select_list,
            distinct=self._distinct,
            limit=self._limit,
            offset=self._offset,
            orderings=self._orderings,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            group_bys=self._group_bys,
            force_indexes=self._force_indexes,
            use_indexes=self._use_indexes,
        )

    def values(self, *args: str, **kwargs: str) -> ValuesQuery[Literal[False]]:
        """
        Make QuerySet return dicts instead of objects.

        If called after `.get()`, `.get_or_none()` or `.first()`, returns a dict instead of an object.

        You can specify which fields to include by:
        - Passing field names as positional arguments
        - Using kwargs in the format `field_name='name_in_dict'` to customize the keys in the resulting dict

        If no arguments are passed, it will default to a dict containing all fields.

        :raises FieldError: If duplicate key has been provided.
        """
        if self._fields_for_select:
            raise ValueError(".values() cannot be used with .only()")

        if args or kwargs:
            fields_for_select: dict[str, str] = {}
            for field in args:
                if field in fields_for_select:
                    raise FieldError(f"Duplicate key {field}")
                fields_for_select[field] = field

            for return_as, field in kwargs.items():
                if return_as in fields_for_select:
                    raise FieldError(f"Duplicate key {return_as}")
                fields_for_select[return_as] = field
        else:
            _fields = [
                field
                for field in self.model._meta.fields_map.keys()
                if field in self.model._meta.fields_db_projection.keys()
            ] + list(self._annotations.keys())

            fields_for_select = {field: field for field in _fields}

        return ValuesQuery(
            db=self._db,
            model=self.model,
            q_objects=self._q_objects,
            single=self._single,
            raise_does_not_exist=self._raise_does_not_exist,
            fields_for_select=fields_for_select,
            distinct=self._distinct,
            limit=self._limit,
            offset=self._offset,
            orderings=self._orderings,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            group_bys=self._group_bys,
            force_indexes=self._force_indexes,
            use_indexes=self._use_indexes,
        )

    def delete(self) -> DeleteQuery:
        """
        Delete all objects in QuerySet.
        """
        return DeleteQuery(
            db=self._db,
            model=self.model,
            q_objects=self._q_objects,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            limit=self._limit,
            orderings=self._orderings,
        )

    def update(self, **kwargs: Any) -> UpdateQuery:
        """
        Update all objects in QuerySet with given kwargs.

        .. admonition: Example:

            .. code-block:: py3

                await Employee.filter(occupation='developer').update(salary=5000)

        Will instead of returning a resultset, update the data in the DB itself.
        """
        return UpdateQuery(
            db=self._db,
            model=self.model,
            update_kwargs=kwargs,
            q_objects=self._q_objects,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            limit=self._limit,
            orderings=self._orderings,
        )

    def count(self) -> CountQuery:
        """
        Return count of objects in queryset instead of objects.
        """
        return CountQuery(
            db=self._db,
            model=self.model,
            q_objects=self._q_objects,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            limit=self._limit,
            offset=self._offset,
            force_indexes=self._force_indexes,
            use_indexes=self._use_indexes,
        )

    def exists(self) -> ExistsQuery:
        """
        Return True/False whether queryset exists.
        """
        return ExistsQuery(
            db=self._db,
            model=self.model,
            q_objects=self._q_objects,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            force_indexes=self._force_indexes,
            use_indexes=self._use_indexes,
        )

    def all(self) -> QuerySet[MODEL]:
        """
        Return the whole QuerySet.
        Essentially a no-op except as the only operation.
        """
        return self._clone()

    def raw(self, sql: str) -> RawSQLQuery:
        """
        Return the QuerySet from raw SQL
        """
        return RawSQLQuery(model=self.model, db=self._db, sql=sql)

    def first(self) -> QuerySetSingle[MODEL | None]:
        """
        Limit queryset to one object and return one object instead of list.
        """
        queryset = self._clone()
        return queryset._as_single()

    def last(self) -> QuerySetSingle[MODEL | None]:
        """
        Limit queryset to one object and return the last object instead of list.
        """
        queryset = self._clone()

        if queryset._orderings:
            new_ordering = [
                (field, Order.desc if order_type == Order.asc else Order.asc)
                for field, order_type in queryset._orderings
            ]
        elif pk := self.model._meta.pk:
            new_ordering = [(pk.model_field_name, Order.desc)]
        else:
            raise FieldError(
                f"QuerySet has no ordering and model {self.model.__name__} has no pk defined"
            )
        queryset._orderings = new_ordering
        return queryset._as_single()

    def get(self, *args: Q, **kwargs: Any) -> QuerySetSingle[MODEL]:
        """
        Fetch exactly one object matching the parameters.
        """
        queryset = self.filter(*args, **kwargs)
        queryset._limit = 2
        queryset._single = True
        queryset._raise_does_not_exist = True
        return queryset  # type: ignore

    async def in_bulk(
        self, id_list: Iterable[PRIMARY_KEY], field_name: str
    ) -> dict[PRIMARY_KEY, MODEL]:
        """
        Return a dictionary mapping each of the given IDs to the object with
        that ID. If `id_list` isn't provided, evaluate the entire QuerySet.

        :param id_list: A list of field values
        :param field_name: Must be a unique field
        """
        objs = await self.filter(**{f"{field_name}__in": id_list})
        return {getattr(obj, field_name): obj for obj in objs}

    def bulk_create(
        self,
        objects: Iterable[MODEL],
        batch_size: int | None = None,
        ignore_conflicts: bool = False,
        update_fields: Iterable[str] | None = None,
        on_conflict: Iterable[str] | None = None,
    ) -> BulkCreateQuery[MODEL]:
        """
        This method inserts the provided list of objects into the database in an efficient manner
        (generally only 1 query, no matter how many objects there are).

        :param on_conflict: On conflict index name
        :param update_fields: Update fields when conflicts
        :param ignore_conflicts: Ignore conflicts when inserting
        :param objects: List of objects to bulk create
        :param batch_size: How many objects are created in a single query

        :raises ValueError: If params do not meet specifications
        """
        if ignore_conflicts and update_fields:
            raise ValueError(
                "ignore_conflicts and update_fields are mutually exclusive.",
            )
        if not ignore_conflicts:
            if (update_fields and not on_conflict) or (on_conflict and not update_fields):
                raise ValueError("update_fields and on_conflict need set in same time.")
        return BulkCreateQuery(
            db=self._db,
            model=self.model,
            objects=objects,
            batch_size=batch_size,
            ignore_conflicts=ignore_conflicts,
            update_fields=update_fields,
            on_conflict=on_conflict,
        )

    def bulk_update(
        self,
        objects: Iterable[MODEL],
        fields: Iterable[str],
        batch_size: int | None = None,
    ) -> BulkUpdateQuery[MODEL]:
        """
        Update the given fields in each of the given objects in the database.

        :param objects: List of objects to bulk create
        :param fields: The fields to update
        :param batch_size: How many objects are created in a single query

        :raises ValueError: If objects have no primary key set
        """
        if any(obj.pk is None for obj in objects):
            raise ValueError("All bulk_update() objects must have a primary key set.")
        return BulkUpdateQuery(
            db=self._db,
            model=self.model,
            q_objects=self._q_objects,
            annotations=self._annotations,
            custom_filters=self._custom_filters,
            limit=self._limit,
            orderings=self._orderings,
            objects=objects,
            fields=fields,
            batch_size=batch_size,
        )

    def get_or_none(self, *args: Q, **kwargs: Any) -> QuerySetSingle[MODEL | None]:
        """
        Fetch exactly one object matching the parameters.
        """
        queryset = self.filter(*args, **kwargs)
        queryset._limit = 2
        queryset._single = True
        return queryset  # type: ignore

    def only(self, *fields_for_select: str) -> QuerySet[MODEL]:
        """
        Fetch ONLY the specified fields to create a partial model.

        Persisting changes on the model is allowed only when:

        * All the fields you want to update is specified in ``<model>.save(update_fields=[...])``
        * You included the Model primary key in the `.only(...)``

        To protect against common mistakes we ensure that errors get raised:

        * If you access a field that is not specified, you will get an ``AttributeError``.
        * If you do a ``<model>.save()`` a ``IncompleteInstanceError`` will be raised as the model is, as requested, incomplete.
        * If you do a ``<model>.save(update_fields=[...])`` and you didn't include the primary key in the ``.only(...)``,
          then ``IncompleteInstanceError`` will be raised indicating that updates can't be done without the primary key being known.
        * If you do a ``<model>.save(update_fields=[...])`` and one of the fields in ``update_fields`` was not in the ``.only(...)``,
          then ``IncompleteInstanceError`` as that field is not available to be updated.
        """
        if not fields_for_select:
            raise ValueError(".only() requires at least one field")
        queryset = self._clone()
        queryset._fields_for_select = fields_for_select
        return queryset

    def select_related(self, *fields: str) -> QuerySet[MODEL]:
        """
        Return a new QuerySet instance that will select related objects.

        If fields are specified, they must be ForeignKey fields and only those
        related objects are included in the selection.
        """

        queryset = self._clone()
        for field in fields:
            queryset._select_related.add(field)
        return queryset

    def force_index(self, *index_names: str) -> QuerySet[MODEL]:
        """
        The FORCE INDEX hint acts like USE INDEX (index_list),
        with the addition that a table scan is assumed to be very expensive.
        """
        if self.capabilities.support_index_hint:
            queryset = self._clone()
            for index_name in index_names:
                queryset._force_indexes.add(index_name)
            return queryset
        return self

    def use_index(self, *index_names: str) -> QuerySet[MODEL]:
        """
        The USE INDEX (index_list) hint tells MySQL to use only one of the named indexes to find rows in the table.
        """
        if self.capabilities.support_index_hint:
            queryset = self._clone()
            for index_name in index_names:
                queryset._use_indexes.add(index_name)
            return queryset
        return self

    def prefetch_related(self, *args: str | Prefetch) -> QuerySet[MODEL]:
        """
        Like ``.fetch_related()`` on instance, but works on all objects in QuerySet.

        :raises FieldError: If the field to prefetch on is not a relation, or not found.
        """
        queryset = self._clone()
        queryset._prefetch_map = {}

        for relation in args:
            if isinstance(relation, Prefetch):
                relation.resolve_for_queryset(queryset)
                continue

            first_level_field, __, forwarded_prefetch = relation.partition("__")
            if first_level_field not in self.model._meta.fetch_fields:
                if first_level_field in self.model._meta.fields:
                    raise FieldError(
                        f"Field {first_level_field} on {self.model._meta.full_name} is not a relation"
                    )
                raise FieldError(
                    f"Relation {first_level_field} for {self.model._meta.full_name} not found"
                )
            if first_level_field not in queryset._prefetch_map.keys():
                queryset._prefetch_map[first_level_field] = set()
            if forwarded_prefetch:
                queryset._prefetch_map[first_level_field].add(forwarded_prefetch)
        return queryset

    async def explain(self) -> Any:
        """Fetch and return information about the query execution plan.

        This is done by executing an ``EXPLAIN`` query whose exact prefix depends
        on the database backend, as documented below.

        - PostgreSQL: ``EXPLAIN (FORMAT JSON, VERBOSE) ...``
        - SQLite: ``EXPLAIN QUERY PLAN ...``
        - MySQL: ``EXPLAIN FORMAT=JSON ...``

        .. note::
            This is only meant to be used in an interactive environment for debugging
            and query optimization.
            **The output format may (and will) vary greatly depending on the database backend.**
        """
        self._choose_db_if_not_chosen()
        self._make_query()
        return await self._db.executor_class(model=self.model, db=self._db).execute_explain(
            self.query.get_sql()
        )

    def using_db(self, _db: BaseDBAsyncClient | None) -> QuerySet[MODEL]:
        """
        Executes query in provided db client.
        Useful for transactions workaround.
        """
        queryset = self._clone()
        queryset._db = _db if _db else queryset._db
        return queryset

    def _join_select_related(self, lookup_expression: str) -> tuple[type[Model], Table]:
        fields = expand_lookup_expression(self.model, lookup_expression)
        model = self.model
        table = self.model._meta.basetable
        path: tuple[str | None, ...] = (None,)
        for field in fields:
            field = cast(RelationalField, field)
            path = path + (field.model_field_name,)
            table = self._join_table_by_field(table, field.model_field_name, field)

            # do not select related fields if we are only selecting a subset of fields
            if self._fields_for_select:
                continue

            related_fields = field.related_model._meta.db_fields
            append_item = (
                field.related_model,
                len(related_fields),
                field.model_field_name,
                model,
                path,
            )
            model = field.related_model
            if append_item not in self._select_related_idx:
                self._select_related_idx.append(append_item)
            self.query = self.query.select(
                *[
                    table[related_field].as_(f"{table.get_table_name()}.{related_field}")
                    for related_field in related_fields
                ]
            )
        return model, table

    def _resolve_only(self, only_lookup_expressions: tuple[str, ...]) -> None:
        # Group fields by fetch fields, e.g. ["a__b", "a__c"] -> {"a": ["b", "c"]}.
        # The direct fields of the model are the ones that would have the key "".
        fetch_to_fields = defaultdict(list)
        # the order is important here, we need to process the shallowest fields first
        # because we want to populate _select_related_idx with actual items that need to be
        # selected, not "filler" items tha just tell the executor that an empty instance has
        # to be created
        for expression in sorted(only_lookup_expressions, key=lambda x: x.count("__")):
            fetch_fields_lookup, __, field_name = expression.rpartition("__")
            fetch_to_fields[fetch_fields_lookup].append(field_name)

        # select direct model fields which would have the key "": {"": ["a", "b"]}
        data_fields = fetch_to_fields.pop("", None)
        if data_fields:
            table = self.model._meta.basetable
            self._select_related_idx.append(
                (
                    self.model,
                    len(data_fields),
                    table,
                    self.model,
                    (None,),
                )
            )
            try:
                self.query = self.query.select(
                    *[
                        table[self.model._meta.fields_db_projection[field]].as_(field)
                        for field in data_fields
                        if field not in self._annotations
                    ]
                )
            except KeyError as e:
                raise FieldError(
                    f'Unknown field "{e.args[0]}" for model "{self.model.__name__}"'
                ) from e

        else:
            # even though no data fields are selected, we need to let the executor know
            # that an empty instance of the model has to be created
            self._select_related_idx.append(
                (
                    self.model,
                    0,
                    self.model._meta.basetable,
                    self.model,
                    (None,),
                )
            )

        # Select fields of related models, e.g. {"a": ["b", "c"]}
        added_paths = set()
        for fetch_fields_lookup, data_fields in fetch_to_fields.items():
            fetch_fields = expand_lookup_expression(self.model, fetch_fields_lookup)
            referring_model = model = self.model
            table = self.model._meta.basetable
            path: tuple[str | None, ...] = (None,)
            for i, fetch_field in enumerate(fetch_fields):
                field = cast(RelationalField, fetch_field)
                path = path + (field.model_field_name,)
                table = self._join_table_by_field(table, field.model_field_name, field)
                referring_model = model
                model = field.related_model

                if path in added_paths:
                    continue

                self._select_related_idx.append(
                    (
                        model,
                        # we need 0 items for letting know the executor that instances need to
                        # be created even though no their fields are selected, e.g.
                        # .only("a__b__field")
                        len(data_fields) if i == len(fetch_fields) - 1 else 0,
                        table,
                        referring_model,
                        path,
                    )
                )
                added_paths.add(path)
            self.query = self.query.select(
                *[table[field].as_(f"{table.get_table_name()}.{field}") for field in data_fields]
            )

    def _make_query(self) -> None:
        # clean tmp records first
        self._select_related_idx = []
        self._joined_tables = []
        table = self.model._meta.basetable
        if self._fields_for_select:
            # select .only() fields
            self.query = self.model._meta.basequery.select()
            self._resolve_only(self._fields_for_select)
        else:
            # select all fields
            self.query = copy(self.model._meta.basequery_all_fields)  # type:ignore[assignment]
            append_item = (
                self.model,
                len(self.model._meta.db_fields) + len(self._annotations),
                table,
                self.model,
                (None,),
            )
            self._select_related_idx.append(append_item)
        self.resolve_ordering(
            self.model,
            self.model._meta.basetable,
            self._orderings,
            self._annotations,
            self._fields_for_select,
        )
        self.resolve_filters()
        if self._limit is not None:
            self.query._limit = self.query._wrapper_cls(self._limit)
        if self._offset is not None:
            self.query._offset = self.query._wrapper_cls(self._offset)
        if self._distinct:
            self.query._distinct = True
        if self._select_for_update:
            self.query = self.query.for_update(
                self._select_for_update_nowait,
                self._select_for_update_skip_locked,
                self._select_for_update_of,
                self._select_for_update_no_key,
            )
        if self._select_related:
            for select_related in self._select_related:
                self._join_select_related(select_related)
        if self._force_indexes:
            self.query._force_indexes = []
            self.query = self.query.force_index(*self._force_indexes)
        if self._use_indexes:
            self.query._use_indexes = []
            self.query = self.query.use_index(*self._use_indexes)

    def __await__(self) -> Generator[Any, None, list[MODEL]]:
        if self._db is None:
            self._db = self._choose_db(self._select_for_update)  # type: ignore
        self._make_query()
        return self._execute().__await__()

    async def __aiter__(self) -> AsyncIterator[MODEL]:
        for val in await self:
            yield val

    async def _execute(self) -> list[MODEL]:
        instance_list = await self._db.executor_class(
            model=self.model,
            db=self._db,
            prefetch_map=self._prefetch_map,
            prefetch_queries=self._prefetch_queries,
            select_related_idx=self._select_related_idx,  # type: ignore
        ).execute_select(
            *self.query.get_parameterized_sql(),
            custom_fields=list(self._annotations.keys()),
        )
        if self._single:
            if len(instance_list) == 1:
                return instance_list[0]
            if not instance_list:
                if self._raise_does_not_exist:
                    raise DoesNotExist(self.model)
                return None  # type: ignore
            raise MultipleObjectsReturned(self.model)
        return instance_list


class UpdateQuery(AwaitableQuery):
    __slots__ = (
        "update_kwargs",
        "_orderings",
        "_limit",
        "values",
    )

    def __init__(
        self,
        model: type[MODEL],
        update_kwargs: dict[str, Any],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        limit: int | None,
        orderings: list[tuple[str, str]],
    ) -> None:
        super().__init__(model)
        self.update_kwargs = update_kwargs
        self._q_objects = q_objects
        self._annotations = annotations
        self._custom_filters = custom_filters
        self._db = db
        self._limit = limit
        self._orderings = orderings

    def _make_query(self) -> None:
        table = self.model._meta.basetable
        self.query = self._db.query_class.update(table)
        if self.capabilities.support_update_limit_order_by and self._limit:
            self.query._limit = self.query._wrapper_cls(self._limit)
            self.resolve_ordering(self.model, table, self._orderings, self._annotations)

        self.resolve_filters()
        for key, value in self.update_kwargs.items():
            field_object = self.model._meta.fields_map.get(key)
            if not field_object:
                raise FieldError(f"Unknown keyword argument {key} for model {self.model}")
            if field_object.pk:
                raise IntegrityError(f"Field {key} is PK and can not be updated")
            if field_object.generated:
                raise IntegrityError(f"Field {key} is generated and can not be updated")
            if isinstance(field_object, (ForeignKeyFieldInstance, OneToOneFieldInstance)):
                self.model._validate_relation_type(key, value)
                fk_field: str = field_object.source_field  # type: ignore
                db_field = self.model._meta.fields_map[fk_field].source_field
                value = self.model._meta.fields_map[fk_field].to_db_value(
                    getattr(value, field_object.to_field_instance.model_field_name),
                    None,
                )
            else:
                try:
                    db_field = self.model._meta.fields_db_projection[key]
                except KeyError:
                    raise FieldError(f"Field {key} is virtual and can not be updated")

                if isinstance(value, Expression):
                    value = value.resolve(
                        ResolveContext(
                            model=self.model,
                            table=table,
                            annotations=self._annotations,
                            custom_filters=self._custom_filters,
                        )
                    ).term
                else:
                    value = self.model._meta.fields_map[key].to_db_value(value, None)

            self.query = self.query.set(db_field, value)

    def __await__(self) -> Generator[Any, None, int]:
        self._choose_db_if_not_chosen(True)
        self._make_query()
        return self._execute().__await__()

    async def _execute(self) -> int:
        return (await self._db.execute_query(*self.query.get_parameterized_sql()))[0]


class DeleteQuery(AwaitableQuery):
    __slots__ = (
        "_annotations",
        "_custom_filters",
        "_orderings",
        "_limit",
    )

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        limit: int | None,
        orderings: list[tuple[str, str]],
    ) -> None:
        super().__init__(model)
        self._q_objects = q_objects
        self._annotations = annotations
        self._custom_filters = custom_filters
        self._db = db
        self._limit = limit
        self._orderings = orderings

    def _make_query(self) -> None:
        self.query = copy(self.model._meta.basequery)
        if self.capabilities.support_update_limit_order_by and self._limit:
            self.query._limit = self.query._wrapper_cls(self._limit)
            self.resolve_ordering(
                model=self.model,
                table=self.model._meta.basetable,
                orderings=self._orderings,
                annotations=self._annotations,
            )
        self.resolve_filters()
        self.query._delete_from = True
        return

    def __await__(self) -> Generator[Any, None, int]:
        self._choose_db_if_not_chosen(True)
        self._make_query()
        return self._execute().__await__()

    async def _execute(self) -> int:
        return (await self._db.execute_query(*self.query.get_parameterized_sql()))[0]


class ExistsQuery(AwaitableQuery):
    __slots__ = (
        "_force_indexes",
        "_use_indexes",
    )

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        force_indexes: set[str],
        use_indexes: set[str],
    ) -> None:
        super().__init__(model)
        self._q_objects = q_objects
        self._db = db
        self._annotations = annotations
        self._custom_filters = custom_filters
        self._force_indexes = force_indexes
        self._use_indexes = use_indexes

    def _make_query(self) -> None:
        self.query = copy(self.model._meta.basequery)
        self.resolve_filters()
        self.query._limit = self.query._wrapper_cls(1)
        self.query._select_other(ValueWrapper(1, allow_parametrize=False))  # type:ignore[arg-type]

        if self._force_indexes:
            self.query._force_indexes = []
            self.query = self.query.force_index(*self._force_indexes)
        if self._use_indexes:
            self.query._use_indexes = []
            self.query = self.query.use_index(*self._use_indexes)

    def __await__(self) -> Generator[Any, None, bool]:
        self._choose_db_if_not_chosen()
        self._make_query()
        return self._execute().__await__()

    async def _execute(
        self,
    ) -> bool:
        result, _ = await self._db.execute_query(*self.query.get_parameterized_sql())
        return bool(result)


class CountQuery(AwaitableQuery):
    __slots__ = (
        "_limit",
        "_offset",
        "_force_indexes",
        "_use_indexes",
    )

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        limit: int | None,
        offset: int | None,
        force_indexes: set[str],
        use_indexes: set[str],
    ) -> None:
        super().__init__(model)
        self._q_objects = q_objects
        self._annotations = annotations
        self._custom_filters = custom_filters
        self._limit = limit
        self._offset = offset or 0
        self._db = db
        self._force_indexes = force_indexes
        self._use_indexes = use_indexes

    def _make_query(self) -> None:
        self.query = copy(self.model._meta.basequery)
        self.resolve_filters()
        count_term = Count(Star())
        if self.query._groupbys:
            count_term = count_term.over()

        # remove annotations
        self.query._selects = []
        self.query._select_other(count_term)

        if self._force_indexes:
            self.query._force_indexes = []
            self.query = self.query.force_index(*self._force_indexes)
        if self._use_indexes:
            self.query._use_indexes = []
            self.query = self.query.use_index(*self._use_indexes)

    def __await__(self) -> Generator[Any, None, int]:
        self._choose_db_if_not_chosen()
        self._make_query()
        return self._execute().__await__()

    async def _execute(self) -> int:
        _, result = await self._db.execute_query(*self.query.get_parameterized_sql())
        if not result:
            return 0
        count = list(dict(result[0]).values())[0] - self._offset
        if self._limit and count > self._limit:
            return self._limit
        return count


class FieldSelectQuery(AwaitableQuery):
    # pylint: disable=W0223

    def __init__(self, model: type[MODEL], annotations: dict[str, Any]) -> None:
        super().__init__(model)
        self._annotations = annotations

    def _join_table_with_forwarded_fields(
        self, model: type[MODEL], table: Table, field: str, forwarded_fields: str
    ) -> tuple[Table, str]:
        if field in model._meta.fields_db_projection and not forwarded_fields:
            return table, model._meta.fields_db_projection[field]

        if field in model._meta.fields_db_projection and forwarded_fields:
            raise FieldError(f'Field "{field}" for model "{model.__name__}" is not relation')

        if field in self.model._meta.fetch_fields and not forwarded_fields:
            raise ValueError(
                f'Selecting relation "{field}" is not possible, select concrete '
                "field on related model"
            )

        field_object = cast(RelationalField, model._meta.fields_map.get(field))
        if not field_object:
            raise FieldError(f'Unknown field "{field}" for model "{model.__name__}"')

        table = self._join_table_by_field(table, field, field_object)
        field, __, forwarded_fields_ = forwarded_fields.partition("__")

        return self._join_table_with_forwarded_fields(
            model=field_object.related_model,
            table=table,
            field=field,
            forwarded_fields=forwarded_fields_,
        )

    def add_field_to_select_query(self, field: str, return_as: str) -> None:
        table = self.model._meta.basetable

        if field in self._annotations:
            self._annotations[return_as] = self._annotations[field]
            return

        if field in self.model._meta.fields_db_projection:
            db_field = self.model._meta.fields_db_projection[field]
            self.query._select_field(table[db_field].as_(return_as))
            return

        if field in self.model._meta.fetch_fields:
            raise ValueError(
                f'Selecting relation "{field}" is not possible, select '
                "concrete field on related model"
            )

        field_, __, forwarded_fields = field.partition("__")
        if field_ in self.model._meta.fetch_fields:
            related_table, related_db_field = self._join_table_with_forwarded_fields(
                model=self.model,
                table=table,
                field=field_,
                forwarded_fields=forwarded_fields,
            )
            self.query._select_field(related_table[related_db_field].as_(return_as))
            return

        raise FieldError(f'Unknown field "{field}" for model "{self.model.__name__}"')

    def resolve_to_python_value(self, model: type[MODEL], field: str) -> Callable:
        if field in model._meta.fetch_fields:
            # return as is to get whole model objects
            return lambda x: x

        if field in (x[1] for x in model._meta.db_native_fields):
            return lambda x: x

        if field in self._annotations:
            annotation = self._annotations[field]
            field_object = getattr(annotation, "field_object", None)
            if field_object:
                return field_object.to_python_value
            return lambda x: x

        if field in model._meta.fields_map:
            return model._meta.fields_map[field].to_python_value

        field_, __, forwarded_fields = field.partition("__")
        if field_ in model._meta.fetch_fields:
            new_model = model._meta.fields_map[field_].related_model  # type: ignore
            return self.resolve_to_python_value(new_model, forwarded_fields)

        raise FieldError(f'Unknown field "{field}" for model "{model}"')

    def _resolve_group_bys(self, *field_names: str) -> list:
        group_bys = []
        for field_name in field_names:
            if field_name in self._annotations:
                group_bys.append(Term(field_name))
                continue
            field, __, forwarded_fields = field_name.partition("__")
            related_table, related_db_field = self._join_table_with_forwarded_fields(
                model=self.model,
                table=self.model._meta.basetable,
                field=field,
                forwarded_fields=forwarded_fields,
            )
            field = related_table[related_db_field].as_(
                f"{related_table.get_table_name()}__{field_name}"
            )
            group_bys.append(field)
        return group_bys


class ValuesListQuery(FieldSelectQuery, Generic[SINGLE]):
    __slots__ = (
        "fields",
        "_limit",
        "_offset",
        "_distinct",
        "_orderings",
        "_single",
        "_raise_does_not_exist",
        "_fields_for_select_list",
        "_flat",
        "_group_bys",
        "_force_indexes",
        "_use_indexes",
        "_fields_to_select_sql",
    )

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        single: bool,
        raise_does_not_exist: bool,
        fields_for_select_list: tuple[str, ...] | list[str],
        limit: int | None,
        offset: int | None,
        distinct: bool,
        orderings: list[tuple[str, str]],
        flat: bool,
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        group_bys: tuple[str, ...],
        force_indexes: set[str],
        use_indexes: set[str],
    ) -> None:
        super().__init__(model, annotations)
        if flat and (len(fields_for_select_list) != 1):
            raise TypeError("You can flat value_list only if contains one field")

        fields_for_select = {str(i): field for i, field in enumerate(fields_for_select_list)}
        self.fields = fields_for_select
        self._limit = limit
        self._offset = offset
        self._distinct = distinct
        self._orderings = orderings
        self._custom_filters = custom_filters
        self._q_objects = q_objects
        self._single = single
        self._raise_does_not_exist = raise_does_not_exist
        self._fields_for_select_list = fields_for_select_list
        self._flat = flat
        self._db = db
        self._group_bys = group_bys
        self._force_indexes = force_indexes
        self._use_indexes = use_indexes
        self._fields_to_select_sql = {
            *self._fields_for_select_list,
            *(key for key, value in self.fields.items() if value in self._fields_for_select_list),
        }

    def _make_query(self) -> None:
        self._joined_tables = []

        self.query = copy(self.model._meta.basequery)
        for positional_number, field in self.fields.items():
            self.add_field_to_select_query(field, positional_number)

        self.resolve_ordering(
            model=self.model,
            table=self.model._meta.basetable,
            orderings=self._orderings,
            annotations=self._annotations,
            fields_for_select=self._fields_for_select_list,
        )
        self.resolve_filters(self._fields_to_select_sql)
        if self._limit:
            self.query._limit = self.query._wrapper_cls(self._limit)
        if self._offset:
            self.query._offset = self.query._wrapper_cls(self._offset)
        if self._distinct:
            self.query._distinct = True
        if self._group_bys:
            self.query._groupbys = self._resolve_group_bys(*self._group_bys)

        if self._force_indexes:
            self.query._force_indexes = []
            self.query = self.query.force_index(*self._force_indexes)
        if self._use_indexes:
            self.query._use_indexes = []
            self.query = self.query.use_index(*self._use_indexes)

    @overload
    def __await__(
        self: ValuesListQuery[Literal[False]],
    ) -> Generator[Any, None, list[tuple[Any, ...]]]: ...

    @overload
    def __await__(
        self: ValuesListQuery[Literal[True]],
    ) -> Generator[Any, None, tuple[Any, ...]]: ...

    def __await__(self) -> Generator[Any, None, list[Any] | tuple[Any, ...]]:
        self._choose_db_if_not_chosen()
        self._make_query()
        return self._execute().__await__()  # pylint: disable=E1101

    async def __aiter__(self: ValuesListQuery[Any]) -> AsyncIterator[Any]:
        for val in await self:
            yield val

    async def _execute(self) -> list[Any] | tuple:
        _, result = await self._db.execute_query(*self.query.get_parameterized_sql())
        columns = [
            (key, self.resolve_to_python_value(self.model, name))
            for key, name in self.fields.items()
        ]
        if self._flat:
            func = columns[0][1]
            flatmap = lambda entry: func(entry["0"])  # noqa
            lst_values = list(map(flatmap, result))
        else:
            listmap = lambda entry: tuple(func(entry[column]) for column, func in columns)  # noqa
            lst_values = list(map(listmap, result))

        if self._single:
            if len(lst_values) == 1:
                return lst_values[0]
            if not lst_values:
                if self._raise_does_not_exist:
                    raise DoesNotExist(self.model)
                return None  # type: ignore
            raise MultipleObjectsReturned(self.model)
        return lst_values


class ValuesQuery(FieldSelectQuery, Generic[SINGLE]):
    __slots__ = (
        "_fields_for_select",
        "_limit",
        "_offset",
        "_distinct",
        "_orderings",
        "_single",
        "_raise_does_not_exist",
        "_group_bys",
        "_force_indexes",
        "_use_indexes",
    )

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        single: bool,
        raise_does_not_exist: bool,
        fields_for_select: dict[str, str],
        limit: int | None,
        offset: int | None,
        distinct: bool,
        orderings: list[tuple[str, str]],
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        group_bys: tuple[str, ...],
        force_indexes: set[str],
        use_indexes: set[str],
    ) -> None:
        super().__init__(model, annotations)
        self._fields_for_select = fields_for_select
        self._limit = limit
        self._offset = offset
        self._distinct = distinct
        self._orderings = orderings
        self._custom_filters = custom_filters
        self._q_objects = q_objects
        self._single = single
        self._raise_does_not_exist = raise_does_not_exist
        self._db = db
        self._group_bys = group_bys
        self._force_indexes = force_indexes
        self._use_indexes = use_indexes

    def _make_query(self) -> None:
        self._joined_tables = []

        self.query = copy(self.model._meta.basequery)
        for return_as, field in self._fields_for_select.items():
            self.add_field_to_select_query(field, return_as)

        self.resolve_ordering(
            model=self.model,
            table=self.model._meta.basetable,
            orderings=self._orderings,
            annotations=self._annotations,
            fields_for_select=self._fields_for_select.keys(),
        )
        self.resolve_filters()

        # remove annotations that are not in fields_for_select
        self.query._selects = [
            select for select in self.query._selects if select.alias in self._fields_for_select
        ]

        if self._limit:
            self.query._limit = self.query._wrapper_cls(self._limit)
        if self._offset:
            self.query._offset = self.query._wrapper_cls(self._offset)
        if self._distinct:
            self.query._distinct = True
        if self._group_bys:
            self.query._groupbys = self._resolve_group_bys(*self._group_bys)

        if self._force_indexes:
            self.query._force_indexes = []
            self.query = self.query.force_index(*self._force_indexes)
        if self._use_indexes:
            self.query._use_indexes = []
            self.query = self.query.use_index(*self._use_indexes)

    @overload
    def __await__(
        self: ValuesQuery[Literal[False]],
    ) -> Generator[Any, None, list[dict[str, Any]]]: ...

    @overload
    def __await__(
        self: ValuesQuery[Literal[True]],
    ) -> Generator[Any, None, dict[str, Any]]: ...

    def __await__(
        self,
    ) -> Generator[Any, None, list[dict[str, Any]] | dict[str, Any]]:
        self._choose_db_if_not_chosen()
        self._make_query()
        return self._execute().__await__()  # pylint: disable=E1101

    async def __aiter__(self: ValuesQuery[Any]) -> AsyncIterator[dict[str, Any]]:
        for val in await self:
            yield val

    async def _execute(self) -> list[dict] | dict:
        result = await self._db.execute_query_dict(*self.query.get_parameterized_sql())
        columns = [
            val
            for val in [
                (alias, self.resolve_to_python_value(self.model, field_name))
                for alias, field_name in self._fields_for_select.items()
            ]
            if not isinstance(val[1], types.LambdaType)
        ]

        if columns:
            for row in result:
                for col, func in columns:
                    row[col] = func(row[col])

        if self._single:
            if len(result) == 1:
                return result[0]
            if not result:
                if self._raise_does_not_exist:
                    raise DoesNotExist(self.model)
                return None  # type: ignore
            raise MultipleObjectsReturned(self.model)
        return result


class RawSQLQuery(AwaitableQuery):
    __slots__ = ("_sql", "_db")

    def __init__(self, model: type[MODEL], db: BaseDBAsyncClient, sql: str) -> None:
        super().__init__(model)
        self._sql = sql
        self._db = db

    async def _execute(self) -> Any:
        instance_list = await self._db.executor_class(
            model=self.model,
            db=self._db,
        ).execute_select(RawSQL(self._sql).get_sql(self._db.query_class.SQL_CONTEXT), [])
        return instance_list

    def __await__(self) -> Generator[Any, None, list[MODEL]]:
        self._choose_db_if_not_chosen()
        return self._execute().__await__()


class BulkUpdateQuery(UpdateQuery, Generic[MODEL]):
    __slots__ = ("fields", "_objects", "_batch_size", "_queries")

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        q_objects: list[Q],
        annotations: dict[str, Any],
        custom_filters: dict[str, FilterInfoDict],
        limit: int | None,
        orderings: list[tuple[str, str]],
        objects: Iterable[MODEL],
        fields: Iterable[str],
        batch_size: int | None = None,
    ):
        super().__init__(
            model,
            update_kwargs={},
            db=db,
            q_objects=q_objects,
            annotations=annotations,
            custom_filters=custom_filters,
            limit=limit,
            orderings=orderings,
        )
        self.fields = fields
        self._objects = objects
        self._batch_size = batch_size
        self._queries: list[QueryBuilder] = []

    def _make_queries(self) -> list[tuple[str, list[Any]]]:
        table = self.model._meta.basetable
        self.query = self._db.query_class.update(table)
        if self.capabilities.support_update_limit_order_by and self._limit:
            self.query._limit = self.query._wrapper_cls(self._limit)
            self.resolve_ordering(
                model=self.model,
                table=table,
                orderings=self._orderings,
                annotations=self._annotations,
            )

        self.resolve_filters()
        pk_attr = self.model._meta.pk_attr
        source_pk_attr = self.model._meta.fields_map[pk_attr].source_field or pk_attr
        pk = Field(source_pk_attr)
        for objects_item in chunk(self._objects, self._batch_size):
            query = copy(self.query)
            for field in self.fields:
                case = Case()
                pk_list = []
                for obj in objects_item:
                    pk_value = self.model._meta.fields_map[pk_attr].to_db_value(obj.pk, None)
                    field_obj = obj._meta.fields_map[field]
                    field_value = field_obj.to_db_value(getattr(obj, field), obj)
                    case.when(
                        pk == pk_value,
                        (
                            Cast(
                                self.query._wrapper_cls(field_value),
                                field_obj.get_for_dialect(
                                    self._db.schema_generator.DIALECT, "SQL_TYPE"
                                ),
                            )
                            if self._db.schema_generator.DIALECT == "postgres"
                            else self.query._wrapper_cls(field_value)
                        ),
                    )
                    pk_list.append(pk_value)
                query = query.set(field, case)
                query = query.where(pk.isin(pk_list))
            self._queries.append(query)
        return [query.get_parameterized_sql() for query in self._queries]

    async def _execute_many(self, queries_with_params: list[tuple[str, list[Any]]]) -> int:
        count = 0
        for sql, values in queries_with_params:
            count += (await self._db.execute_query(sql, values))[0]
        return count

    def __await__(self) -> Generator[Any, Any, int]:
        self._choose_db_if_not_chosen(True)
        queries = self._make_queries()
        return self._execute_many(queries).__await__()

    def sql(self, params_inline=False) -> str:
        self._choose_db_if_not_chosen()
        queries = self._make_queries()
        return ";".join([sql for sql, _ in queries])


class BulkCreateQuery(AwaitableQuery, Generic[MODEL]):
    __slots__ = (
        "_objects",
        "_ignore_conflicts",
        "_batch_size",
        "_db",
        "_executor",
        "_update_fields",
        "_on_conflict",
    )

    def __init__(
        self,
        model: type[MODEL],
        db: BaseDBAsyncClient,
        objects: Iterable[MODEL],
        batch_size: int | None = None,
        ignore_conflicts: bool = False,
        update_fields: Iterable[str] | None = None,
        on_conflict: Iterable[str] | None = None,
    ):
        super().__init__(model)
        self._objects = objects
        self._ignore_conflicts = ignore_conflicts
        self._batch_size = batch_size
        self._db = db
        self._update_fields = update_fields
        self._on_conflict = on_conflict

    def _make_queries(self) -> tuple[str, str]:
        self._executor = self._db.executor_class(model=self.model, db=self._db)
        if self._ignore_conflicts or self._update_fields:
            _, columns = self._executor._prepare_insert_columns()
            insert_query = self._executor._prepare_insert_statement(
                columns, ignore_conflicts=self._ignore_conflicts
            )
            insert_query_all = insert_query
            if self.model._meta.generated_db_fields:
                _, columns_all = self._executor._prepare_insert_columns(include_generated=True)
                insert_query_all = self._executor._prepare_insert_statement(
                    columns_all,
                    has_generated=False,
                    ignore_conflicts=self._ignore_conflicts,
                )
            if self._update_fields:
                alias = f"new_{self.model._meta.db_table}"
                insert_query_all = insert_query_all.as_(alias).on_conflict(
                    *(self._on_conflict or [])
                )
                insert_query = insert_query.as_(alias).on_conflict(*(self._on_conflict or []))
                for update_field in self._update_fields:
                    insert_query_all = insert_query_all.do_update(update_field)
                    insert_query = insert_query.do_update(update_field)
            return insert_query.get_sql(), insert_query_all.get_sql()
        else:
            return self._executor.insert_query, self._executor.insert_query_all

    async def _execute_many(self, insert_sql: str, insert_sql_all: str) -> None:
        fields_map = self.model._meta.fields_map
        for instance_chunk in chunk(self._objects, self._batch_size):
            values_lists_all = []
            values_lists = []
            for instance in instance_chunk:
                if instance._custom_generated_pk:
                    values_lists_all.append(
                        [
                            fields_map[field_name].to_db_value(
                                getattr(instance, field_name), instance
                            )
                            for field_name in self._executor.regular_columns_all
                        ]
                    )
                else:
                    values_lists.append(
                        [
                            fields_map[field_name].to_db_value(
                                getattr(instance, field_name), instance
                            )
                            for field_name in self._executor.regular_columns
                        ]
                    )
            if values_lists_all:
                await self._db.execute_many(insert_sql_all, values_lists_all)
            if values_lists:
                await self._db.execute_many(insert_sql, values_lists)

    def __await__(self) -> Generator[Any, None, None]:
        self._choose_db_if_not_chosen(True)
        insert_sql, insert_sql_all = self._make_queries()
        return self._execute_many(insert_sql, insert_sql_all).__await__()

    def sql(self, params_inline=False) -> str:
        self._choose_db_if_not_chosen()
        insert_sql, insert_sql_all = self._make_queries()

        if all(o._custom_generated_pk for o in self._objects):
            return insert_sql_all

        if all(not o._custom_generated_pk for o in self._objects):
            return insert_sql

        return ";".join([insert_sql, insert_sql_all])
