from django.db.models import Q

from ishtar_common.serializers_utils import generic_get_results, archive_serialization
from archaeological_finds import models


FIND_MODEL_LIST = [models.BaseFind, models.Find]

# TODO: associated documents, property, findbasket, treatments


def generate_warehouse_queryset(ids):
    q_find, q_basefind = None, None

    for container_key in ("container", "container_ref"):
        for warehouse_key in ("location", "responsible"):
            q_f = Q(**{"{}__{}__id__in".format(container_key, warehouse_key): ids})
            q_bf = Q(
                **{"find__{}__{}__id__in".format(container_key, warehouse_key): ids}
            )
        if not q_find:
            q_find = q_f
            q_basefind = q_bf
        else:
            q_find |= q_f
            q_basefind |= q_bf

    result_queryset = {
        models.BaseFind.__name__: models.BaseFind.objects.filter(q_basefind),
        models.Find.__name__: models.Find.objects.filter(q_find),
    }
    return result_queryset


def find_serialization(
    archive=False,
    return_empty_types=False,
    archive_name=None,
    operation_queryset=None,
    site_queryset=None,
    cr_queryset=None,
    find_queryset=None,
    warehouse_queryset=None,
    get_queryset=False,
    no_geo=True,
    put_locks=False,
    lock_user=None,
):
    result_queryset = {}
    if operation_queryset:
        operation_ids = operation_queryset.values_list("id", flat=True)
        result_queryset = {
            models.BaseFind.__name__: models.BaseFind.objects.filter(
                context_record__operation_id__in=operation_ids
            ),
            models.Find.__name__: models.Find.objects.filter(
                base_finds__context_record__operation_id__in=operation_ids
            ),
        }
    elif site_queryset:
        sites = site_queryset.values_list("id", flat=True)
        f_q = {
            "base_finds__context_record__operation__archaeological_sites__id__in": sites
        }
        result_queryset = {
            models.BaseFind.__name__: models.BaseFind.objects.filter(
                context_record__operation__archaeological_sites__id__in=sites
            ),
            models.Find.__name__: models.Find.objects.filter(**f_q),
        }
    elif cr_queryset:
        cr_ids = cr_queryset.values_list("id", flat=True)
        result_queryset = {
            models.BaseFind.__name__: models.BaseFind.objects.filter(
                context_record__in=cr_ids
            ),
            models.Find.__name__: models.Find.objects.filter(
                base_finds__context_record__in=cr_ids
            ),
        }
    elif find_queryset:
        find_ids = find_queryset.values_list("id", flat=True)
        result_queryset = {
            models.BaseFind.__name__: models.BaseFind.objects.filter(
                find__id__in=find_ids
            ),
            models.Find.__name__: find_queryset,
        }
    elif warehouse_queryset:
        warehouse_ids = warehouse_queryset.values_list("id", flat=True)
        result_queryset = generate_warehouse_queryset(warehouse_ids)

    if get_queryset:
        return result_queryset

    result = generic_get_results(
        FIND_MODEL_LIST, "finds", result_queryset=result_queryset, no_geo=no_geo
    )
    if put_locks:
        for model in FIND_MODEL_LIST:
            if not hasattr(model, "locked"):
                continue
            q = model.objects
            if result_queryset and model.__name__ in result_queryset:
                q = result_queryset[model.__name__]
            q.update(locked=True, lock_user=lock_user)

    full_archive = archive_serialization(
        result,
        archive_dir="operations",
        archive=archive,
        return_empty_types=return_empty_types,
        archive_name=archive_name,
    )
    return full_archive