"""Custom query filters for the Stock models."""

from typing import Optional

from django.db.models import F, Func, IntegerField, OuterRef, Q, Subquery
from django.db.models.functions import Coalesce

import stock.models


def annotate_location_items(filter: Optional[Q] = None):
    """Construct a queryset annotation which returns the number of stock items in a particular location.

    - Includes items in subcategories also
    - Requires subquery to perform annotation
    """
    # Construct a subquery to provide all items in this location and any sublocations
    subquery = stock.models.StockItem.objects.exclude(location=None).filter(
        location__tree_id=OuterRef('tree_id'),
        location__lft__gte=OuterRef('lft'),
        location__rght__lte=OuterRef('rght'),
        location__level__gte=OuterRef('level'),
    )

    # Optionally apply extra filter to returned results
    if filter is not None:
        subquery = subquery.filter(filter)

    return Coalesce(
        Subquery(
            subquery
            .annotate(
                total=Func(F('pk'), function='COUNT', output_field=IntegerField())
            )
            .values('total')
            .order_by()
        ),
        0,
        output_field=IntegerField(),
    )


def annotate_sub_locations():
    """Construct a queryset annotation which returns the number of sub-locations below a certain StockLocation node in a StockLocation tree."""
    subquery = stock.models.StockLocation.objects.filter(
        tree_id=OuterRef('tree_id'),
        lft__gt=OuterRef('lft'),
        rght__lt=OuterRef('rght'),
        level__gt=OuterRef('level'),
    )

    return Coalesce(
        Subquery(
            subquery
            .annotate(
                count=Func(F('pk'), function='COUNT', output_field=IntegerField())
            )
            .values('count')
            .order_by()
        ),
        0,
        output_field=IntegerField(),
    )
