#!/usr/bin/env impala-python
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

# This module implements helpers for representing queries to be executed by the
# stress test, loading them and generating them.

import logging
import os
from textwrap import dedent

from tests.comparison.db_types import Int, TinyInt, SmallInt, BigInt
from tests.comparison.query_generator import QueryGenerator
from tests.comparison.query_profile import DefaultProfile
from tests.comparison.model_translator import SqlWriter
from tests.util.parse_util import match_memory_estimate, parse_mem_to_mb
import tests.util.test_file_parser as test_file_parser

LOG = logging.getLogger(os.path.splitext(os.path.basename(__file__))[0])


class QueryType(object):
  COMPUTE_STATS, DELETE, INSERT, SELECT, UPDATE, UPSERT = range(6)


class Query(object):
  """Contains a SQL statement along with expected runtime information.
  This class is used as a struct, with the fields filled out by calling
  classes."""

  def __init__(self):
    self.name = None
    self.sql = None
    # In order to be able to make good estimates for DML queries in the binary search,
    # we need to bring the table to a good initial state before excuting the sql. Running
    # set_up_sql accomplishes this task.
    self.set_up_sql = None
    self.db_name = None
    self.result_hash = None
    self.required_mem_mb_with_spilling = None
    self.required_mem_mb_without_spilling = None
    self.solo_runtime_profile_with_spilling = None
    self.solo_runtime_profile_without_spilling = None
    self.solo_runtime_secs_with_spilling = None
    self.solo_runtime_secs_without_spilling = None
    # Query options to set before running the query.
    self.options = {}
    # Determines the order in which we will populate query runtime info. Queries with the
    # lowest population_order property will be handled first.
    self.population_order = 0
    # Type of query. Can have the following values: SELECT, COMPUTE_STATS, INSERT, UPDATE,
    # UPSERT, DELETE.
    self.query_type = QueryType.SELECT

    self._logical_query_id = None

  def __repr__(self):
    return dedent("""
        <Query
        Mem: %(required_mem_mb_with_spilling)s
        Mem no-spilling: %(required_mem_mb_without_spilling)s
        Solo Runtime: %(solo_runtime_secs_with_spilling)s
        Solo Runtime no-spilling: %(solo_runtime_secs_without_spilling)s
        DB: %(db_name)s
        Options: %(options)s
        Set up SQL: %(set_up_sql)s>
        SQL: %(sql)s>
        Population order: %(population_order)r>
        """.strip() % self.__dict__)

  @property
  def logical_query_id(self):
    """
    Return a meanginful unique str identifier for the query.

    Example: "tpcds_300_decimal_parquet_q21"
    """
    if self._logical_query_id is None:
      self._logical_query_id = '{0}_{1}'.format(self.db_name, self.name)
    return self._logical_query_id

  def write_runtime_info_profiles(self, directory):
    """Write profiles for spilling and non-spilling into directory (str)."""
    profiles_to_write = [
        (self.logical_query_id + "_profile_with_spilling.txt",
         self.solo_runtime_profile_with_spilling),
        (self.logical_query_id + "_profile_without_spilling.txt",
         self.solo_runtime_profile_without_spilling),
    ]
    for filename, profile in profiles_to_write:
      if profile is None:
        LOG.debug("No profile recorded for {0}".format(filename))
        continue
      with open(os.path.join(directory, filename), "w") as fh:
        fh.write(profile)


def load_tpc_queries(workload):
  """Returns a list of TPC queries. 'workload' should either be 'tpch' or 'tpcds'."""
  LOG.info("Loading %s queries", workload)
  queries = []
  for query_name, query_sql in test_file_parser.load_tpc_queries(workload,
      include_stress_queries=True).iteritems():
    query = Query()
    query.name = query_name
    query.sql = query_sql
    queries.append(query)
  return queries


def load_queries_from_test_file(file_path, db_name=None):
  LOG.debug("Loading queries from %s", file_path)
  test_cases = test_file_parser.parse_query_test_file(file_path)
  queries = list()
  for test_case in test_cases:
    query = Query()
    query.sql = test_file_parser.remove_comments(test_case["QUERY"])
    query.db_name = db_name
    queries.append(query)
  return queries


def generate_compute_stats_queries(cursor):
  """For each table in the database that cursor is connected to, generate several compute
  stats queries. Each query will have a different value for the MT_DOP query option.
  """
  LOG.info("Generating Compute Stats queries")
  tables = [cursor.describe_table(t) for t in cursor.list_table_names()
            if not t.endswith("_original")]
  result = []
  mt_dop_values = [str(2**k) for k in range(5)]
  for table in tables:
    for mt_dop_value in mt_dop_values:
      compute_query = Query()
      compute_query.population_order = 1
      compute_query.query_type = QueryType.COMPUTE_STATS
      compute_query.sql = "COMPUTE STATS {0}".format(table.name)
      compute_query.options["MT_DOP"] = mt_dop_value
      compute_query.db_name = cursor.db_name
      compute_query.name = "compute_stats_{0}_mt_dop_{1}".format(
          table.name, compute_query.options["MT_DOP"])
      result.append(compute_query)
      LOG.debug("Added compute stats query: {0}".format(compute_query))
  return result


def generate_DML_queries(cursor, dml_mod_values):
  """Generate insert, upsert, update, delete DML statements.

  For each table in the database that cursor is connected to, create 4 DML queries
  (insert, upsert, update, delete) for each mod value in 'dml_mod_values'. This value
  controls which rows will be affected. The generated queries assume that for each table
  in the database, there exists a table with a '_original' suffix that is never modified.

  This function has some limitations:
  1. Only generates DML statements against Kudu tables, and ignores non-Kudu tables.
  2. Requires that the type of the first column of the primary key is an integer type.
  """
  LOG.info("Generating DML queries")
  tables = [cursor.describe_table(t) for t in cursor.list_table_names()
            if not t.endswith("_original")]
  result = []
  for table in tables:
    if not table.primary_keys:
      # Skip non-Kudu tables. If a table has no primary keys, then it cannot be a Kudu
      # table.
      LOG.debug("Skipping table '{0}' because it has no primary keys.".format(table.name))
      continue
    if len(table.primary_keys) > 1:
      # TODO(IMPALA-4665): Add support for tables with multiple primary keys.
      LOG.debug("Skipping table '{0}' because it has more than "
                "1 primary key column.".format(table.name))
      continue
    primary_key = table.primary_keys[0]
    if primary_key.exact_type not in (Int, TinyInt, SmallInt, BigInt):
      # We want to be able to apply the modulo operation on the primary key. If the
      # the first primary key column happens to not be an integer, we will skip
      # generating queries for this table
      LOG.debug("Skipping table '{0}' because the first column '{1}' in the "
                "primary key is not an integer.".format(table.name, primary_key.name))
      continue
    for mod_value in dml_mod_values:
      # Insert
      insert_query = Query()
      # Populate runtime info for Insert and Upsert queries before Update and Delete
      # queries because tables remain in original state after running the Insert and
      # Upsert queries. During the binary search in runtime info population for the
      # Insert query, we first delete some rows and then reinsert them, so the table
      # remains in the original state. For the delete, the order is reversed, so the table
      # is not in the original state after running the the delete (or update) query. This
      # is why population_order is smaller for Insert and Upsert queries.
      insert_query.population_order = 1
      insert_query.query_type = QueryType.INSERT
      insert_query.name = "insert_{0}".format(table.name)
      insert_query.db_name = cursor.db_name
      insert_query.sql = (
          "INSERT INTO TABLE {0} SELECT * FROM {0}_original "
          "WHERE {1} % {2} = 0").format(table.name, primary_key.name, mod_value)
      # Upsert
      upsert_query = Query()
      upsert_query.population_order = 1
      upsert_query.query_type = QueryType.UPSERT
      upsert_query.name = "upsert_{0}".format(table.name)
      upsert_query.db_name = cursor.db_name
      upsert_query.sql = (
          "UPSERT INTO TABLE {0} SELECT * "
          "FROM {0}_original WHERE {1} % {2} = 0").format(
              table.name, primary_key.name, mod_value)
      # Update
      update_query = Query()
      update_query.population_order = 2
      update_query.query_type = QueryType.UPDATE
      update_query.name = "update_{0}".format(table.name)
      update_query.db_name = cursor.db_name
      update_list = ', '.join(
          'a.{0} = b.{0}'.format(col.name)
          for col in table.cols if not col.is_primary_key)
      update_query.sql = (
          "UPDATE a SET {update_list} FROM {table_name} a JOIN {table_name}_original b "
          "ON a.{pk} = b.{pk} + 1 WHERE a.{pk} % {mod_value} = 0").format(
              table_name=table.name, pk=primary_key.name, mod_value=mod_value,
              update_list=update_list)
      # Delete
      delete_query = Query()
      delete_query.population_order = 2
      delete_query.query_type = QueryType.DELETE
      delete_query.name = "delete_{0}".format(table.name)
      delete_query.db_name = cursor.db_name
      delete_query.sql = ("DELETE FROM {0} WHERE {1} % {2} = 0").format(
          table.name, primary_key.name, mod_value)

      if table.name + "_original" in set(table.name for table in tables):
        insert_query.set_up_sql = "DELETE FROM {0} WHERE {1} % {2} = 0".format(
            table.name, primary_key.name, mod_value)
        upsert_query.set_up_sql = insert_query.set_up_sql
        update_query.set_up_sql = (
            "UPSERT INTO TABLE {0} SELECT * FROM {0}_original "
            "WHERE {1} % {2} = 0").format(table.name, primary_key.name, mod_value)
        delete_query.set_up_sql = update_query.set_up_sql

      result.append(insert_query)
      LOG.debug("Added insert query: {0}".format(insert_query))
      result.append(update_query)
      LOG.debug("Added update query: {0}".format(update_query))
      result.append(upsert_query)
      LOG.debug("Added upsert query: {0}".format(upsert_query))
      result.append(delete_query)
      LOG.debug("Added delete query: {0}".format(delete_query))
  assert len(result) > 0, "No DML queries were added."
  return result


def generate_random_queries(impala, random_db):
  """Generator function to produce random queries. 'impala' is the Impala service
  object. random_db is the name of the database that queries should be
  generated for."""
  with impala.cursor(db_name=random_db) as cursor:
    tables = [cursor.describe_table(t) for t in cursor.list_table_names()]
  query_generator = QueryGenerator(DefaultProfile())
  model_translator = SqlWriter.create()
  while True:
    query_model = query_generator.generate_statement(tables)
    sql = model_translator.write_query(query_model)
    query = Query()
    query.sql = sql
    query.db_name = random_db
    yield query


def estimate_query_mem_mb_usage(query, impalad_conn):
  """Runs an explain plan then extracts and returns the estimated memory needed to run
  the query.
  """
  with impalad_conn.cursor() as cursor:
    LOG.debug("Using %s database", query.db_name)
    if query.db_name:
      cursor.execute('USE ' + query.db_name)
    if query.query_type == QueryType.COMPUTE_STATS:
      # Running "explain" on compute stats is not supported by Impala.
      return
    LOG.debug("Explaining query\n%s", query.sql)
    cursor.execute('EXPLAIN ' + query.sql)
    explain_rows = cursor.fetchall()
    explain_lines = [row[0] for row in explain_rows]
    mem_limit, units = match_memory_estimate(explain_lines)
    return parse_mem_to_mb(mem_limit, units)