# Copyright 2015 Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can
# be found in the LICENSE file.
"""Utility module for proto3-python conversions.

This module defines the conversion functions from proto3 to python,
and utility methods / classes to convert requests / responses for any
python connector using the proto3 requests / responses.
"""

import datetime
from decimal import Decimal

from vtproto import query_pb2
from vtproto import topodata_pb2
from vtproto import vtgate_pb2
from vtproto import vtrpc_pb2

from vtdb import field_types
from vtdb import keyrange_constants
from vtdb import keyspace
from vtdb import times
from vtdb import vtgate_utils

# conversions is a map of type to the conversion function that needs
# to be used to convert the incoming array of bytes to the
# corresponding native python type.
# If a type doesn't need conversion, it's not in the map.
conversions = {
    query_pb2.INT8: int,
    query_pb2.UINT8: int,
    query_pb2.INT16: int,
    query_pb2.UINT16: int,
    query_pb2.INT24: int,
    query_pb2.UINT24: int,
    query_pb2.INT32: int,
    query_pb2.UINT32: int,
    query_pb2.INT64: int,
    query_pb2.UINT64: long,
    query_pb2.FLOAT32: float,
    query_pb2.FLOAT64: float,
    query_pb2.TIMESTAMP: times.DateTimeOrNone,
    query_pb2.DATE: times.DateOrNone,
    query_pb2.TIME: times.TimeDeltaOrNone,
    query_pb2.DATETIME: times.DateTimeOrNone,
    query_pb2.YEAR: int,
    query_pb2.DECIMAL: Decimal,
    # query_pb2.TEXT: no conversion
    # query_pb2.BLOB: no conversion
    # query_pb2.VARCHAR: no conversion
    # query_pb2.VARBINARY: no conversion
    # query_pb2.CHAR: no conversion
    # query_pb2.BINARY: no conversion
    # query_pb2.BIT: no conversion
    # query_pb2.ENUM: no conversion
    # query_pb2.SET: no conversion
    # query_pb2.TUPLE: no conversion
}

# legacy_code_to_code_map maps legacy error codes
# to the new code that matches grpc's cannonical error codes.
legacy_code_to_code_map = {
    vtrpc_pb2.SUCCESS_LEGACY: vtrpc_pb2.OK,
    vtrpc_pb2.CANCELLED_LEGACY: vtrpc_pb2.CANCELED,
    vtrpc_pb2.UNKNOWN_ERROR_LEGACY: vtrpc_pb2.UNKNOWN,
    vtrpc_pb2.BAD_INPUT_LEGACY: vtrpc_pb2.INVALID_ARGUMENT,
    vtrpc_pb2.DEADLINE_EXCEEDED_LEGACY: vtrpc_pb2.DEADLINE_EXCEEDED,
    vtrpc_pb2.INTEGRITY_ERROR_LEGACY: vtrpc_pb2.ALREADY_EXISTS,
    vtrpc_pb2.PERMISSION_DENIED_LEGACY: vtrpc_pb2.PERMISSION_DENIED,
    vtrpc_pb2.RESOURCE_EXHAUSTED_LEGACY: vtrpc_pb2.RESOURCE_EXHAUSTED,
    vtrpc_pb2.QUERY_NOT_SERVED_LEGACY: vtrpc_pb2.FAILED_PRECONDITION,
    vtrpc_pb2.NOT_IN_TX_LEGACY: vtrpc_pb2.ABORTED,
    vtrpc_pb2.INTERNAL_ERROR_LEGACY: vtrpc_pb2.INTERNAL,
    vtrpc_pb2.TRANSIENT_ERROR_LEGACY: vtrpc_pb2.UNAVAILABLE,
    vtrpc_pb2.UNAUTHENTICATED_LEGACY: vtrpc_pb2.UNAUTHENTICATED,
}


INT_UPPERBOUND_PLUS_ONE = 1<<63


def make_row(row, convs):
  """Builds a python native row from proto3 row, and conversion array.

  Args:
    row: proto3 query.Row object
    convs: conversion function array

  Returns:
    an array of converted rows.
  """
  converted_row = []
  offset = 0
  for i, l in enumerate(row.lengths):
    if l == -1:
      converted_row.append(None)
    elif convs[i]:
      converted_row.append(convs[i](row.values[offset:offset+l]))
      offset += l
    else:
      converted_row.append(row.values[offset:offset+l])
      offset += l
  return converted_row


def build_value(v):
  """Build a proto value from any valid input."""
  val = query_pb2.Value()
  convert_value(v, val)
  return val


def convert_value(value, proto_value, allow_lists=False):
  """Convert a variable from python type to proto type+value.

  Args:
    value: the python value.
    proto_value: the proto3 object, needs a type and value field.
    allow_lists: allows the use of python lists.
  """
  if isinstance(value, bool):
    proto_value.type = query_pb2.INT64
    proto_value.value = str(int(value))
  elif isinstance(value, int):
    proto_value.type = query_pb2.INT64
    proto_value.value = str(value)
  elif isinstance(value, long):
    if value < INT_UPPERBOUND_PLUS_ONE:
      proto_value.type = query_pb2.INT64
    else:
      proto_value.type = query_pb2.UINT64
    proto_value.value = str(value)
  elif isinstance(value, float):
    proto_value.type = query_pb2.FLOAT64
    proto_value.value = str(value)
  elif hasattr(value, '__sql_literal__'):
    proto_value.type = query_pb2.VARBINARY
    proto_value.value = str(value.__sql_literal__())
  elif isinstance(value, datetime.datetime):
    proto_value.type = query_pb2.VARBINARY
    proto_value.value = times.DateTimeToString(value)
  elif isinstance(value, datetime.date):
    proto_value.type = query_pb2.VARBINARY
    proto_value.value = times.DateToString(value)
  elif isinstance(value, str):
    proto_value.type = query_pb2.VARBINARY
    proto_value.value = value
  elif isinstance(value, field_types.NoneType):
    proto_value.type = query_pb2.NULL_TYPE
  elif allow_lists and isinstance(value, (set, tuple, list)):
    # this only works for bind variables, not for entities.
    proto_value.type = query_pb2.TUPLE
    for v in list(value):
      proto_v = proto_value.values.add()
      convert_value(v, proto_v)
  else:
    proto_value.type = query_pb2.VARBINARY
    proto_value.value = str(value)


def convert_bind_vars(bind_variables, request_bind_variables):
  """Convert binding variables to proto3.

  Args:
    bind_variables: a map of strings to python native types.
    request_bind_variables: the proto3 object to add bind variables to.
  """
  if not bind_variables:
    return
  for key, val in bind_variables.iteritems():
    convert_value(val, request_bind_variables[key], allow_lists=True)


def convert_stream_event_statement(statement):
  """Converts encoded rows inside a StreamEvent.Statement to native types.

  Args:
    statement: the StreamEvent.Statement object.

  Returns:
    fields: array of names for the primary key columns.
    rows: array of tuples for each primary key value.
  """
  fields = []
  rows = []
  if statement.primary_key_fields:
    convs = []
    for field in statement.primary_key_fields:
      fields.append(field.name)
      convs.append(conversions.get(field.type))

    for r in statement.primary_key_values:
      row = tuple(make_row(r, convs))
      rows.append(row)

  return fields, rows


class Proto3Connection(object):
  """A base class for proto3-based python connectors.

  It assumes the derived object will contain a proto3 self.session object.
  """

  def __init__(self):
    self._effective_caller_id = None
    self.event_token = None
    self.fresher = None

  def _add_caller_id(self, request, caller_id):
    """Adds the vtgate_client.CallerID to the proto3 request, if any.

    Args:
      request: proto3 request (any of the {,stream,batch} execute queries).
      caller_id: vtgate_client.CallerID object.
    """
    if caller_id:
      if caller_id.principal:
        request.caller_id.principal = caller_id.principal
      if caller_id.component:
        request.caller_id.component = caller_id.component
      if caller_id.subcomponent:
        request.caller_id.subcomponent = caller_id.subcomponent

  def _add_session(self, request):
    """Adds self.session to the request, if any.

    Args:
      request: the proto3 request to add session to.
    """
    if self.session:
      request.session.CopyFrom(self.session)

  def update_session(self, response):
    """Updates the current session from the response, if it has one.

    Args:
      response: a proto3 response that may contain a session object.
    """
    if response.HasField('session') and response.session:
      self.session = response.session

  def _convert_entity_ids(self, entity_keyspace_ids, request_eki):
    """Convert external entity id map to ProtoBuffer.

    Args:
      entity_keyspace_ids: map of entity_keyspace_id.
      request_eki: destination proto3 list.
    """
    for xid, kid in entity_keyspace_ids.iteritems():
      eid = request_eki.add()
      eid.keyspace_id = kid
      convert_value(xid, eid, allow_lists=False)

  def _add_key_ranges(self, request, key_ranges):
    """Adds the provided keyrange.KeyRange objects to the proto3 request.

    Args:
      request: proto3 request.
      key_ranges: list of keyrange.KeyRange objects.
    """
    for kr in key_ranges:
      encoded_kr = request.key_ranges.add()
      encoded_kr.start = kr.Start
      encoded_kr.end = kr.End

  def _extract_rpc_error(self, exec_method, error):
    """Raises a VitessError for a proto3 vtrpc.RPCError structure, if set.

    Args:
      exec_method: name of the method to use in VitessError.
      error: vtrpc.RPCError structure.

    Raises:
      vtgate_utils.VitessError: if an error was set.
    """
    if error.code:
      raise vtgate_utils.VitessError(exec_method, error.code, error.message)
    elif error.legacy_code:
      raise vtgate_utils.VitessError(
          exec_method,
          legacy_code_to_code_map[error.legacy_code],
          error.message)