"""DB Connector. Manages interactions with Snowflake."""

from contextlib import contextmanager
import os
import re
import urllib

from cached_property import cached_property
import requests
from snowflake import connector

from ssavva.sos_snowflake_performance import config


class SnowflakeSQLExecutor(object):
    """Helper base class to abstract Snowflake operations within ETLs.

    Operates with the connection to Snowflake, provides the methods for
    retrieving information about Snowflake tables, etc.
    Could be used as is, if provided methods are enough, and also could be
    extended with more specific method to incapsulate db logic.
    """

    def __init__(self, sf_config, autocommit=True):
        """Open a connection to database, set up the db operations.

        Args:
            sf_config (dict): Dict with all the required credentials.
            autocommit (bool): If True, use autocommit mode.
        """
        self.sf_config = sf_config
        self.autocommit = autocommit
        # self.validator = validator.BaseValidator()
        self.validator = None

        self.snowflake_conn = self.get_connection()

    def __enter__(self):
        """Implement the context manager protocol.

        Returns:
            (SnowflakeSQLExecutor obj): An instance of SnowflakeSQLExecutor.
        """
        return self

    def __exit__(self, ext_type, exc_value, traceback):
        """Implement the context manager protocol.

        Calls methods to free resources.
        """
        self.snowflake_conn.close()

    def get_connection(self):
        """Create connection to the Snowflake.

        If autocommit == False, then automatically rollback all the statements
        of the current connection in case if db exception occurs within any of
        them.
        Returns:
            Connection: Connection that supports DB API v2 interface.
        """
        return connector.connect(
            user=self.sf_config['user'],
            password=self.sf_config['password'],
            account=self.sf_config['account'],
            warehouse=self.sf_config['warehouse'],
            db=self.sf_config['db'],
            schema=self.sf_config['schema'],
            role=self.sf_config['role'],
            autocommit=self.autocommit)

    @contextmanager
    def get_cursor(self, dict_cursor=None):
        """Convenience context manager to provide a cursor.

        Usage example:
            with self.get_cursor() as cursor:
                cursor.execute('DROP TABLE super_important_stuff')
        Args:
            dict_cursor (bool): If True, use DictCursor.
        Yields:
            Cursor: initialized cursor object as per DB API v2.
        """
        if dict_cursor:
            cursor = self.snowflake_conn.cursor(connector.DictCursor)
        else:
            cursor = self.snowflake_conn.cursor()

        if self.autocommit:
            try:
                yield cursor
                cursor.close()
            except:
                cursor.close()
                self.snowflake_conn.close()
                raise
        else:
            try:
                cursor.execute('BEGIN')
                yield cursor
                self.snowflake_conn.commit()
                cursor.close()
            except:
                self.snowflake_conn.rollback()
                cursor.close()
                self.snowflake_conn.close()
                raise

    def execute(self, sql_template, params=None):
        """Execute an SQL statement.

        Args:
            sql_template (str): A template to be parametrized. All the
                identifiers (db name, schema name, table name, column names)
                should be already present.
            params (dict): A dict of params to bind to sql_template.
        """
        with self.get_cursor() as cursor:
            return cursor.execute(sql_template, params)

    def executemany(self, sql_template, params_list):
        """Execute an SQL statement many times with different params to bind.

        Args:
            sql_template (str): A template to be parametrized. All the
                identifiers (db name, schema name, table name, column names)
                should be already present.
            params_list (list): A list of params dicts to bind to sql_template.
        """
        with self.get_cursor() as cursor:
            cursor.executemany(sql_template, params_list)

    def fetchone(self, sql_template, params=None):
        """Execute SQL and fetch first tuple.

        Useful for SQL statements which are supposed to return just one row,
        like COUNT, etc..
        Args:
            sql_template (str): A template to be parametrized. All the
                identifiers (db name, schema name, table name, column names)
                should be already present.
            params (dict): A dict of params to bind to sql_template.
        Returns:
            tuple: A first row produced by executing an SQL statement.
        """
        with self.get_cursor() as cursor:
            cursor.execute(sql_template, params)
            return cursor.fetchone()

    def fetchall(self, sql_template, params=None):
        """Execute query and fetch all results.

        Args:
            sql_template (str): A template to be parametrized. All the
                identifiers (db name, schema name, table name, column names)
                should be already present.
            params (dict): A dict of params to bind to sql_template.
        Returns:
            list(tuple): List of rows with query result.
        """
        with self.get_cursor() as cursor:
            cursor.execute(sql_template, params)
            return cursor.fetchall()

    def fetchmany(self, sql_template, size, params=None):
        """Generator fetches set of n rows of a query result (n=size).

        Stops when no more results available.
        Args:
            sql_template (str): A template to be parametrized. All the
                identifiers (db name, schema name, table name, column names)
                should be already present.
            params (dict): A dict of params to bind to sql_template.
            size (int): Number of rows to fetch within one batch.
        Yields:
            list(tuple): List of maximum n rows with query result (n = size).
        """
        with self.get_cursor() as cursor:
            cursor.execute(sql_template, params)
            while True:
                batch = cursor.fetchmany(size)
                if not batch:
                    break
                yield batch

    def table_exists(self, table):
        """Check if table exists in the Snowflake.

        Schema name will be taken from sf_config.
        Args:
            table (str): Table name in Snowflake.
        Returns:
            bool: True if table exists, False othewise.
        """
        sql_template = (
            'SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES '
            'WHERE TABLE_NAME = %(table)s '
            'AND TABLE_SCHEMA = %(schema)s '
            'AND TABLE_CATALOG = %(db)s;')
        params = {
            'table': table.upper(),
            'schema': self.sf_config['schema'].upper(),
            'db': self.sf_config['db'].upper()}
        with self.get_cursor() as cursor:
            cursor.execute(sql_template, params)
            return bool(cursor.fetchone()[0])

    def drop_table(self, table, db=None, schema=None):
        """Drop table if it exists in the Snowflake.

        Default db and schema names will be taken from sf_config.
        Args:
            table (str): Table name in Snowflake.
            db (str): Optional db name.
            schema (str): Optional schema name.
        """
        sql_template = 'DROP TABLE IF EXISTS %(db)i.%(schema)i.%(table)i;'
        params = dict(
            db=db or self.sf_config['db'],
            schema=schema or self.sf_config['schema'],
            table=table)
        sql, _ = self.validator.format_identifiers(sql_template, params)

        with self.get_cursor() as cursor:
            cursor.execute(sql)

    def create_table(self, table, transient=False, db=None, schema=None):
        """Create table in the Snowflake.

        This is a simple helper, if you want more parameters, please use custom
        SQL statement and execute() method.
        Default db and schema names will be taken from sf_config.
        Args:
            table (str): Table name in Snowflake.
            transient (bool): True if you want to create a transient table.
            db (str): Optional db name.
            schema (str): Optional schema name.
        """
        sql_template = (
            'CREATE %(transient)i TABLE IF NOT EXISTS '
            '%(db)i.%(schema)i.%(table)i;')
        params = dict(
                transient='' if not transient else 'TRANSIENT',
                db=self.sf_config['db'] or db,
                schema=self.sf_config['schema'] or schema,
                table=table)
        sql, _ = self.validator.format_identifiers(sql_template, params)
        with self.get_cursor() as cursor:
            cursor.execute(sql)

    def create_table_like(
            self, table, source_table, db=None, schema=None, source_db=None,
            source_schema=None, transient=False):
        """Create table LIKE other table.

        Args:
            table (str): Destination table name in Snowflake.
            source_table (str): Source table name in Snowflake.
            db (str): Optional destination db name.
            schema (str): Optional destination schema name.
            source_db (str): Optional source db name.
            source_schema (str): Optional source schema name.
            transient (bool): True if you want to create a transient table.
        """
        sql_template = (
            'CREATE %(transient)i TABLE IF NOT EXISTS '
            '%(db)i.%(schema)i.%(table)i LIKE '
            '%(source_db)i.%(source_schema)i.%(source_table)i;')
        params = dict(
                transient='' if not transient else 'TRANSIENT',
                table=table,
                source_table=source_table,
                db=db or self.sf_config['db'],
                schema=schema or self.sf_config['schema'],
                source_db=source_db or self.sf_config['db'],
                source_schema=source_schema or self.sf_config['schema'])
        sql, _ = self.validator.format_identifiers(sql_template, params)
        with self.get_cursor() as cursor:
            cursor.execute(sql)

    def truncate_table(self, table, db=None, schema=None):
        """Create table in the Snowflake.

        Default db and schema names will be taken from sf_config.
        Args:
            table (str): Table name in Snowflake.
            db (str): Optional db name.
            schema (str): Optional schema name.
        """
        sql_template = 'TRUNCATE %(db)i.%(schema)i.%(table)i;'
        params = dict(
                db=db or self.sf_config['db'],
                schema=schema or self.sf_config['schema'],
                table=table)
        sql, _ = self.validator.format_identifiers(sql_template, params)
        with self.get_cursor() as cursor:
            cursor.execute(sql)

    def get_column_names(self, table, db=None, schema=None):
        """Get a list of column names of a Snowflake table.

        Args:
            table (str): A table name to extract column names for.
            db (str): Optional db name.
            schema (str): Optional schema name.
        Returns:
            list: A list of column names.
        """
        sql_template = 'DESC TABLE %(db)i.%(schema)i.%(table)i;'
        params = dict(
            db=db or self.sf_config['db'],
            schema=schema or self.sf_config['schema'],
            table=table)
        sql, _ = self.validator.format_identifiers(sql_template, params)
        result = self.fetchall(sql)
        column_names = []
        for table_desc in result:
            column_names.append(table_desc[0])
        return column_names

    def swap_tables(self, table1, table2):
        """Swap two tables.

        Args:
            table1 (str): A first table name.
            table2 (str): A second table name.
        """
        sql_template = (
            'ALTER TABLE %(db)i.%(schema)i.%(table1)i SWAP WITH '
            '%(db)i.%(schema)i.%(table2)i;')
        params = dict(
                db=self.sf_config['db'],
                schema=self.sf_config['schema'],
                table1=table1,
                table2=table2)
        sql, _ = self.validator.format_identifiers(sql_template, params)
        with self.get_cursor() as cursor:
            cursor.execute(sql)


class SQLLoader:
    """Loads SQL queries/templates from text files."""

    def __init__(self, sql_files_root='queries'):
        """Constructor for SQL loader.

        Args:
             sql_files_root (str): path to the sql directory
        """
        self.sql_files_root = os.path.realpath(sql_files_root)
        self._query_cash = {}

    def _load_query(self, query_name):
        """Load query from the disk.

        Args:
            query_name (str): name of the query (filename without extension).
        """
        path = os.path.join(self.sql_files_root, '{}.sql'.format(query_name))
        with open(path, 'r') as query_file:
            query = query_file.read()
        # get rid of docstring in the beginning of the SQL file
        # this is not required, but useful when debugging
        query = re.sub(r'^/\*.*\*/\n+', '', query, flags=re.DOTALL)
        return query

    def load_query(self, query_name):
        """Access queries by name.

        This method uses cache to reduce I/O operations.
        Usage example:
            >>> sql_loader = SQLLoader('sql/files/root')
            >>> sql_loader.load_query('query_name')
        Args:
            query_name (str): name of the query (filename without extension).
        Returns:
            str: SQL query template.
        """
        if query_name not in self._query_cash:
            self._query_cash[query_name] = self._load_query(query_name)
        return self._query_cash[query_name]


class SnowflakeBadResponse(Exception):
    """Snowflake Bad Response.

    Raised in case of bad response from Snowflake server
    """

    def __init__(self, response=None):
        """Create an instance of Snowflake Bad Response exception.

        Args:
            response (requests.Response): Bad response from Snowflake
        """
        self.response = response

    def __str__(self):
        """Get human-readable representation of error."""
        return 'Bad response from Snowflake server. Got {}'.format(
            self.response.status_code or 'unknown error'
        )


class SnowflakeMetadataConnector:
    """Snowflake Metadata connector.

    Provides a way to check Snowflake query metadata
    """

    def __init__(self):
        """Create a Snowflake API connector."""
        self.session = requests.Session()
        self.session.headers.update({
            'Accept': 'application/json',
            'Content-Type': 'application/json',
        })

    def authenticate(self):
        """Authenticate connector.

        Raises:
            SnowflakeAPIError: invalid credentials.
        """
        self.session.headers.update({
            'Authorization': 'Snowflake Token="{}"'.format(self._token),
        })

    def _get_url(self, path):
        return urllib.parse.urljoin(self._host, path)

    @cached_property
    def _host(self):
        return 'https://{}.snowflakecomputing.com'.format(
            config.SNOWFLAKE_CONFIG['account'])

    @cached_property
    def _credentials(self):
        response = self.session.post(
            self._get_url('/session/v1/login-request'),
            json={
                'data': {
                    'ACCOUNT_NAME': config.SNOWFLAKE_CONFIG['account'],
                    'LOGIN_NAME': config.SNOWFLAKE_CONFIG['user'],
                    'PASSWORD': config.SNOWFLAKE_CONFIG['password'],
                }
            }
        )
        if not response.ok:
            raise SnowflakeBadResponse(response)
        return response.json()['data']

    @property
    def _token(self):
        return self._credentials['token']

    def get_query_data(self, snowflake_query_id):
        response = self.session.get(
            self._get_url('/monitoring/queries/{}'.format(snowflake_query_id))
        )
        if not response.ok:
            raise SnowflakeBadResponse(response)
        return response.json()['data']

    def get_query_scan_bytes_number(self, cursor):
        """Get number of bytes scanned not from cache.

        Args:
            cursor (snowflake.connector.cursor.SnowflakeCursor):
                Snowflake query cursor

        Returns:
            int: number of bytes scanned not from cache
        """
        data = self.get_query_data(cursor.sfqid)
        return sum(
            query['stats'].get('ioRemoteFdnReadBytes', 0)
            for query in data['queries'])