# Copyright 2014 Google LLC
#
# Licensed 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.

"""Script to populate datastore with system test data."""

from __future__ import print_function

import os
import sys

import six

from google.cloud import datastore


FETCH_MAX = 20
ALL_KINDS = (
    "Character",
    "Company",
    "Kind",
    "Person",
    "Post",
    "uuid_key",
    "timestamp_key",
)
TRANSACTION_MAX_GROUPS = 5
MAX_DEL_ENTITIES = 500


def print_func(message):
    if os.getenv("GOOGLE_CLOUD_NO_PRINT") != "true":
        print(message)


def get_ancestors(entities):
    # NOTE: A key will always have at least one path element.
    key_roots = [entity.key.flat_path[:2] for entity in entities]
    # Return the unique roots.
    return list(set(key_roots))


def delete_chunks(client, results):
    while results:
        chunk, results = results[:MAX_DEL_ENTITIES], results[MAX_DEL_ENTITIES:]
        client.delete_multi([result.key for result in chunk])


def remove_kind(kind, client):
    query = client.query(kind=kind)
    query.keys_only()
    results = list(query.fetch())

    if not results:
        return

    delete_outside_transaction = False
    with client.transaction():
        # Now that we have all results, we seek to delete.
        print_func("Deleting keys:")
        print_func(results)

        ancestors = get_ancestors(results)
        if len(ancestors) > TRANSACTION_MAX_GROUPS:
            delete_outside_transaction = True
        else:
            delete_chunks(client, results)

    if delete_outside_transaction:
        delete_chunks(client, results)


def main():
    client = datastore.Client()
    kinds = sys.argv[1:]

    if len(kinds) == 0:
        kinds = ALL_KINDS

    print_func("This command will remove all entities for " "the following kinds:")
    print_func("\n".join("- " + val for val in kinds))
    response = six.moves.input("Is this OK [y/n]? ")

    if response.lower() == "y":

        for kind in kinds:
            remove_kind(kind, client)

    else:
        print_func("Doing nothing.")


if __name__ == "__main__":
    main()
