# Copyright 2014 Google Inc. All rights reserved.
#
# 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.

"""Create / interact with a batch of updates / deletes.

Batches provide the ability to execute multiple operations
in a single request to the Cloud Datastore API.

See
https://cloud.google.com/datastore/docs/concepts/entities#Datastore_Batch_operations
"""

from gcloud.datastore import helpers
from gcloud.datastore._generated import datastore_pb2 as _datastore_pb2


class Batch(object):
    """An abstraction representing a collected group of updates / deletes.

    Used to build up a bulk mutuation.

    For example, the following snippet of code will put the two ``save``
    operations and the ``delete`` operation into the same mutation, and send
    them to the server in a single API request::

      >>> from gcloud import datastore
      >>> client = datastore.Client()
      >>> batch = client.batch()
      >>> batch.put(entity1)
      >>> batch.put(entity2)
      >>> batch.delete(key3)
      >>> batch.commit()

    You can also use a batch as a context manager, in which case
    :meth:`commit` will be called automatically if its block exits without
    raising an exception::

      >>> with batch:
      ...     batch.put(entity1)
      ...     batch.put(entity2)
      ...     batch.delete(key3)

    By default, no updates will be sent if the block exits with an error::

      >>> with batch:
      ...     do_some_work(batch)
      ...     raise Exception()  # rolls back

    :type client: :class:`gcloud.datastore.client.Client`
    :param client: The client used to connect to datastore.
    """

    _id = None  # "protected" attribute, always None for non-transactions

    _INITIAL = 0
    """Enum value for _INITIAL status of batch/transaction."""

    _IN_PROGRESS = 1
    """Enum value for _IN_PROGRESS status of batch/transaction."""

    _ABORTED = 2
    """Enum value for _ABORTED status of batch/transaction."""

    _FINISHED = 3
    """Enum value for _FINISHED status of batch/transaction."""

    def __init__(self, client):
        self._client = client
        self._commit_request = _datastore_pb2.CommitRequest()
        self._partial_key_entities = []
        self._status = self._INITIAL

    def current(self):
        """Return the topmost batch / transaction, or None."""
        return self._client.current_batch

    @property
    def project(self):
        """Getter for project in which the batch will run.

        :rtype: :class:`str`
        :returns: The project in which the batch will run.
        """
        return self._client.project

    @property
    def namespace(self):
        """Getter for namespace in which the batch will run.

        :rtype: :class:`str`
        :returns: The namespace in which the batch will run.
        """
        return self._client.namespace

    @property
    def connection(self):
        """Getter for connection over which the batch will run.

        :rtype: :class:`gcloud.datastore.connection.Connection`
        :returns: The connection over which the batch will run.
        """
        return self._client.connection

    def _add_partial_key_entity_pb(self):
        """Adds a new mutation for an entity with a partial key.

        :rtype: :class:`gcloud.datastore._generated.entity_pb2.Entity`
        :returns: The newly created entity protobuf that will be
                  updated and sent with a commit.
        """
        new_mutation = self.mutations.add()
        return new_mutation.insert

    def _add_complete_key_entity_pb(self):
        """Adds a new mutation for an entity with a completed key.

        :rtype: :class:`gcloud.datastore._generated.entity_pb2.Entity`
        :returns: The newly created entity protobuf that will be
                  updated and sent with a commit.
        """
        # We use ``upsert`` for entities with completed keys, rather than
        # ``insert`` or ``update``, in order not to create race conditions
        # based on prior existence / removal of the entity.
        new_mutation = self.mutations.add()
        return new_mutation.upsert

    def _add_delete_key_pb(self):
        """Adds a new mutation for a key to be deleted.

        :rtype: :class:`gcloud.datastore._generated.entity_pb2.Key`
        :returns: The newly created key protobuf that will be
                  deleted when sent with a commit.
        """
        new_mutation = self.mutations.add()
        return new_mutation.delete

    @property
    def mutations(self):
        """Getter for the changes accumulated by this batch.

        Every batch is committed with a single commit request containing all
        the work to be done as mutations. Inside a batch, calling :meth:`put`
        with an entity, or :meth:`delete` with a key, builds up the request by
        adding a new mutation. This getter returns the protobuf that has been
        built-up so far.

        :rtype: iterable
        :returns: The list of :class:`._generated.datastore_pb2.Mutation`
                  protobufs to be sent in the commit request.
        """
        return self._commit_request.mutations

    def put(self, entity):
        """Remember an entity's state to be saved during :meth:`commit`.

        .. note::
           Any existing properties for the entity will be replaced by those
           currently set on this instance.  Already-stored properties which do
           not correspond to keys set on this instance will be removed from
           the datastore.

        .. note::
           Property values which are "text" ('unicode' in Python2, 'str' in
           Python3) map to 'string_value' in the datastore;  values which are
           "bytes" ('str' in Python2, 'bytes' in Python3) map to 'blob_value'.

        When an entity has a partial key, calling :meth:`commit` sends it as
        an ``insert`` mutation and the key is completed. On return,
        the key for the ``entity`` passed in is updated to match the key ID
        assigned by the server.

        :type entity: :class:`gcloud.datastore.entity.Entity`
        :param entity: the entity to be saved.

        :raises: ValueError if entity has no key assigned, or if the key's
                 ``project`` does not match ours.
        """
        if entity.key is None:
            raise ValueError("Entity must have a key")

        if self.project != entity.key.project:
            raise ValueError("Key must be from same project as batch")

        if entity.key.is_partial:
            entity_pb = self._add_partial_key_entity_pb()
            self._partial_key_entities.append(entity)
        else:
            entity_pb = self._add_complete_key_entity_pb()

        _assign_entity_to_pb(entity_pb, entity)

    def delete(self, key):
        """Remember a key to be deleted during :meth:`commit`.

        :type key: :class:`gcloud.datastore.key.Key`
        :param key: the key to be deleted.

        :raises: ValueError if key is not complete, or if the key's
                 ``project`` does not match ours.
        """
        if key.is_partial:
            raise ValueError("Key must be complete")

        if self.project != key.project:
            raise ValueError("Key must be from same project as batch")

        key_pb = key.to_protobuf()
        self._add_delete_key_pb().CopyFrom(key_pb)

    def begin(self):
        """Begins a batch.

        This method is called automatically when entering a with
        statement, however it can be called explicitly if you don't want
        to use a context manager.

        Overridden by :class:`gcloud.datastore.transaction.Transaction`.

        :raises: :class:`ValueError` if the batch has already begun.
        """
        if self._status != self._INITIAL:
            raise ValueError('Batch already started previously.')
        self._status = self._IN_PROGRESS

    def _commit(self):
        """Commits the batch.

        This is called by :meth:`commit`.
        """
        # NOTE: ``self._commit_request`` will be modified.
        _, updated_keys = self.connection.commit(
            self.project, self._commit_request, self._id)
        # If the back-end returns without error, we are guaranteed that
        # :meth:`Connection.commit` will return keys that match (length and
        # order) directly ``_partial_key_entities``.
        for new_key_pb, entity in zip(updated_keys,
                                      self._partial_key_entities):
            new_id = new_key_pb.path[-1].id
            entity.key = entity.key.completed_key(new_id)

    def commit(self):
        """Commits the batch.

        This is called automatically upon exiting a with statement,
        however it can be called explicitly if you don't want to use a
        context manager.
        """
        try:
            self._commit()
        finally:
            self._status = self._FINISHED

    def rollback(self):
        """Rolls back the current batch.

        Marks the batch as aborted (can't be used again).

        Overridden by :class:`gcloud.datastore.transaction.Transaction`.
        """
        self._status = self._ABORTED

    def __enter__(self):
        self._client._push_batch(self)
        self.begin()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type is None:
                self.commit()
            else:
                self.rollback()
        finally:
            self._client._pop_batch()


def _assign_entity_to_pb(entity_pb, entity):
    """Copy ``entity`` into ``entity_pb``.

    Helper method for ``Batch.put``.

    :type entity_pb: :class:`gcloud.datastore._generated.entity_pb2.Entity`
    :param entity_pb: The entity owned by a mutation.

    :type entity: :class:`gcloud.datastore.entity.Entity`
    :param entity: The entity being updated within the batch / transaction.
    """
    bare_entity_pb = helpers.entity_to_protobuf(entity)
    bare_entity_pb.key.CopyFrom(bare_entity_pb.key)
    entity_pb.CopyFrom(bare_entity_pb)