# Copyright 2017 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.
"""Common methods to display parts of SQL query results."""

from __future__ import absolute_import
from __future__ import unicode_literals
from functools import partial
from apitools.base.py import encoding
from googlecloudsdk.core.resource import resource_printer
from googlecloudsdk.core.util import text


def _GetAdditionalProperty(properties, property_key, not_found_value='Unknown'):
  """Gets the value for the given key in a list of properties.

  Looks through a list of properties and tries to find the value for the given
  key. If it's not found, not_found_value is returned.

  Args:
    properties: A dictionary of key string, value string pairs.
    property_key: The key string for which we want to get the value.
    not_found_value: The string value to return if the key is not found.

  Returns:
    A string containing the value for the given key, or `not_found_value` if
    the key is not found.
  """
  for prop in properties:
    if prop.key == property_key:
      if hasattr(prop, 'value'):
        return prop.value
      break
  return not_found_value


def _ConvertToTree(plan_nodes):
  """Creates tree of Node objects from the plan_nodes in server response.

  Args:
    plan_nodes (spanner_v1_messages.PlanNode[]): The plan_nodes from the server
      response. Plan nodes are topologically sorted.

  Returns:
    A Node, root of a tree built from `plan_nodes`.
  """
  # A dictionary from the index in plan_nodes to Node object.
  node_dict = {}
  # If we start at the end then by the time we get to a node's parent, that
  # node must already exist in node_dict.
  for node in reversed(plan_nodes):
    n = Node(node)
    if node.childLinks:
      for link in node.childLinks:
        n.children.append(node_dict[link.childIndex])
    if node.index:
      node_dict[node.index] = n
    # The only node without an index is the root.
    else:
      return n


def _ConvertToStringValue(prop):
  """Converts the prop to a string if it exists.

  Args:
    prop (object_value): The value returned from _GetAdditionalProperty.

  Returns:
    A string value for the given prop, or the `not_found_value` if the prop does
    not exist.
  """
  return getattr(prop, 'string_value', prop)


def QueryHasAggregateStats(result):
  """Checks if the given results have aggregate statistics.

  Args:
    result (spanner_v1_messages.ResultSetStats): The stats for a query.

  Returns:
    A boolean indicating whether 'results' contain aggregate statistics.
  """
  return hasattr(
      result, 'stats') and getattr(result.stats, 'queryStats', None) is not None


def DisplayQueryAggregateStats(query_stats, out):
  """Displays the aggregate stats for a Spanner SQL query.

  Looks at the queryStats portion of the query response and prints some of
  the aggregate statistics.

  Args:
    query_stats (spanner_v1_messages.ResultSetStats.QueryStatsValue): The query
      stats taken from the server response to a query.
    out: Output stream to which we print.
  """
  get_prop = partial(_GetAdditionalProperty, query_stats.additionalProperties)
  stats = {
      'total_elapsed_time': _ConvertToStringValue(get_prop('elapsed_time')),
      'cpu_time': _ConvertToStringValue(get_prop('cpu_time')),
      'rows_returned': _ConvertToStringValue(get_prop('rows_returned')),
      'rows_scanned': _ConvertToStringValue(get_prop('rows_scanned'))
  }
  resource_printer.Print(
      stats,
      'table[box](total_elapsed_time, cpu_time, rows_returned, rows_scanned)',
      out=out)


def DisplayQueryPlan(result, out):
  """Displays a graphical query plan for a query.

  Args:
    result (spanner_v1_messages.ResultSet): The server response to a query.
    out: Output stream to which we print.
  """
  node_tree_root = _ConvertToTree(result.stats.queryPlan.planNodes)
  node_tree_root.PrettyPrint(out)


def DisplayQueryResults(result, out):
  """Prints the result rows for a query.

  Args:
    result (spanner_v1_messages.ResultSet): The server response to a query.
    out: Output stream to which we print.
  """
  # Print "(Unspecified)" for computed columns.
  fields = [
      field.name or '(Unspecified)' for field in result.metadata.rowType.fields
  ]

  # Create the format string we pass to the table layout.
  table_format = ','.join('row.slice({0}).join():label="{1}"'.format(i, f)
                          for i, f in enumerate(fields))
  rows = [{'row': encoding.MessageToPyValue(row.entry)} for row in result.rows]

  # Can't use the PrintText method because we want special formatting.
  resource_printer.Print(rows, 'table({0})'.format(table_format), out=out)


class Node(object):
  """Represents a single node in a Spanner query plan.

  Attributes:
    properties (spanner_v1_messages.PlanNode): The details about a given node
      as returned from the server.
    children: A list of children in the query plan of type Node.
  """

  def __init__(self, properties, children=None):
    self.children = children or []
    self.properties = properties

  def _DisplayKindAndName(self, out, prepend, stub):
    """Prints the kind of the node (SCALAR or RELATIONAL) and its name."""
    kind_and_name = '{}{} {} {}'.format(prepend, stub, self.properties.kind,
                                        self.properties.displayName)
    out.Print(kind_and_name)

  def _GetNestedStatProperty(self, prop_name, nested_prop_name):
    """Gets a nested property name on this object's executionStats.

    Args:
      prop_name: A string of the key name for the outer property on
        executionStats.
      nested_prop_name: A string of the key name of the nested property.

    Returns:
      The string value of the nested property, or None if the outermost
      property or nested property don't exist.
    """
    prop = _GetAdditionalProperty(
        self.properties.executionStats.additionalProperties, prop_name, '')
    if not prop:
      return None

    nested_prop = _GetAdditionalProperty(prop.object_value.properties,
                                         nested_prop_name, '')
    if nested_prop:
      return nested_prop.string_value

    return None

  def _DisplayExecutionStats(self, out, prepend, beneath_stub):
    """Prints the relevant execution statistics for a node.

    More specifically, print out latency information and the number of
    executions. This information only exists when query is run in 'PROFILE'
    mode.

    Args:
      out: Output stream to which we print.
      prepend: String that precedes any information about this node to maintain
        a visible hierarchy.
      beneath_stub: String that preserves the indentation of the vertical lines.
    """
    if not self.properties.executionStats:
      return None

    stat_props = []

    num_executions = self._GetNestedStatProperty('execution_summary',
                                                 'num_executions')
    if num_executions:
      num_executions = int(num_executions)
      executions_str = '{} {}'.format(num_executions,
                                      text.Pluralize(num_executions,
                                                     'execution'))
      stat_props.append(executions_str)

    # Total latency and latency unit are always expected to be present when
    # latency exists. Latency exists when the query is run in PROFILE mode.
    mean_latency = self._GetNestedStatProperty('latency', 'mean')
    total_latency = self._GetNestedStatProperty('latency', 'total')
    unit = self._GetNestedStatProperty('latency', 'unit')
    if mean_latency:
      stat_props.append('{} {} average latency'.format(mean_latency, unit))
    elif total_latency:
      stat_props.append('{} {} total latency'.format(total_latency, unit))

    if stat_props:
      executions_stats_str = '{}{} ({})'.format(prepend, beneath_stub,
                                                ', '.join(stat_props))
      out.Print(executions_stats_str)

  def _DisplayMetadata(self, out, prepend, beneath_stub):
    """Prints the keys and values of the metadata for a node.

    Args:
      out: Output stream to which we print.
      prepend: String that precedes any information about this node to maintain
        a visible hierarchy.
      beneath_stub: String that preserves the indentation of the vertical lines.
    """
    if self.properties.metadata:
      additional_props = []
      # additionalProperties looks like: [key: {value: {string_value: str}}]
      for prop in self.properties.metadata.additionalProperties:
        additional_props.append(
            '{}: {}'.format(prop.key, prop.value.string_value))
      metadata = '{}{} {}'.format(prepend, beneath_stub,
                                  ', '.join(sorted(additional_props)))
      out.Print(metadata)

  def _DisplayShortRepresentation(self, out, prepend, beneath_stub):
    if self.properties.shortRepresentation:
      short_rep = '{}{} {}'.format(
          prepend, beneath_stub,
          self.properties.shortRepresentation.description)
      out.Print(short_rep)

  def _DisplayBreakLine(self, out, prepend, beneath_stub, is_root):
    """Displays an empty line between nodes for visual breathing room.

    Keeps in tact the vertical lines connecting all immediate children of a
    node to each other.

    Args:
      out: Output stream to which we print.
      prepend: String that precedes any information about this node to maintain
        a visible hierarchy.
      beneath_stub: String that preserves the indentation of the vertical lines.
      is_root: Boolean indicating whether this node is the root of the tree.
    """
    above_child = '  ' if is_root else ''
    above_child += '  |' if self.children else ''
    break_line = '{}{}{}'.format(prepend, beneath_stub, above_child)
    # It could be the case the beneath_stub adds spaces but above_child doesn't
    # add an additional vertical line, in which case we want to remove the
    # extra trailing spaces.
    out.Print(break_line.rstrip())

  def PrettyPrint(self, out, prepend=None, is_last=True, is_root=True):
    """Prints a string representation of this node in the tree.

    Args:
      out: Output stream to which we print.
      prepend: String that precedes any information about this node to maintain
        a visible hierarchy.
      is_last: Boolean indicating whether this node is the last child of its
        parent.
      is_root: Boolean indicating whether this node is the root of the tree.
    """
    prepend = prepend or ''
    # The symbol immediately before node kind to indicate that this is a child
    # of its parents. All nodes except the root get one.
    stub = '' if is_root else (r'\-' if is_last else '+-')

    # To list additional properties beneath the name, figure out how they should
    # be indented relative to the name's stub.
    beneath_stub = '' if is_root else ('  ' if is_last else '| ')

    self._DisplayKindAndName(out, prepend, stub)
    self._DisplayExecutionStats(out, prepend, beneath_stub)
    self._DisplayMetadata(out, prepend, beneath_stub)
    self._DisplayShortRepresentation(out, prepend, beneath_stub)
    self._DisplayBreakLine(out, prepend, beneath_stub, is_root)

    for idx, child in enumerate(self.children):
      is_last_child = idx == len(self.children) - 1
      # The amount each subsequent level in the tree is indented.
      indent = '   '
      # Connect all immediate children to each other with a vertical line
      # of '|'. Don't extend this line down past the last child node. It's
      # cleaner.
      child_prepend = prepend + (' ' if is_last else '|') + indent
      child.PrettyPrint(
          out, prepend=child_prepend, is_last=is_last_child, is_root=False)
