# Copyright 2021 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Module for querying BigQuery."""

import collections
import json
import os
import subprocess
from typing import Any, Dict, List

from flake_suppressor_common import common_typing as ct
from flake_suppressor_common import results as results_module
from flake_suppressor_common import tag_utils

MAX_ROWS = (2**31) - 1

# A note about the try version of the queries: The submitted builds subquery is
# included in each query instead of running it once by itself and including the
# returned data in other queries because we can end up getting a very large
# number of build IDs, which can push the query over BigQuery's hard query size
# limit. The query runs pretty quickly (within a couple of seconds), so
# duplicating it does not add much runtime.

# Subquery for getting all builds used for CL submission in the past
# |sample_period| days. Will be inserted into other queries.
SUBMITTED_BUILDS_TEMPLATE = """\
  SELECT
    CONCAT("build-", CAST(unnested_builds.id AS STRING)) as id
  FROM
    `commit-queue.chromium.attempts`,
    UNNEST(builds) as unnested_builds,
    UNNEST(gerrit_changes) as unnested_changes
  WHERE
    unnested_builds.host = "cr-buildbucket.appspot.com"
    AND unnested_changes.submit_status = "SUCCESS"
    AND start_time > TIMESTAMP_SUB(CURRENT_TIMESTAMP(),
                                    INTERVAL @sample_period DAY)
"""

# Subquery for getting all ci builds that are in sheriff rotations in the past
# |sample_period| days. Will be inserted into other queries.
SHERIFF_ROTATIONS_CI_BUILDS_TEMPLATE = """\
  SELECT DISTINCT builder.builder,
  FROM
    `cr-buildbucket.chrome.builds_30d`
  WHERE
    input.properties LIKE '%sheriff_rotations%'
    AND JSON_VALUE_ARRAY(input.properties, '$.sheriff_rotations')[OFFSET(0)]
      IN ("chromium", "android")
    AND start_time > TIMESTAMP_SUB(CURRENT_TIMESTAMP(),
                                    INTERVAL @sample_period DAY)
"""

QueryParameters = Dict[str, Dict[str, Any]]


class BigQueryQuerier():
  def __init__(self, sample_period: int, billing_project: str,
               result_processor: results_module.ResultProcessor):
    """Class for making calls to BigQuery.

    Args:
      sample_period: An int denoting the number of days that data should be
          queried over.
      billing_project: A string containing the billing project to use for
          BigQuery queries.
    """
    self._sample_period = sample_period
    self._billing_project = billing_project
    self._result_processor = result_processor

  def GetFlakyOrFailingCiTests(self) -> ct.QueryJsonType:
    """Gets all flaky or failing tests from CI.

    Returns:
      A JSON representation of the BigQuery results containing all found flaky
      or failing test results that came from CI bots.
    """
    return self._GetJsonResultsFromBigQuery(self.GetFlakyOrFailingCiQuery())

  def GetFailingCiBuildCulpritTests(self) -> ct.QueryJsonType:
    """Gets all failing build culprit tests from CI builders.

    Returns:
      A JSON representation of the BigQuery results containing all found
      all failing build culprit results that came from CI bots.
    """
    return self._GetJsonResultsFromBigQuery(
        self.GetFailingBuildCulpritFromCiQuery())

  def GetFlakyOrFailingTestsFromCiBuilders(
      self, builder_names: List[str]) -> ct.QueryJsonType:
    """Gets all flaky or failing tests from input CI builders.

    Returns:
      A JSON representation of the BigQuery results containing all found
      all failing results that came from input CI builders.
    """
    return self._GetJsonResultsFromBigQuery(
        self.GetFlakyOrFailingFromCIBuildersQuery(builder_names))

  def GetFailingBuildCulpritFromCiBuilders(
      self, builder_names: List[str]) -> ct.QueryJsonType:
    """Gets all failing build culprit tests from input CI builders.

    Returns:
      A JSON representation of the BigQuery results containing all found
      all failing results that came from input CI builders.
    """
    return self._GetJsonResultsFromBigQuery(
        self.GetFailingBuildCulpritFromCIBuildersQuery(builder_names))

  def GetFlakyOrFailingTryTests(self) -> ct.QueryJsonType:
    """Gets all flaky or failing tests from the trybots.

    Limits results to those that came from builds used for CL submission.

    Returns:
      A JSON representation of the BigQuery results containing all found flaky
      or failing test results that came from trybots AND came from builds that
      were used for CL submission.
    """
    return self._GetJsonResultsFromBigQuery(self.GetFlakyOrFailingTryQuery())

  def GetResultCounts(self) -> ct.ResultCountType:
    """Gets the result count for each test/config combination.

    Returns:
      A dict in the format:
      {
        typ_tags (tuple): {
          test_name (str): result_count (int)
        }
      }
    """
    # A default dict of default dicts of ints.
    result_counts = collections.defaultdict(lambda: collections.defaultdict(int)
                                            )
    self._GetResultCountWithQuery(self.GetResultCountCIQuery(), result_counts)
    self._GetResultCountWithQuery(self.GetResultCountTryQuery(), result_counts)
    return result_counts

  def GetResultCountFromCiBuilders(
      self, builder_names: List[str]) -> ct.ResultCountType:
    """Gets the result count for the input CI builders.

    Returns:
      A dict in the format:
      {
        typ_tags (tuple): {
          test_name (str): result_count (int)
        }
      }
    """
    result_counts = collections.defaultdict(
        lambda: collections.defaultdict(int))
    self._GetResultCountWithQuery(
        self.GetResultCountFromCIBuildersQuery(builder_names), result_counts)
    return result_counts

  def GetFlakyOrFailingCiQuery(self) -> str:
    """
    Returns:
      Query string to get all the failing or flaky results from CI bots.
    """
    raise NotImplementedError

  def GetFailingBuildCulpritFromCiQuery(self) -> str:
    """
    Returns:
      Query string to get all failing build culprit results from CI bots.
    """
    raise NotImplementedError

  def GetFlakyOrFailingFromCIBuildersQuery(self,
                                           builder_names: List[str]) -> str:
    """
    Returns:
      Query string to get all the failing or flaky results from input CI
      builders.
    """
    raise NotImplementedError

  def GetFailingBuildCulpritFromCIBuildersQuery(
      self, builder_names: List[str]) -> str:
    """
    Returns:
      Query string to get all failing build culprit results from input CI
      builders.
    """
    raise NotImplementedError

  def GetFlakyOrFailingTryQuery(self) -> str:
    """
    Returns:
      Query string to get all the failing or flaky results from Try bots.
    """
    raise NotImplementedError

  def GetResultCountCIQuery(self) -> str:
    """
    Returns:
      Query string to get the result count for test/tag combination from CI
      bots.
    """
    raise NotImplementedError

  def GetResultCountTryQuery(self) -> str:
    """
    Returns:
      Query string to get result count for test/tag combination from Try
      bots.
    """
    raise NotImplementedError

  def GetResultCountFromCIBuildersQuery(self, builder_names: List[str]) -> str:
    """
    Returns:
      Query string to get the result count for test/tag combination from input
      CI builders.
    """
    raise NotImplementedError

  def _GetJsonResultsFromBigQuery(self, query: str) -> ct.QueryJsonType:
    """Gets the JSON results from a BigQuery query.

    Automatically passes in the "@sample_period" parameterized argument to
    BigQuery.

    Args:
      query: A string containing the SQL query to run in BigQuery.

    Returns:
      The loaded JSON results from running |query|.
    """
    cmd = GenerateBigQueryCommand(
        self._billing_project,
        {'INT64': {
            'sample_period': self._sample_period
        }},
        batch=False)

    with open(os.devnull, 'w') as devnull:
      completed_process = subprocess.run(cmd,
                                         input=query,
                                         stdout=subprocess.PIPE,
                                         stderr=devnull,
                                         check=True,
                                         text=True)

    return json.loads(completed_process.stdout)

  def _GetResultCountWithQuery(self, query: str,
                               result_counts: ct.ResultCountType) -> None:
    """Helper to get result counts using a particular query.

    Args:
      query: A string containing a SQL query to run.
      result_counts: A defaultdict of defaultdict of ints that will be modified
          in place to tally result counts.
    """
    json_results = self._GetJsonResultsFromBigQuery(query)

    for r in json_results:
      typ_tags = tuple(tag_utils.TagUtils.RemoveIgnoredTags(r['typ_tags']))
      test_name = r['test_name']
      _, test_name = self._result_processor.GetTestSuiteAndNameFromResultDbName(
          test_name)
      count = int(r['result_count'])
      result_counts[typ_tags][test_name] += count


# TODO(crbug.com/343248818): Switch off this and use the bigquery module
# directly.
def GenerateBigQueryCommand(project: str,
                            parameters: QueryParameters,
                            batch: bool = True) -> List[str]:
  """Generate a BigQuery commandline.

  Does not contain the actual query, as that is passed in via stdin.

  Args:
    project: A string containing the billing project to use for BigQuery.
    parameters: A dict specifying parameters to substitute in the query in
        the format {type: {key: value}}. For example, the dict:
        {'INT64': {'num_builds': 5}}
        would result in --parameter=num_builds:INT64:5 being passed to BigQuery.
    batch: Whether to run the query in batch mode or not. Batching adds some
        random amount of overhead since it means the query has to wait for idle
        resources, but also allows for much better parallelism.

  Returns:
    A list containing the BigQuery commandline, suitable to be passed to a
    method from the subprocess module.
  """
  cmd = [
      'bq',
      'query',
      '--max_rows=%d' % MAX_ROWS,
      '--format=json',
      '--project_id=%s' % project,
      '--use_legacy_sql=false',
  ]

  if batch:
    cmd.append('--batch')

  for parameter_type, parameter_pairs in parameters.items():
    for k, v in parameter_pairs.items():
      cmd.append('--parameter=%s:%s:%s' % (k, parameter_type, v))
  return cmd