from typing import Tuple, Any

from django.db.models.query import BaseIterable, QuerySet


class NestedValuesIterable(BaseIterable):
    """
    Вместо {"user__username": <username>} возвращает {"user": {"username": <username>}}

    Если в списке values есть `a__b`, то обязательно должно быть и `a`, причем после всех a__*
    Иначе при a = None получим {'a': {'b': None}}, а нужно просто {'a': None}
    """

    @staticmethod
    def insert_value_into_nested_dict(obj: dict, keys_path: Tuple[str], value: Any):
        *path, field = keys_path
        for key in path:
            if key not in obj:
                obj[key] = {}
            obj = obj[key]

        if field not in obj or value is None:
            obj[field] = value

    def __iter__(self):
        queryset = self.queryset
        query = queryset.query
        compiler = query.get_compiler(queryset.db)

        field_names = list(query.values_select)
        extra_names = list(query.extra_select)
        annotation_names = list(query.annotation_select)

        # непонятно как аннотации встраивать в словарь, так что не разрешаем такое
        assert not any((extra_names, annotation_names))

        names = [x.split('__') for x in field_names]

        for row in compiler.results_iter():
            serialized_row = {}

            for keys_path, value in zip(names, row):
                self.insert_value_into_nested_dict(serialized_row, keys_path, value)

            yield serialized_row


class NestedValuesMixin:
    def nested_values(self: QuerySet, *fields):
        clone = self._values(*fields)
        clone._iterable_class = NestedValuesIterable
        return clone
