"""
Generic SQLAlchemy DataWarehouse classes
"""
from __future__ import annotations
from abc import abstractmethod
import logging
from typing import List, Optional, Tuple, Union
from urllib.parse import quote_plus as urlquote

import pandas as pd

from rasgoql.data.base import DataWarehouse, DWCredentials
from rasgoql.errors import (
    DWConnectionError,
    SQLWarning,
    TableAccessError,
    TableConflictException,
)
from rasgoql.imports import alchemy_engine, alchemy_session, alchemy_exceptions
from rasgoql.primitives.enums import (
    check_response_type,
    check_table_type,
    check_write_method,
)
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("SQLAlchemy DataWarehouse")
logger.setLevel(logging.INFO)

# Each DB derived from the generic SQLAlchemy class will have slightly
# different connection strings. Build a unique Credentials class for each DB


class SQLAlchemyDataWarehouse(DataWarehouse):
    """
    Base SQLAlchemy DataWarehouse
    """

    dw_type = None
    credentials_class = None

    def __init__(self):
        super().__init__()
        self.credentials: Optional[Union[dict, DWCredentials]]
        self.connection: alchemy_session
        self.database = None
        self.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)
        """
        raise NotImplementedError(
            "Connecting to a new Database in a single session is not supported "
            "by the SQLAlchemy Engine. Please build a new connection."
        )

    @abstractmethod
    def connect(self, credentials: Union[dict, DWCredentials]):
        """
        Connect to DB
        Params:
        `credentials`: dict:
            dict  holding the connection credentials
        """
        if isinstance(credentials, DWCredentials):
            credentials = credentials.to_dict()

        try:
            self.credentials = credentials
            self.database = credentials.get("database")
            self.schema = credentials.get("schema")
            self.connection = alchemy_session(self._engine, autocommit=True)
        except Exception as e:
            self._error_handler(e)

    def close_connection(self):
        """
        Close connection
        """
        try:
            if self.connection:
                self.connection.close()
            self.connection = None
            verbose_message("Connection 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
        If a derived class's DB cannot execute against a FQTN (e.g., Postgres),
        override this method or SQL execution will fail
        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 = self.magic_fqtn_handler(fqtn, self.default_namespace)
        if self._table_exists(fqtn=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} 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.database}.{self.schema}"

    @default_namespace.setter
    def default_namespace(self, new_namespace: str):
        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,
        **kwargs,
    ) -> Union[list[dict], pd.DataFrame, list[tuple], None]:
        """
        Run a query against DB 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._query_into_dict(sql)
        if response == "DF":
            return self._query_into_df(sql)
        return self._execute_string(sql, ignore_results=(response == "NONE"))

    @abstractmethod
    def get_ddl(self, fqtn: str) -> pd.DataFrame:
        """
        Returns a DataFrame describing the column in the table
        This method should be overridden by derived classes since optimal DDL
        selection will vary by DB type
        `fqtn`: str:
            Fully-qualified Table Name (database.schema.table)
        """
        fqtn = self.magic_fqtn_handler(fqtn, self.default_namespace)
        db_name, schema_name, table_name = self.parse_fqtn(fqtn)
        sql = (
            f"select table_schema, table_name, column_name, data_type, "
            f"character_maximum_length, column_default, is_nullable from "
            f"INFORMATION_SCHEMA.COLUMNS where table_name = '{table_name}' "
            f"and table_schema = '{schema_name}';"
        )
        query_response = self.execute_query(sql, response="DF")
        return query_response

    @abstractmethod
    def get_object_details(self, fqtn: str) -> tuple:
        """
        Return details of a table or view
        This method should be overridden by derived classes since optimal DDL
        selection will vary by DB type
        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.
        """
        query_sql = (
            "SELECT * FROM INFORMATION_SCHEMA.TABLES "
            "WHERE table_catalog = '{database}' "
            "AND table_schema = '{schema}' "
            "AND table_name = '{table}'"
        )
        # Check for SQL
        if self._is_select_statement(fqtn_or_sql):
            self.create(fqtn_or_sql, "temp_schema", table_type="view", overwrite=True)
            query_response = self.execute_query(
                query_sql.format(database=self.database, schema=self.schema, table="temp_schema"),
                response="dict",
            )
            self.execute_query(
                "DROP VIEW temp_schema",
                response="none",
                acknowledge_risk=True,
            )
            return [(row["name"], row["type"]) for row in query_response]
        # Otherwise assume fqtn:
        fqtn = self.magic_fqtn_handler(fqtn_or_sql, self.default_namespace)
        database, schema, table = self.parse_fqtn(fqtn)
        if self._table_exists(fqtn):
            query_response = self.execute_query(
                query_sql.format(database=database, schema=schema, table=table), 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"
        )
        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 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)
        database, schema, table = self.parse_fqtn(fqtn)
        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)
                self.execute_query(create_stmt, response="None", acknowledge_risk=True)
            df.to_sql(
                table,
                self._engine,
                schema=schema,
                if_exists=method.lower(),
                index=False,
                chunksize=1000,
            )
            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

    # --------------------------
    # SQLAlchemy and derived class helpers
    # --------------------------
    @property
    @abstractmethod
    def _engine(self):
        """
        Returns a SQLAlchemy engine
        """
        engine_url = (
            f"{self.credentials.get('dw_type')}://"
            f"{self.credentials.get('username')}:"
            f"{urlquote(self.credentials.get('password'))}"
            f"@{self.credentials.get('host')}:"
            f"{self.credentials.get('port')}/"
            f"{self.credentials.get('database')}"
        )
        return alchemy_engine(engine_url)

    def _error_handler(self, exception: Exception, query: str = None) -> None:
        """
        Handle SQLAlchemy exceptions that need additional info
        """
        verbose_message(f"Exception occurred while running query: {query}", logger)
        if exception is None:
            return
        if isinstance(exception, alchemy_exceptions.DisconnectionError):
            raise DWConnectionError(
                "Disconnected from DataWarehouse. Please validate connection or reconnect."
            ) from exception
        raise exception

    def _execute_string(self, query: str, ignore_results: bool = False) -> list[tuple]:
        """
        Execute a query string against the DataWarehouse connection and fetch all results
        """
        try:
            results = self.connection.execute(query)
            if ignore_results:
                return
            return list(results)
        except Exception as e:
            self._error_handler(e)

    def _query_into_dict(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
        """
        try:
            query_return = self.connection.execute(query).__dict__
            return query_return
        except Exception as e:
            self._error_handler(e)

    def _query_into_df(self, query: str) -> pd.DataFrame:
        """
        Run
        """
        try:
            query_result = self.connection.execute(query)
            query_return_df = pd.DataFrame(query_result.all())
            response_cols = list(query_result.keys())
            query_return_df.columns = response_cols
            return query_return_df
        except Exception as e:
            self._error_handler(e)
