from django.db.models import QuerySet
from django.db.models.expressions import BaseExpression, Combinable
from django.db.models.query import ValuesIterable
from django.db.models.manager import Manager
from .adjuncts import BaseAdjunct
from .handlers import RecordDataclassHandler, RecordHandler


class RecordIterable(ValuesIterable):
    """
        Iterable returned by records() and attached to it's queryset, that yields a record class for each row.
    """

    def __iter__(self):
        queryset = self.queryset
        model = self.queryset.model
        query = queryset.query
        compiler = query.get_compiler(queryset.db)
        record_data = getattr(queryset, '_record_kwargs', {})
        try:
            record_handler = queryset._record
        except AttributeError:
            raise AttributeError("The queryset lacks a _record entry, Is _clone copying that field on the QuerySet class?")

        # extra(select=...) cols are always at the start of the row.
        names = [
            *query.extra_select,
            *query.values_select,
            *query.annotation_select,
        ]
        indexes = range(len(names))

        for row in compiler.results_iter(chunked_fetch=self.chunked_fetch, chunk_size=self.chunk_size):
            dbdata = {names[i]: row[i] for i in indexes}
            # post-processors will be able to rewrite the whole dictionary.
            post_processors = []
            # we overwrite db data bluntly for now. actually we would provide callbacks the current dict.
            for k, v in record_data.items():
                if v.resolves_field:
                    dbdata[k] = v.resolve(model, dbdata)
                if v.post_processing:
                    post_processors.append(v)
            if post_processors:
                for processor in post_processors:
                    processed = processor.post_process(model, dbdata)
                    if processed is not None:
                        dbdata = processed
            yield record_handler.create(**dbdata)


class RecordQuerySetMixin:
    """
        Actual records() implementation.
        As records() calls values(), the queryset is chain-cloned there.
        
        You can mix this in into your QuerySet and Manager classes.
        However make sure in QuerySet classes to implement _clone properly (or use RecordQuerySet instead)
    """
    
    _record_handler = RecordDataclassHandler

    def records(self, *args, **kwargs):
        """
            generates record objects
        
            Acts like values(), however:
                - you can pass a record type or RecordHandler as first argument.
                - if record type is not defined in records(), you have to define it on the queryset, or the model, with _default_record,
                  otherwise it will raise a RuntimeError.
                - keyword arguments of type "Adjunct" are used as deferred values, and resolved independently.
                - values() is called with every required_argument on the dataclass not handled by an Adjunct
        """
        if len(args) and not isinstance(args[0], str):
            # we assume this is our dataclass
            # @TODO better checks.
            handler = args[0]
            args = args[1:]
        else:
            # determine dataclass.
            handler = getattr(self, '_default_record', getattr(self.model, '_default_record', getattr(self, '_record', None)))
            if not handler:
                raise RuntimeError("Trying to records a class without destination class.")

        if not isinstance(handler, RecordHandler):
            handler = self._record_handler.wrap(handler)

        all_keys = [*args, *kwargs.keys()]
        unhandled_keys = list(set(handler.required_arguments) - set(all_keys))
        args = [*args, *unhandled_keys]

        # rebuild keyword arguments for values, by filtering out our adjuncts
        new_kw = {}
        adjuncts = {}
        for k, v in kwargs.items():
            if isinstance(v, BaseAdjunct):
                # skip allows an adjunct to completely ignore a key.
                if not v.skip:
                    adjuncts[k] = v
                # check if we have to add to values. adjuncts can define a field to add here.
                add_to_values = v.values_field()
                if isinstance(add_to_values, str) and add_to_values not in args:
                    args.append(add_to_values)
                elif isinstance(add_to_values, tuple):
                    new_kw[add_to_values[0]] = add_to_values[1]
            elif isinstance(v, BaseExpression) or isinstance(v, Combinable) or hasattr(v, 'resolve_expression'):
                new_kw[k] = v
            else:
                # this will fail in values() for now, but i do not want to hijack future django functionality here.
                new_kw[k] = v

        # copy ourself with values() and save the results on the cloned queryset values produces.
        values = self.values(*args, **new_kw)
        values._iterable_class = RecordIterable
        values._record_kwargs = adjuncts
        values._record = handler
        return values


class RecordQuerySet(RecordQuerySetMixin, QuerySet):
    # overwrite cloning. important.
    def _clone(self):
        c = super()._clone()
        for key in ['_record', '_record_kwargs', '_record_handler', '_default_record']:
            if hasattr(self, key):
                setattr(c, key, getattr(self, key))
        return c

# i use a mixin instead for better clarity with intellisense systems. records is completely safe, as it does not call _chain.
# however you can also simply do:
#class RecordManager(BaseManager.from_queryset(RecordQuerySet)):
#    pass


class RecordManager(RecordQuerySetMixin, Manager):
    def get_queryset(self):
        return RecordQuerySet(self.model, using=self._db)