"""
Snowflake DataWarehouse classes
"""
import logging
import os
from typing import List, Optional, Union

import json
import pandas as pd

from rasgoql.data.base import DataWarehouse, DWCredentials
from rasgoql.errors import (
    DWCredentialsWarning, DWConnectionError, DWQueryError,
    PackageDependencyWarning, ParameterException,
    SQLWarning, TableConflictException
)
from rasgoql.imports import sf_connector, write_pandas
from rasgoql.primitives.enums import (
    check_response_type, check_table_type, check_write_method
)
from rasgoql.utils.creds import load_env, save_env
from rasgoql.utils.df import cleanse_sql_dataframe, generate_dataframe_ddl
from rasgoql.utils.messaging import verbose_message
from rasgoql.utils.sql import (
    is_scary_sql, magic_fqtn_handler,
    parse_fqtn, parse_namespace,
    validate_namespace
)

logging.basicConfig()
logger = logging.getLogger('Snowflake DataWarehouse')
logger.setLevel(logging.INFO)


class SnowflakeCredentials(DWCredentials):
    """
    Snowflake Credentials
    """
    dw_type = 'snowflake'

    def __init__(
            self,
            account: str,
            user: str,
            password: str,
            role: str,
            warehouse: str,
            database: str,
            schema: str
        ):
        if sf_connector is None:
            raise PackageDependencyWarning(
                'Missing a required python package to run Snowflake. '
                'Please download the Snowflake package by running: '
                'pip install rasgoql[snowflake]')
        self.account = account
        self.user = user
        self.password = password
        self.role = role
        self.warehouse = warehouse
        self.database = database
        self.schema = schema

    def __repr__(self) -> str:
        return json.dumps(
            {
                "account": self.account,
                "user": self.user,
                "role": self.role,
                "warehouse": self.warehouse,
            }
        )

    @classmethod
    def from_env(
            cls,
            filepath: str = None
        ) -> 'SnowflakeCredentials':
        """
        Creates an instance of this Class from a .env file on your machine
        """
        load_env(filepath)
        if not all([
                os.getenv('snowflake_account'),
                os.getenv('snowflake_user'),
                os.getenv('snowflake_password'),
                os.getenv('snowflake_role'),
                os.getenv('snowflake_warehouse'),
                os.getenv('snowflake_database'),
                os.getenv('snowflake_schema')
        ]):
            raise DWCredentialsWarning(
                'Your env file is missing expected credentials. Consider running '
                'SnowflakeCredentials(*args).to_env() to repair this.'
            )
        return cls(
            os.getenv('snowflake_account'),
            os.getenv('snowflake_user'),
            os.getenv('snowflake_password'),
            os.getenv('snowflake_role'),
            os.getenv('snowflake_warehouse'),
            os.getenv('snowflake_database'),
            os.getenv('snowflake_schema')
        )

    def to_dict(self) -> dict:
        """
        Returns a dict of the credentials
        """
        return {
            "account": self.account,
            "user": self.user,
            "password": self.password,
            "role": self.role,
            "warehouse": self.warehouse,
            "database": self.database,
            "schema": self.schema
        }

    def to_env(
            self,
            filepath: str = None,
            overwrite: bool = False
        ):
        """
        Saves credentials to a .env file on your machine
        """
        creds = {
            "snowflake_account": self.account,
            "snowflake_user": self.user,
            "snowflake_password": self.password,
            "snowflake_role": self.role,
            "snowflake_warehouse": self.warehouse,
            "snowflake_database": self.database,
            "snowflake_schema": self.schema,
        }
        return save_env(creds, filepath, overwrite)


class SnowflakeDataWarehouse(DataWarehouse):
    """
    Snowflake DataWarehouse
    """
    dw_type = 'snowflake'
    credentials_class = SnowflakeCredentials

    def __init__(self):
        super().__init__()
        self.credentials: dict = None
        self.connection: sf_connector.SnowflakeConnection = None
        self.default_database = None
        self.default_schema = None

    # ---------------------------
    # Core Data Warehouse methods
    # ---------------------------
    def change_namespace(
            self,
            namespace: str
        ) -> None:
        """
        Changes the default namespace of your connection

        Params:
        `namespace`: str:
            namespace (database.schema)
        """
        namespace = self._validate_namespace(namespace)
        database, schema = parse_namespace(namespace)
        try:
            self.execute_query(f'USE DATABASE {database}')
            self.execute_query(f'USE SCHEMA {schema}')
            self.default_namespace = namespace
            self.default_database = database
            self.default_schema = schema
            verbose_message(
                f"Namespace reset to {self.default_namespace}",
                logger
            )
        except Exception as e:
            self._error_handler(e)

    def connect(
            self,
            credentials: Union[dict, SnowflakeCredentials]
        ):
        """
        Connect to Snowflake

        Params:
        `credentials`: dict:
            dict (or DWCredentials class) holding the connection credentials
        """
        if isinstance(credentials, SnowflakeCredentials):
            credentials = credentials.to_dict()

        # This allows you to track what queries were run by RasgoQL in your history tab
        credentials.update({
            "application": "rasgoql",
            "session_parameters": {
                "QUERY_TAG": "rasgoql"
            }
        })
        try:
            self.credentials = credentials
            self.default_database = credentials.get('database')
            self.default_schema = credentials.get('schema')
            self.connection = sf_connector.connect(**credentials)
            verbose_message(
                "Connected to Snowflake",
                logger
            )
        except Exception as e:
            self._error_handler(e)

    def close_connection(self):
        """
        Close connection to Snowflake
        """
        try:
            if self.connection:
                self.connection.close()
            self.connection = None
            verbose_message(
                "Connection to Snowflake closed",
                logger
            )
        except Exception as e:
            self._error_handler(e)

    def create(
            self,
            sql: str,
            fqtn: str,
            table_type: str = 'VIEW',
            overwrite: bool = False
        ):
        """
        Create a view or table from given SQL

        Params:
        `sql`: str:
            query that returns data (i.e. can be wrapped in a CREATE TABLE statement)
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)
            Name for the new table
        `table_type`: str:
            One of values: [view, table]
        `overwrite`: bool
            pass True when this table name already exists in your DataWarehouse
            and you know you want to overwrite it
            WARNING: This will completely overwrite data in the existing table
        """
        table_type = check_table_type(table_type)
        fqtn = magic_fqtn_handler(fqtn, self.default_namespace)
        if self._table_exists(fqtn) and not overwrite:
            msg = f'A table or view named {fqtn} already exists. ' \
                   'If you are sure you want to overwrite it, ' \
                   'pass in overwrite=True and run this function again'
            raise TableConflictException(msg)
        query = f"CREATE OR REPLACE {table_type} {fqtn} COMMENT='rasgoql' AS {sql}"
        self.execute_query(query, acknowledge_risk=True, response='None')
        return fqtn

    @property
    def default_namespace(self) -> str:
        """
        Returns the default database.schema of this connection
        """
        return f'{self.default_database}.{self.default_schema}'

    def execute_query(
            self,
            sql: str,
            response: str = 'tuple',
            acknowledge_risk: bool = False
        ):
        """
        Run a query against Snowflake and return all results

        `sql`: str:
            query text to execute
        `response`: str:
            Possible values: [dict, df, None]
        `acknowledge_risk`: bool:
            pass True when you know your SQL statement contains
            a potentially dangerous or data-altering operation
            and still want to run it against your DataWarehouse
        """
        response = check_response_type(response)
        if is_scary_sql(sql) and not acknowledge_risk:
            msg = 'It looks like your SQL statement contains a ' \
                  'potentially dangerous or data-altering operation.' \
                  'If you are positive you want to run this, ' \
                  'pass in acknowledge_risk=True and run this function again.'
            raise SQLWarning(msg)
        verbose_message(
            f"Executing query: {sql}",
            logger
        )
        if response == 'DICT':
            return self._execute_dict_cursor(sql)
        if response == 'DF':
            return self._execute_df_cursor(sql)
        return self._execute_string(sql, ignore_results=(response == 'NONE'))

    def get_ddl(
            self,
            fqtn: str
        ) -> str:
        """
        Returns the create statement for a table or view

        `fqtn`: str:
            Fully-qualified Table Name (database.schema.table)
        """
        fqtn = magic_fqtn_handler(fqtn, self.default_namespace)
        sql = f"SELECT GET_DDL('TABLE', '{fqtn}') AS DDL"
        query_response = self.execute_query(sql, response='dict')
        return query_response[0]['DDL']

    def get_object_details(
            self,
            fqtn: str
        ) -> tuple:
        """
        Return details of a table or view in Snowflake

        Params:
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)

        Response:
            object exists: bool
            is rasgo object: bool
            object type: [table|view|unknown]
        """
        fqtn = magic_fqtn_handler(fqtn, self.default_namespace)
        database, schema, table = parse_fqtn(fqtn)
        sql = f"SHOW OBJECTS LIKE '{table}' IN {database}.{schema}"
        result = self.execute_query(sql, response='dict')
        obj_exists = len(result) > 0
        is_rasgo_obj = False
        obj_type = 'unknown'
        if obj_exists:
            is_rasgo_obj = (result[0].get('comment') == 'rasgoql')
            obj_type = result[0].get('kind')
        return obj_exists, is_rasgo_obj, obj_type

    def get_schema(
            self,
            fqtn: str
        ) -> dict:
        """
        Return the schema of a table or view

        Params:
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)
        """
        fqtn = magic_fqtn_handler(fqtn, self.default_namespace)
        sql = f"DESC TABLE {fqtn}"
        query_response = self.execute_query(sql, response='dict')
        return query_response

    def list_tables(
            self,
            database: str = None,
            schema: str = None
        ) -> pd.DataFrame:
        """
        List all tables and views available in default namespace

        Params:
        `database`: str:
            override database
        `schema`: str:
            override schema
        """
        select_clause = "SELECT TABLE_NAME, " \
                        "TABLE_CATALOG||'.'||TABLE_SCHEMA||'.'||TABLE_NAME AS FQTN, " \
                        "CASE TABLE_TYPE WHEN 'BASE TABLE' THEN 'TABLE' ELSE TABLE_TYPE END AS TABLE_TYPE, " \
                        "ROW_COUNT, CREATED, LAST_ALTERED "
        from_clause = " FROM INFORMATION_SCHEMA.TABLES "
        if database:
            from_clause = f" FROM {database.upper()}.INFORMATION_SCHEMA.TABLES "
        where_clause = f" WHERE TABLE_SCHEMA = '{schema.upper()}'" if schema else ""
        sql = select_clause + from_clause + where_clause
        return self.execute_query(sql, response='df', acknowledge_risk=True)

    def preview(
            self,
            sql: str,
            limit: int = 10
        ) -> pd.DataFrame:
        """
        Returns 10 records into a pandas DataFrame

        Params:
        `sql`: str:
            SQL statment to run
        `limit`: int:
            Records to return
        """
        return self.execute_query(
            f'{sql} LIMIT {limit}',
            response='df',
            acknowledge_risk=True
        )

    def save_df(
            self,
            df: pd.DataFrame,
            fqtn: str,
            method: str = None
        ):
        """
        Creates a table in Snowflake from a pandas Dataframe

        Params:
        `df`: pandas DataFrame:
            DataFrame to upload
        `fqtn`: str:
            Fully-qualied table name (database.schema.table)
            Name for the new table
        `method`: str
            Values: [append, replace]
            when this table already exists in your DataWarehouse,
            pass append: to add dataframe rows to it
            pass replace: to overwrite it with dataframe rows
                WARNING: This will completely overwrite data in the existing table
        """
        if method:
            method = check_write_method(method)
        fqtn = magic_fqtn_handler(fqtn, self.default_namespace)
        table_exists = self._table_exists(fqtn)
        if table_exists and not method:
            msg = f"A table named {fqtn} already exists. " \
                   "If you are sure you want to write over it, pass in " \
                   "method='append' or method='replace' and run this function again"
            raise TableConflictException(msg)
        try:
            cleanse_sql_dataframe(df)
            # If the table does not exist or we've received instruction to replace
            # Issue a create or replace statement before we insert data
            if not table_exists or method == 'REPLACE':
                create_stmt = generate_dataframe_ddl(df, fqtn)
                create_stmt += " COMMENT='rasgoql' "
                self.execute_query(create_stmt, response='None', acknowledge_risk=True)
            success, chunks, rows, output = write_pandas(
                conn=self.connection,
                df=df,
                table_name=fqtn,
                quote_identifiers=False
            )
            return success, chunks, rows, output
        except Exception as e:
            self._error_handler(e)

    # ---------------------------
    # Core Data Warehouse helpers
    # ---------------------------
    def _table_exists(
            self,
            fqtn: str
        ) -> bool:
        """
        Check for existence of fqtn in the Data Warehouse and return a boolean

        Params:
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)
        """
        fqtn = magic_fqtn_handler(fqtn, self.default_namespace)
        do_i_exist, _, _ = self.get_object_details(fqtn)
        return do_i_exist

    def _validate_namespace(
            self,
            namespace: str
        ) -> None:
        """
        Checks a namespace string for compliance with Snowflake format

        Params:
        `namespace`: str:
            namespace (database.schema)
        """
        try:
            validate_namespace(namespace)
            return namespace.upper()
        except ValueError:
            raise ParameterException("Snowflake namespaces should be format: DATABASE.SCHEMA")

    # --------------------------
    # Snowflake specific helpers
    # --------------------------
    def _error_handler(
            self,
            exception: Exception,
            query: str = None
        ) -> None:
        """
        Handle Snowflake exceptions that need additional info
        """
        verbose_message(
            f"Exception occurred while running query: {query}",
            logger
        )
        if exception is None:
            return
        if isinstance(exception, sf_connector.errors.ProgrammingError):
            # TODO:
            # TableAccessError
            # Insufficient privileges to write to namespace
            if exception.errno == ...:
                raise DWQueryError(
                    'You do not have access to operate on this object. '
                    'Two possible ways to resolve: '
                    'Connect with different credentials that have the proper access.'
                    'Or run `.change_namespace` on your SQLChain to write it to a '
                    'namespace your credentials can access')
            raise DWQueryError(exception)
        if isinstance(exception, sf_connector.errors.DatabaseError):
            if exception.errno == 250001:
                raise DWConnectionError(
                    'Invalid username / password, please check that your '
                    ' credentials are correct and try to reconnect'
                )
            raise DWConnectionError(exception)
        if isinstance(exception, sf_connector.errors.ServiceUnavailableError):
            raise DWConnectionError('Snowflake service is unavailable')
        raise exception

    def _execute_dict_cursor(
            self,
            query: str
        ) -> List[dict]:
        """
        Run a query string and return results in a Snowflake DictCursor

        PRO:
        Results are callable by column name
        for row in data:
            row['COL_NAME']

        CON:
        Query string must be a single statement (only one ;) or Snowflake returns an error
        """
        cursor = None
        try:
            cursor = self.connection.cursor(sf_connector.DictCursor)
            query_return = cursor.execute(query).fetchall()
            return query_return
        except Exception as e:
            self._error_handler(e)
        finally:
            if cursor:
                cursor.close()

    def _execute_df_cursor(
            self,
            query: str,
            params: Optional[dict] = None
        ) -> pd.DataFrame:
        """
        Run a query string and return results in a pandas DataFrame
        """
        cursor = None
        try:
            cursor = self.connection.cursor()
            query_return = cursor.execute(query, params).fetch_pandas_all()
            return query_return
        except Exception as e:
            self._error_handler(e)
        finally:
            if cursor:
                cursor.close()

    def _execute_string(
            self,
            query: str,
            ignore_results: bool = False
        ) -> List[tuple]:
        """
        Execute a query string against the Data Warehouse connection and fetch all results
        """
        query_returns = []
        cursor = None
        try:
            for cursor in self.connection.execute_string(query, return_cursors=(not ignore_results)):
                for query_return in cursor:
                    query_returns.append(query_return)
            return query_returns
        except Exception as e:
            self._error_handler(e)
        finally:
            if cursor:
                cursor.close()
