#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Implements a Cloud Datastore query splitter.

For internal use only. No backwards compatibility guarantees.
"""
# pytype: skip-file

from apache_beam.io.gcp.datastore.v1new import types
from apache_beam.options.value_provider import ValueProvider

__all__ = ['QuerySplitterError', 'SplitNotPossibleError', 'get_splits']

SCATTER_PROPERTY_NAME = '__scatter__'
KEY_PROPERTY_NAME = '__key__'
# The number of keys to sample for each split.
KEYS_PER_SPLIT = 32


class QuerySplitterError(Exception):
  """Top-level error type."""


class SplitNotPossibleError(QuerySplitterError):
  """Raised when some parameter of the query does not allow splitting."""


def get_splits(client, query, num_splits):
  """Returns a list of sharded queries for the given Cloud Datastore query.

  This will create up to the desired number of splits, however it may return
  less splits if the desired number of splits is unavailable. This will happen
  if the number of split points provided by the underlying Datastore is less
  than the desired number, which will occur if the number of results for the
  query is too small.

  This implementation of the QuerySplitter uses the __scatter__ property to
  gather random split points for a query.

  Note: This implementation is derived from the java query splitter in
  https://github.com/GoogleCloudPlatform/google-cloud-datastore/blob/master/java/datastore/src/main/java/com/google/datastore/v1/client/QuerySplitterImpl.java

  Args:
    client: the datastore client.
    query: the query to split.
    num_splits: the desired number of splits.

  Returns:
    A list of split queries, of a max length of `num_splits`

  Raises:
    QuerySplitterError: if split could not be performed owing to query or split
      parameters.
  """
  if num_splits <= 1:
    raise SplitNotPossibleError('num_splits must be > 1, got: %d' % num_splits)
  validate_split(query)

  splits = []
  client_scatter_keys = _get_scatter_keys(client, query, num_splits)
  last_client_key = None
  for next_client_key in _get_split_key(client_scatter_keys, num_splits):
    splits.append(_create_split(last_client_key, next_client_key, query))
    last_client_key = next_client_key

  splits.append(_create_split(last_client_key, None, query))
  return splits


def validate_split(query):
  """
  Verifies that the given query can be properly scattered.