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

import json
import pandas as pd

from rasgoql.data.base import DataWarehouse, DWCredentials
from rasgoql.errors import (
    DWCredentialsWarning,
    DWConnectionError,
    DWQueryError,
    PackageDependencyWarning,
    ParameterException,
    SQLWarning,
    TableAccessError,
    TableConflictException,
)
from rasgoql.imports import sf_connector, write_pandas
from rasgoql.primitives.enums import check_response_type, check_write_method, check_write_table_type
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

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,
                "database": self.database,
                "schema": self.schema,
                "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)
        account = os.getenv('SNOWFLAKE_ACCOUNT')
        user = os.getenv('SNOWFLAKE_USER')
        password = os.getenv('SNOWFLAKE_PASSWORD')
        role = os.getenv('SNOWFLAKE_ROLE')
        warehouse = os.getenv('SNOWFLAKE_WAREHOUSE')
        database = os.getenv('SNOWFLAKE_DATABASE')
        schema = os.getenv('SNOWFLAKE_SCHEMA')
        if not all([account, user, password, role, warehouse, database, schema]):
            raise DWCredentialsWarning(
                'Your env file is missing expected credentials. Consider running '
                'SnowflakeCredentials(*args).to_env() to repair this.'
            )
        return cls(
            account,
            user,
            password,
            role,
            warehouse,
            database,
            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 = self.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_write_table_type(table_type)
        fqtn = self.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}'

    @default_namespace.setter
    def default_namespace(
        self,
        new_namespace: str,
    ):
        """
        Setter method for the `default_namespace` property
        """
        namespace = self.validate_namespace(new_namespace)
        db, schema = self.parse_namespace(namespace)
        self.default_database = db
        self.default_schema = schema

    def execute_query(
        self,
        sql: str,
        response: str = 'tuple',
        acknowledge_risk: bool = False,
        batches: 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, batches=batches)
        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 = self.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 = self.magic_fqtn_handler(fqtn, self.default_namespace)
        database, schema, table = self.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_or_sql: str,
    ) -> List[Tuple[str, str]]:
        """
        Return the schema of a table, view, or select statement

        Params:
        `fqtn_or_sql`: str:
            Either a Fully-qualified table name (database.schema.table)
            or a SQL select statement that will create a view.
        """
        # Check for SQL
        if self._is_select_statement(fqtn_or_sql):
            query_response = self.connection.cursor().describe(fqtn_or_sql)
            return [(row[0], convert_to_type(row[1], row[4], row[5])) for row in query_response]
        # Otherwise assume fqtn:
        fqtn = self.magic_fqtn_handler(fqtn_or_sql, self.default_namespace)
        if self._table_exists(fqtn):
            query_response = self.execute_query(f"DESC TABLE {fqtn}", response='dict')
            return [(row['name'], row['type']) for row in query_response]
        else:
            raise TableAccessError(f'Table {fqtn} does not exist or cannot be accessed.')

    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,
    ) -> str:
        """
        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 = self.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 fqtn
        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 = self.magic_fqtn_handler(fqtn, self.default_namespace)
        do_i_exist, _, _ = self.get_object_details(fqtn)
        return do_i_exist

    # --------------------------
    # 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):
            if exception.errno == 3001:
                raise TableAccessError(
                    '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'
                ) from 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.'
                ) from exception
        if isinstance(exception, sf_connector.errors.ServiceUnavailableError):
            raise DWConnectionError(
                'Snowflake is unavailable. Please check that your are using '
                'a valid account identifier, that you have internet access, and '
                'that http connections to Snowflake are whitelisted in your env. '
                'Finally check https://status.snowflake.com/ for outage status.'
            ) from exception
        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,
        batches: Optional[bool] = False,
    ) -> pd.DataFrame:
        """
        Run a query string and return results in a pandas DataFrame
        """
        cursor = None
        try:
            cursor = self.connection.cursor()
            cursor.execute(query, params)
            if batches:
                return cursor.fetch_pandas_batches()
            return cursor.fetch_pandas_all()
        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()


def convert_to_type(type_code: int, precision: int, scale: int) -> str:
    """
    Accepts values from a Snowflake ResultMetadata record
    and returns a Snowflake data type

    https://docs.snowflake.com/en/user-guide/python-connector-api.html#type-codes
    ResultMetadata(name='', type_code=, display_size=, internal_size=, precision=, scale=, is_nullable=)
    """
    type_code_map = {
        0: "FIXED",  # NUMBER/INT
        1: "REAL",  # REAL
        2: "TEXT",  # VARCHAR/STRING
        3: "DATE",  # DATE
        4: "TIMESTAMP",  # TIMESTAMP
        5: "VARIANT",  # VARIANT
        6: "TIMESTAMP_LTZ",  # TIMESTAMP_LTZ
        7: "TIMESTAMP_TZ",  # TIMESTAMP_TZ
        8: "TIMESTAMP_NTZ",  # TIMESTAMP_TZ
        9: "OBJECT",  # OBJECT
        10: "ARRAY",  # ARRAY
        11: "BINARY",  # BINARY
        12: "TIME",  # TIME
        13: "BOOLEAN",  # BOOLEAN
    }
    data_type = type_code_map[type_code]
    if data_type == 'REAL':
        return 'FLOAT'
    if data_type == 'FIXED':
        return f'NUMBER({precision or 0},{scale or 0})'
    if 'TIMESTAMP' in data_type:
        return f'{data_type}({scale or 0})'
    return data_type
