from math import ceil
from typing import Callable

from django.db.models.query import QuerySet
from gcloudc.db.backends.datastore.expressions import Scatter


def _find_random_keys(queryset: QuerySet, shard_count: int) -> list:
    OVERSAMPLING_FACTOR = 32

    return list(
        queryset.model.objects.order_by(Scatter()).values_list("pk", flat=True)[
            :(shard_count * OVERSAMPLING_FACTOR)
        ]
    )


def sequential_int_key_ranges(queryset, shard_count):
    """
        Given a queryset and a number of shards.
        This function generate key-ranges for a model with
        an integer, dense, sequential primary key, which is usually the default
        when using a SQL backend with autogenerated pks.
    """
    qs = queryset.order_by("pk").values_list("pk", flat=True)
    smallest = qs.first()
    biggest = qs.last()
    max_gap = biggest - smallest
    size = ceil(max_gap / shard_count)
    if biggest < shard_count:
        shard_count = max_gap
        size = 1

    current_min = smallest
    current_max = smallest + size
    key_ranges = [(current_min, current_max)]

    while current_max < biggest:
        current_min += size
        current_max += size
        key_ranges.append((current_min, current_max))

    # If biggest/shard is a whole number, we'd lose the last element (otherwise "ceil" will fix it)
    # e.g. 1000/10 = 100, last range would be (900, 1000), which is off by one
    key_ranges[-1] = (biggest - size, biggest + 1)
    return key_ranges


def datastore_key_ranges(
    queryset: QuerySet,
    shard_count: int,
    random_keys_getter: Callable[[QuerySet, int], list] = _find_random_keys,
) -> list:
    """
        Given a queryset and a number of shard. This function makes use of the
        __scatter__ property to return a list of key ranges for sharded iteration.

        `random_keys_getter` is a callable used to generate random keys.
        It defaults to `_find_random_keys`, but can be used to customise how random keys are
        generated for a given model, e.g. `_find_random_keys` uses the `objects` model manager,
        other implementations can use a different model manager.
        This is especially useful on AppEngine Python 3 which no longer allows `__scatter__` indexes.
    """

    if shard_count > 1:
        # Use the scatter property to generate shard points
        random_keys = random_keys_getter(queryset, shard_count)

        if not random_keys:
            # No random keys? Don't shard
            key_ranges = [(None, None)]
        else:
            random_keys.sort()

            # We have enough random keys to shard things
            if len(random_keys) >= shard_count:
                index_stride = len(random_keys) / float(shard_count)
                split_keys = [random_keys[int(round(index_stride * i))] for i in range(1, shard_count)]
            else:
                split_keys = random_keys

            key_ranges = [(None, split_keys[0])] + [
                (split_keys[i], split_keys[i + 1]) for i in range(len(split_keys) - 1)
            ] + [(split_keys[-1], None)]
    else:
        # Don't shard
        key_ranges = [(None, None)]

    return key_ranges
