# Copyright 2019 Verily Life Sciences 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.
"""Base class for a library for interacting with BigQuery."""

# Workaround for https://github.com/GoogleCloudPlatform/google-cloud-python/issues/2366
from __future__ import absolute_import

import logging

from google.api_core.future import polling
from google.api_core.retry import Retry
from google.cloud.bigquery import retry as bq_retry

# Bigquery uses . to separate project, dataset, and table parts.
BQ_PATH_DELIMITER = '.'

# The call to datasets to list all tables requires you to set a maximum number of tables
# (or use the unspecified API default). We always want to list all the tables so we set
# a really high number for it. If ever there would be a dataset with more than 5000 tables,
# this constant should be adjusted. Also applies to max datasets to list in a project.
MAX_TABLES = 5000

# Constants for classes of datatypes. Some Bigquery functions accept or return legacy SQL types
# while others accept or return standard SQL types. These sets map equivalences between them.
INTEGER_CLASS = set(['INTEGER', 'INT64'])
FLOAT_CLASS = set(['FLOAT', 'FLOAT64'])
TIME_CLASS = set(['TIMESTAMP', 'DATE', 'TIME', 'DATETIME'])
BOOLEAN_CLASS = set(['BOOLEAN', 'BOOL'])
DATATYPE_CLASSES = [INTEGER_CLASS, FLOAT_CLASS, TIME_CLASS, BOOLEAN_CLASS]

# Values in the QueryJob.error_result dict that correspond to query validation errors.
VALIDATION_ERROR_REASONS = ['invalidQuery', 'notFound', 'duplicate']

# The default timeout for any BigQuery operations executed in real BQ.
DEFAULT_TIMEOUT_SEC = 1200


def _transient_string_in_exception_message(exc):
    # type: (Exception) -> bool
    """Determines whether an exception's message contains a common message for transient errors.

    The exception's message containing one of these substrings is sufficient to determine that it is
    transient, but there can be transient exceptions whose messages do not contain these substrings.
    """
    return ('The job encountered an internal error during execution' in str(exc) or
            'Retrying the job may solve the problem' in str(exc))

# Retry object for errors encountered in making API calls (executing jobs, etc.)
DEFAULT_RETRY_FOR_API_CALLS = Retry(
    # The predicate takes an exception and returns whether it is transient.
    predicate=lambda exc: (bq_retry.DEFAULT_RETRY._predicate(exc) or
                           _transient_string_in_exception_message(exc)),
    deadline=DEFAULT_TIMEOUT_SEC)

# Retry object for errors encountered while polling jobs in progress.
# See https://github.com/googleapis/google-cloud-python/issues/6301
DEFAULT_RETRY_FOR_ASYNC_JOBS = Retry(
            # The predicate takes an exception and returns whether it is transient.
            predicate=lambda exc: (polling.DEFAULT_RETRY._predicate(exc) or
                                   _transient_string_in_exception_message(exc)),
    deadline=DEFAULT_TIMEOUT_SEC)


class BigqueryBaseClient(object):
    """Stores credentials and pointers to a BigQuery project.

    Args:
      project_id: The id of the project to associate with the client.
      default_dataset_id: If specified, use this dataset as the default.
      maximum_billing_tier: The maximum billing tier of this client.
      default_max_api_call_tries: The maximum number of tries for any REST API call.
    """

    def __init__(self, project_id, default_dataset_id=None, maximum_billing_tier=None):
        self.maximum_billing_tier = maximum_billing_tier
        self.project_id = project_id
        self.default_dataset_id = default_dataset_id

    def get_delimiter(self):
        """ Returns the delimiter used to separate project, dataset, and table in a table path. """
        raise NotImplementedError("get_delimiter is not implemented.")

    def get_query_results(self, query, use_legacy_sql=False, max_wait_secs=None):
        # type: (str, Optional[Bool], Optional[int]) -> List[Tuple[Any]]
        """Returns a list or rows, each of which is a tuple of values.

        Args:
            query: A string with a complete SQL query.
            use_legacy_sql: Whether to use legacy SQL
            max_wait_secs: The maximum number of seconds to wait for the query to complete. If not
                set, the class default will be used.

        Returns:
            A list of tuples of values.
        """
        raise NotImplementedError("get_query_results is not implemented.")

    def create_table_from_query(self,
                                query,  # type: str
                                table_path,  # type: str
                                write_disposition='WRITE_EMPTY',  # type: Optional[str]
                                use_legacy_sql=False,  # type: Optional[bool]
                                max_wait_secs=None,  # type: Optional[int]
                                expected_schema=None  # type: Optional[List[SchemaField]]
                                ):
        # type: (...) -> None
        """Creates a table in BigQuery from a specified query.

        Args:
          query: The query to run.
          table_path: The path to the table (in the client's project) to write
              the results to.
          write_disposition: Specifies behavior if table already exists. See options here:
              https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs under
              configuration.query.writeDisposition
          use_legacy_sql: Whether the query is written in standard or legacy sql.
          max_wait_secs: Seconds to wait for the query before timing out. If not
                set, the class default will be used.
          expected_schema: The expected schema of the resulting table (only required for mocking).
        """
        raise NotImplementedError("create_table_from_query is not implemented.")

    def create_tables_from_dict(self,
                                table_names_to_schemas,  # type: Dict[str, List[SchemaField]]
                                dataset_id=None,  # type: Optional[str]
                                replace_existing_tables=False  # type: Optional[bool]
                                ):
        # type: (...) -> None
        """Creates a set of tables from a dictionary of table names to their schemas.

        Args:
          table_names_to_schemas: A dictionary of:
            key: The table name.
            value: A list of SchemaField objects.
          dataset_id: The dataset in which to create tables. If not specified, use default dataset.
          replace_existing_tables: If True, delete and re-create tables. Otherwise, checks to see
              if any of the requested tables exist. If they do, it will raise a RuntimeError.

        Raises:
            RuntimeError if replace_existing_tables is False and any of the tables requested for
                creation already exist
        """
        raise NotImplementedError("create_tables_from_dict is not implemented.")

    def create_dataset_by_name(self, name, expiration_hours=None):
        # type: (str, Optional[float]) -> None
        """Create a new dataset within the current project.

        Args:
          name: The name of the new dataset.
          expiration_hours: The default expiration time for tables within the dataset.
        """
        raise NotImplementedError("create_dataset_by_name is not implemented.")

    def delete_dataset_by_name(self, name, delete_all_tables=False):
        # type: (str, Optional[bool]) -> None
        """Delete a dataset within the current project.

        Args:
          name: The name of the dataset to delete.
          delete_all_tables: If True, will delete all tables in the dataset before attempting to
              delete the dataset. You can't delete a dataset until it contains no tables.

        Raises:
            RuntimeError if there are still tables in the dataset and you try to delete it (with
                delete_all_tables set to False)
        """
        raise NotImplementedError("delete_dataset_by_name is not implemented.")

    def delete_table_by_name(self, table_path):
        # type: (str) -> None
        """Delete a table within the current project.

        Args:
          table_path: A string of the form '<dataset id>.<table name>' or
              '<project id>.<dataset_id>.<table_name>'
        """
        raise NotImplementedError("delete_table_by_name is not implemented.")

    def dataset_exists_with_name(self, dataset_name):
        # type: (str) -> bool
        """Determines whether a dataset exists with the given name.

        Args:
            dataset_name: The name of the dataset to check.

        Returns:
            True if the dataset exists in this client's project, False otherwise.
        """
        raise NotImplementedError("dataset_exists_with_name is not implemented.")

    def table_exists_with_name(self, table_path):
        # type: (str) -> bool
        """Determines whether a table exists at the given table path.

        Args:
            table_path: The table path of the table to check. Uses the default dataset ID if a
                dataset is not specified as part of the table path.

        Returns:
            True if the table exists at the given path, False otherwise.
        """
        raise NotImplementedError("table_exists_with_name is not implemented.")

    def tables(self, dataset_id):
        # type: (str) -> List[str]
        """Returns a list of table names in a given dataset.

        Args:
          dataset_id: The name of the dataset to query.

        Returns:
          A list of table names (strings).
        """
        raise NotImplementedError("tables is not implemented.")

    def get_schema(self, dataset_id, table_name, project_id=None):
        # type: (str, str, Optional[str]) -> List[SchemaField]
        """Returns the schema of a table.

        Args:
          dataset_id: The dataset to query.
          table_name: The name of the table.
          project_id: The project ID of the table.
        Returns:
          A list of SchemaFields representing the schema.
        """
        raise NotImplementedError("get_schema is not implemented.")

    def get_datasets(self):
        # type: (None) -> List[str]
        """Returns a list of dataset ids in the default project.

        Returns:
            A list of dataset ids/names (strings).
        """
        raise NotImplementedError("get_datasets is not implemented.")

    def populate_table(self, table_path, schema, data=[], make_immediately_available=False,
                       replace_existing_table=False):
        # type: (str, List[SchemaField], Optional[List[Any]], Optional[bool], Optional[bool]) -> None
        """Creates a table and populates it with a list of rows.

        If make_immediately_available is False, the table will be created using streaming inserts.
        Note that streaming inserts are immediately available for querying, but not for exporting or
        copying, so if you need that capability you should set make_immediately_available to True.
        https://cloud.google.com/bigquery/streaming-data-into-bigquery

        If the table is already created, it will raise a RuntimeError, unless replace_existing_table
        is True.

        Args:
          table_path: A string of the form '<dataset id>.<table name>'
              or '<project id>.<dataset id>.<table name>'.
          schema: A list of SchemaFields to represent the table's schema.
          data: A list of rows, each of which corresponds to a row to insert into the table.
          make_immediately_available: If False, the table won't immediately be available for
              copying or exporting, but will be available for querying. If True, after this
              operation returns, it will be available for copying and exporting too.
          replace_existing_table: If set to True, the table at table_path will be deleted and
              recreated if it's already present.

        Raises:
            RuntimeError if the table at table_path is already there and replace_existing_table
                is False
        """
        raise NotImplementedError("populate_table is not implemented.")

    #TODO(Issue 9): Implement a make_immediately_available flag for this function.
    def append_rows(self, table_path, data, schema=None):
        # type: (str, List[Tuple[Any]], Optional[List[SchemaField]]) -> None
        """Appends the rows contained in data to the table at table_path.
        Args:
          table_path: A string of the form '<dataset id>.<table name>'
              or '<project id>.<dataset id>.<table name>'.
          data: A list of rows, each of which is a list of values.
          schema: Optionally, a list of SchemaFields to describe the
              table's expected schema. If this is present, it will check the table's schema against