import json
from dataclasses import dataclass
from json import JSONDecodeError
from typing import Any, ClassVar, Optional, cast

import strawberry
from asgiref.sync import sync_to_async
from django.core.exceptions import ValidationError
from django.db import DEFAULT_DB_ALIAS, models
from django.db.models import Expression, F, OrderBy, Q, QuerySet, Value, Window
from django.db.models.constants import LOOKUP_SEP
from django.db.models.expressions import Col
from django.db.models.functions import RowNumber
from django.db.models.sql.datastructures import BaseTable
from strawberry import Info, relay
from strawberry.relay import NodeType, PageInfo, from_base64
from strawberry.relay.types import NodeIterableType
from strawberry.relay.utils import should_resolve_list_connection_edges
from strawberry.types import get_object_definition
from strawberry.types.base import StrawberryContainer
from strawberry.utils.await_maybe import AwaitableOrValue
from strawberry.utils.inspect import in_async_context
from typing_extensions import Self

from strawberry_django.pagination import apply_window_pagination, get_total_count
from strawberry_django.queryset import get_queryset_config
from strawberry_django.resolvers import django_resolver


def _get_order_by(qs: QuerySet) -> list[OrderBy]:
    return [
        expr
        for expr, _ in qs.query.get_compiler(
            using=qs._db or DEFAULT_DB_ALIAS  # type: ignore
        ).get_order_by()
    ]


@dataclass
class OrderingDescriptor:
    attname: str
    order_by: OrderBy
    # we have to assume everything is nullable by default
    maybe_null: bool = True

    def get_comparator(self, value: Any, before: bool) -> Optional[Q]:
        if value is None:
            # 1. When nulls are first:
            #    1.1 there is nothing before "null"
            #    1.2 after "null" comes everything non-null
            # 2. When nulls are last:
            #    2.1 there is nothing after "null"
            #    2.2 before "null" comes everything non-null
            # => 1.1 and 2.1 require no checks
            # => 1.2 and 2.2 require an "is not null" check
            if bool(self.order_by.nulls_first) ^ before:
                return Q((f"{self.attname}{LOOKUP_SEP}isnull", False))
            return None
        lookup = "lt" if before ^ self.order_by.descending else "gt"
        cmp = Q((f"{self.attname}{LOOKUP_SEP}{lookup}", value))

        if self.maybe_null and bool(self.order_by.nulls_first) == before:
            # if nulls are first, "before any value" can also mean "is null"
            # if nulls are last, "after any value" can also mean "is null"
            cmp |= Q((f"{self.attname}{LOOKUP_SEP}isnull", True))
        return cmp

    def get_eq(self, value) -> Q:
        if value is None:
            return Q((f"{self.attname}{LOOKUP_SEP}isnull", True))
        return Q((f"{self.attname}{LOOKUP_SEP}exact", value))


def annotate_ordering_fields(
    qs: QuerySet,
) -> tuple[QuerySet, list[OrderingDescriptor], list[OrderBy]]:
    annotations = {}
    descriptors = []
    new_defer = None
    new_only = None
    order_bys = _get_order_by(qs)
    pk_in_order = False
    for index, order_by in enumerate(order_bys):
        if isinstance(order_by.expression, Col) and isinstance(
            # Col.alias is missing from django-types
            qs.query.alias_map[order_by.expression.alias],  # type: ignore
            BaseTable,
        ):
            field_name = order_by.expression.field.name
            # if it's a field in the base table, just make sure it is not deferred (e.g. by the optimizer)
            existing, defer = qs.query.deferred_loading
            if defer and field_name in existing:
                # Query is in "defer fields" mode and our field is being deferred
                if new_defer is None:
                    new_defer = set(existing)
                new_defer.discard(field_name)
            elif not defer and field_name not in existing:
                # Query is in "only these fields" mode and our field is not in the set of fields
                if new_only is None:
                    new_only = set(existing)
                new_only.add(field_name)
            descriptors.append(
                OrderingDescriptor(
                    order_by.expression.field.attname,
                    order_by,
                    maybe_null=order_by.expression.field.null,
                )
            )
            if order_by.expression.field.primary_key:
                pk_in_order = True
        else:
            dynamic_field = f"_strawberry_order_field_{index}"
            annotations[dynamic_field] = order_by.expression
            descriptors.append(OrderingDescriptor(dynamic_field, order_by))

    if new_defer is not None:
        # defer is additive, so clear it first
        qs = qs.defer(None).defer(*new_defer)
    elif new_only is not None:
        # only overwrites
        qs = qs.only(*new_only)

    if not pk_in_order:
        # Ensure we always have a clearly defined order by ordering by pk if it isn't in the order already
        # We cannot use QuerySet.order_by, because it validates the order expressions again,
        # but we're operating on the OrderBy expressions which have already been resolved by the compiler
        # In case the user has previously ordered by an aggregate like so:
        # qs.annotate(_c=Count("foo")).order_by("_c")  # noqa: ERA001
        # then the OrderBy we get here would trigger a ValidationError by QuerySet.order_by.
        # But we only want to append to the existing order (and the existing order must be valid already)
        # So this is safe.
        pk_order = F("pk").resolve_expression(qs.query).asc()
        order_bys.append(pk_order)
        descriptors.append(OrderingDescriptor("pk", pk_order, maybe_null=False))
        qs = qs._chain()  # type: ignore
        qs.query.order_by += (pk_order,)
    return qs.annotate(**annotations), descriptors, order_bys


def build_tuple_compare(
    descriptors: list[OrderingDescriptor],
    cursor_values: list[Optional[str]],
    before: bool,
) -> Q:
    current = None
    for descriptor, field_value in zip(reversed(descriptors), reversed(cursor_values)):
        if field_value is None:
            value_expr = None
        else:
            output_field = descriptor.order_by.expression.output_field
            value_expr = Value(field_value, output_field=output_field)
        cmp = descriptor.get_comparator(value_expr, before)
        if current is None:
            current = cmp
        else:
            eq = descriptor.get_eq(value_expr)
            current = cmp | (eq & current) if cmp is not None else eq & current
    return current if current is not None else Q()


class AttrHelper:
    pass


def _extract_expression_value(
    model: models.Model, expr: Expression, attname: str
) -> Optional[str]:
    output_field = expr.output_field
    # Unfortunately Field.value_to_string operates on the object, not a direct value
    # So we have to potentially construct a fake object
    # If the output field's attname doesn't match, we have to construct a fake object
    # Additionally, the output field may not have an attname at all
    # if expressions are used
    field_attname = getattr(output_field, "attname", None)
    if not field_attname:
        # If the field doesn't have an attname, it's a dynamically constructed field,
        # for the purposes of output_field in an expression. Just set its attname, it doesn't hurt anything
        output_field.attname = field_attname = attname
    obj: Any
    if field_attname != attname:
        obj = AttrHelper()
        setattr(obj, output_field.attname, getattr(model, attname))
    else:
        obj = model
    value = output_field.value_from_object(obj)
    if value is None:
        return None
    # value_to_string is missing from django-types
    return output_field.value_to_string(obj)  # type: ignore


def apply_cursor_pagination(
    qs: QuerySet,
    *,
    related_field_id: Optional[str] = None,
    info: Info,
    before: Optional[str],
    after: Optional[str],
    first: Optional[int],
    last: Optional[int],
    max_results: Optional[int],
) -> tuple[QuerySet, list[OrderingDescriptor]]:
    max_results = (
        max_results if max_results is not None else info.schema.config.relay_max_results
    )

    qs, ordering_descriptors, original_order_by = annotate_ordering_fields(qs)
    if after:
        after_cursor = OrderedCollectionCursor.from_cursor(after, ordering_descriptors)
        qs = qs.filter(
            build_tuple_compare(ordering_descriptors, after_cursor.field_values, False)
        )
    if before:
        before_cursor = OrderedCollectionCursor.from_cursor(
            before, ordering_descriptors
        )
        qs = qs.filter(
            build_tuple_compare(ordering_descriptors, before_cursor.field_values, True)
        )

    slice_: Optional[slice] = None
    if first is not None and last is not None:
        if last > max_results:
            raise ValueError(f"Argument 'last' cannot be higher than {max_results}.")
        # if first and last are given, we have to
        # - reverse the order in the DB so we can use slicing to apply [:last],
        #   otherwise we would have to know the total count to apply slicing from the end
        # - We still need to apply forward-direction [:first] slicing, and according to the Relay spec,
        #   it shall happen before [:last] slicing. To do this, we use a window function with a RowNumber ordered
        #   in the original direction, which is opposite the actual query order.
        #   This query is likely not very efficient, but using last _and_ first together is discouraged by the
        #   spec anyway
        qs = (
            qs.reverse()
            .annotate(
                _strawberry_row_number_fwd=Window(
                    RowNumber(),
                    order_by=original_order_by,
                ),
            )
            .filter(
                _strawberry_row_number_fwd__lte=first + 1,
            )
        )
        # we're overfetching by two, in both directions
        slice_ = slice(last + 2)
    elif first is not None:
        if first < 0:
            raise ValueError("Argument 'first' must be a non-negative integer.")
        if first > max_results:
            raise ValueError(f"Argument 'first' cannot be higher than {max_results}.")
        slice_ = slice(first + 1)
    elif last is not None:
        # when using last, optimize by reversing the QuerySet ordering in the DB,
        # then slicing from the end (which is now the start in QuerySet ordering)
        # and then iterating the results in reverse to restore the original order
        if last < 0:
            raise ValueError("Argument 'last' must be a non-negative integer.")
        if last > max_results:
            raise ValueError(f"Argument 'last' cannot be higher than {max_results}.")
        slice_ = slice(last + 1)
        qs = qs.reverse()
    if related_field_id is not None:
        # we always apply window pagination for nested connections,
        # because we want its total count annotation
        offset = slice_.start or 0 if slice_ is not None else 0
        qs = apply_window_pagination(
            qs,
            related_field_id=related_field_id,
            offset=offset,
            limit=slice_.stop - offset if slice_ is not None else None,
        )
    elif slice_ is not None:
        qs = qs[slice_]

    get_queryset_config(qs).ordering_descriptors = ordering_descriptors

    return qs, ordering_descriptors


@dataclass
class OrderedCollectionCursor:
    field_values: list[Any]

    @classmethod
    def from_model(
        cls, model: models.Model, descriptors: list[OrderingDescriptor]
    ) -> Self:
        values = [
            _extract_expression_value(
                model, descriptor.order_by.expression, descriptor.attname
            )
            for descriptor in descriptors
        ]
        return cls(field_values=values)

    @classmethod
    def from_cursor(cls, cursor: str, descriptors: list[OrderingDescriptor]) -> Self:
        type_, values_json = from_base64(cursor)
        if type_ != DjangoCursorEdge.CURSOR_PREFIX:
            raise ValueError("Invalid cursor")
        try:
            string_values = json.loads(values_json)
        except JSONDecodeError as e:
            raise ValueError("Invalid cursor") from e
        if (
            not isinstance(string_values, list)
            or len(string_values) != len(descriptors)
            or any(not (v is None or isinstance(v, str)) for v in string_values)
        ):
            raise ValueError("Invalid cursor")

        try:
            decoded_values = [
                d.order_by.expression.output_field.to_python(v)
                for d, v in zip(descriptors, string_values)
            ]
        except ValidationError as e:
            raise ValueError("Invalid cursor") from e

        return cls(decoded_values)

    def __str__(self):
        return json.dumps(self.field_values, separators=(",", ":"))


@strawberry.type(name="CursorEdge", description="An edge in a connection.")
class DjangoCursorEdge(relay.Edge[relay.NodeType]):
    CURSOR_PREFIX: ClassVar[str] = "orderedcursor"


@strawberry.type(
    name="CursorConnection", description="A connection to a list of items."
)
class DjangoCursorConnection(relay.Connection[relay.NodeType]):
    total_count_qs: strawberry.Private[Optional[QuerySet]] = None
    edges: list[DjangoCursorEdge[NodeType]] = strawberry.field(  # type: ignore
        description="Contains the nodes in this connection"
    )

    @strawberry.field(description="Total quantity of existing nodes.")
    @django_resolver
    def total_count(self) -> int:
        assert self.total_count_qs is not None

        return get_total_count(self.total_count_qs)

    @classmethod
    def resolve_connection(
        cls,
        nodes: NodeIterableType[NodeType],
        *,
        info: Info,
        before: Optional[str] = None,
        after: Optional[str] = None,
        first: Optional[int] = None,
        last: Optional[int] = None,
        max_results: Optional[int] = None,
        **kwargs: Any,
    ) -> AwaitableOrValue[Self]:
        from strawberry_django.optimizer import is_optimized_by_prefetching

        if not isinstance(nodes, QuerySet):
            raise TypeError("DjangoCursorConnection requires a QuerySet")
        total_count_qs: QuerySet = nodes
        qs: QuerySet
        if not is_optimized_by_prefetching(nodes):
            qs, ordering_descriptors = apply_cursor_pagination(
                nodes,
                info=info,
                before=before,
                after=after,
                first=first,
                last=last,
                max_results=max_results,
            )
        else:
            qs = nodes
            ordering_descriptors = get_queryset_config(qs).ordering_descriptors
            assert ordering_descriptors is not None

        type_def = get_object_definition(cls)
        assert type_def
        field_def = type_def.get_field("edges")
        assert field_def

        field = field_def.resolve_type(type_definition=type_def)
        while isinstance(field, StrawberryContainer):
            field = field.of_type

        edge_class = cast("DjangoCursorEdge[NodeType]", field)

        if not should_resolve_list_connection_edges(info):
            return cls(
                edges=[],
                total_count_qs=total_count_qs,
                page_info=PageInfo(
                    start_cursor=None,
                    end_cursor=None,
                    has_previous_page=False,
                    has_next_page=False,
                ),
            )

        def finish_resolving():
            nonlocal qs
            has_previous_page = has_next_page = False

            results = list(qs)

            if first is not None:
                if last is None:
                    has_next_page = len(results) > first
                    results = results[:first]
                # we're paginating forwards _and_ backwards
                # remove the (potentially) overfetched row in the forwards direction first
                elif (
                    results
                    and getattr(results[0], "_strawberry_row_number_fwd", 0) > first
                ):
                    has_next_page = True
                    results = results[1:]

            if last is not None:
                has_previous_page = len(results) > last
                results = results[:last]

            it = reversed(results) if last is not None else results

            edges = [
                edge_class.resolve_edge(
                    cls.resolve_node(v, info=info, **kwargs),
                    cursor=OrderedCollectionCursor.from_model(v, ordering_descriptors),
                )
                for v in it
            ]

            return cls(
                edges=edges,
                total_count_qs=total_count_qs,
                page_info=PageInfo(
                    start_cursor=edges[0].cursor if edges else None,
                    end_cursor=edges[-1].cursor if edges else None,
                    has_previous_page=has_previous_page,
                    has_next_page=has_next_page,
                ),
            )

        if in_async_context() and qs._result_cache is None:  # type: ignore
            return sync_to_async(finish_resolving)()
        return finish_resolving()