#
# Copyright 2013 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.
#
"""googledatastore connection."""

import logging
import httplib2

from googledatastore import helper
from google.datastore.v1beta3 import datastore_pb2
from google.protobuf import timestamp_pb2
from google.rpc import code_pb2
from google.rpc import status_pb2
from google.type import latlng_pb2

__all__ = [
    'Datastore',
    'Error',
    'RPCError',
    'AuthError',
    'BadArgumentError'
]


class Datastore(object):
  """Datastore client connection constructor."""

  def __init__(self, project_id=None, credentials=None, project_endpoint=None):
    """Datastore client connection constructor.

    Args:
      project_id: the Cloud project to use. Exactly one of
          project_endpoint and project_id must be set.
      credentials: oauth2client.Credentials to authorize the
          connection, default to no credentials.
      project_endpoint: the Datastore endpoint to use. Exactly one of
          project_endpoint and project_id must be set.

    Usage: demos/trivial.py for example usages.

    Raises:
      TypeError: when neither or both of project_endpoint and project_id
      are set.
    """
    self._http = httplib2.Http()
    if not project_endpoint and not project_id:
      raise TypeError('project_endpoint or project_id argument is required.')
    if project_endpoint and project_id:
      raise TypeError('only one of project_endpoint or project_id argument '
                      'is allowed.')

    self._url = (project_endpoint
                 or helper.get_project_endpoint_from_env(project_id=project_id))

    if credentials:
      self._credentials = credentials
      credentials.authorize(self._http)
    else:
      logging.warning('no datastore credentials')

  def lookup(self, request):
    """Lookup entities by key.

    Args:
      request: LookupRequest proto message.

    Returns:
      LookupResponse proto message.

    Raises:
      RPCError: The underlying RPC call failed with an HTTP error.
          (See: .response attribute)
    """
    return self._call_method('lookup', request,
                             datastore_pb2.LookupResponse)

  def run_query(self, request):
    """Query for entities.

    Args:
      request: RunQueryRequest proto message.

    Returns:
      RunQueryResponse proto message.

    Raises:
      RPCError: The underlying RPC call failed with an HTTP error.
          (See: .response attribute)
    """
    return self._call_method('runQuery', request,
                             datastore_pb2.RunQueryResponse)

  def begin_transaction(self, request):
    """Begin a new transaction.

    Args:
      request: BeginTransactionRequest proto message.

    Returns:
      BeginTransactionResponse proto message.

    Raises:
      RPCError: The underlying RPC call failed with an HTTP error.
          (See: .response attribute)
    """
    return self._call_method('beginTransaction', request,
                             datastore_pb2.BeginTransactionResponse)

  def commit(self, request):
    """Commit a mutation, transaction or mutation in a transaction.

    Args:
      request: CommitRequest proto message.

    Returns:
      CommitResponse proto message.

    Raises:
      RPCError: The underlying RPC call failed with an HTTP error.
          (See: .response attribute)
    """
    return self._call_method('commit', request,
                             datastore_pb2.CommitResponse)

  def rollback(self, request):
    """Rollback a transaction.

    Args:
      request: RollbackRequest proto message.

    Returns:
      RollbackResponse proto message.

    Raises:
      RPCError: The underlying RPC call failed with an HTTP error.
          (See: .response attribute)
    """
    return self._call_method('rollback', request,
                             datastore_pb2.RollbackResponse)

  def allocate_ids(self, request):
    """Allocate ids for incomplete keys.

    Args:
      request: AllocateIdsRequest proto message.

    Returns:
      AllocateIdsResponse proto message.

    Raises:
      RPCError: The underlying RPC call failed with an HTTP error.
          (See: .response attribute)
    """
    return self._call_method('allocateIds', request,
                             datastore_pb2.AllocateIdsResponse)

  def _call_method(self, method, req, resp_class):
    """_call_method call the given RPC method over HTTP.

    It uses the given protobuf message request as the payload and
    returns the deserialized protobuf message response.

    Args:
      method: RPC method name to be called.
      req: protobuf message for the RPC request.
      resp_class: protobuf message class for the RPC response.

    Returns:
      Deserialized resp_class protobuf message instance.

    Raises:
      BadArgumentError: No project has been defined.
      RPCError: The rpc method call failed.
    """
    payload = req.SerializeToString()
    headers = {
        'Content-Type': 'application/x-protobuf',
        'Content-Length': str(len(payload)),
        'X-Goog-Api-Format-Version': '2'