# coding=utf-8
from __future__ import unicode_literals, print_function, absolute_import

import six
import random
import itertools
from django.db import models


class ITableChunk(object):

    def __iter__(self):
        raise NotImplementedError

    def get_pk_start(self):
        raise NotImplementedError

    def get_pk_stop(self):
        raise NotImplementedError


class TableScannerChunk(ITableChunk):

    def __init__(self, data_list, pk_start, pk_stop):
        self._data_list = data_list
        self._pk_start = pk_start
        self._pk_stop = pk_stop

    def __iter__(self):
        return iter(self._data_list)

    def get_pk_start(self):
        return self._pk_start

    def get_pk_stop(self):
        return self._pk_stop


class TableScannerChunkIterator(object):

    def __init__(self, scanner, last_pk, chunk_size):
        assert isinstance(scanner, FiniteTableScanner)
        self._scanner = scanner
        self._last_pk = last_pk
        self._chunk_size = chunk_size

    def __iter__(self):
        while True:
            last_pk = self._last_pk
            data_list, next_last_pk = self._scanner.get_next_data_list(last_pk=last_pk, chunk_size=self._chunk_size)
            self._last_pk = next_last_pk
            yield TableScannerChunk(data_list=data_list, pk_start=last_pk, pk_stop=next_last_pk)

            if next_last_pk is None:
                break


class FiniteTableScanner(object):

    def __init__(self, queryset):
        assert isinstance(queryset, models.QuerySet)
        self._model = queryset.model
        self._query = queryset.query
        self._db_table = self._model._meta.db_table

    @property
    def queryset(self):
        return models.QuerySet(model=self._model, query=self._query)

    def get_first_pk(self):
        try:
            return self.queryset.values_list('pk', flat=True)[0]
        except IndexError:
            return None

    def get_next_data_list(self, last_pk=None, chunk_size=1):
        qs = self.queryset.order_by('pk')
        if last_pk is not None:
            qs = qs.filter(pk__gt=last_pk)
        data_list = list(qs[:chunk_size])
        if len(data_list) == 0:
            next_last_pk = None
        else:
            next_last_pk = data_list[-1].pk
        return data_list, next_last_pk

    def chunks(self, chunk_size):
        return iter(TableScannerChunkIterator(scanner=self, last_pk=None, chunk_size=chunk_size))


def group(iterator, group_size):
    while True:
        segment = []
        for element in iterator:
            segment.append(element)
            if len(segment) >= group_size:
                break
        yield segment
        if len(segment) == 0:
            break

