from typing import Any, Dict, Iterable, List, Optional, Union, cast

from app_types import S3Path
from config import SNOWFLAKE_ENV, SNOWFLAKE_FORCE_SYNC_QUERY
from constants import CSV_DELIMITER, CSV_NULL_VALUE, CSV_QUOTE_CHAR
from db.sf_connector import SnowflakeConnector

from ..mixins import S3Mixin


__all__ = ["SnowflakeMixin"]


Result = List[Dict[str, Any]]


class SnowflakeMixin(S3Mixin):
    table_location: str = f"DNA.DNA_PUBLIC_{SNOWFLAKE_ENV.upper()}."
    table_location_replace: str = "DNA.DNA_PUBLIC."

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.snowflake_instance: SnowflakeConnector = SnowflakeConnector()

    def _pre_run(self):
        super()._pre_run()
        self.snowflake_instance.init()

    def _post_run(self):
        super()._post_run()
        self.snowflake_instance.close()

    def run_raw_query(
        self, query: str, params: Optional[Dict[str, Any]] = None, is_async: bool = True
    ) -> Union[str, Result]:
        is_async = not SNOWFLAKE_FORCE_SYNC_QUERY and is_async

        query = self._prepare_query(query)
        params = self._prepare_params(params)

        if is_async:
            return self.snowflake_instance.execute_async_query(query, params)

        result = self.snowflake_instance.execute_query(query, params)
        return self._filter_results(result)

    def get_query_result(self, query_id: str) -> List[Dict[str, Any]]:
        results = self.snowflake_instance.get_query_result(query_id)
        return self._filter_results(results)

    def export(
        self, query: str, prefix: S3Path, *, is_async: bool = False, options: Iterable[str] = ("OVERWRITE = TRUE",)
    ) -> Union[str, Result]:
        self.wipe_folder(prefix)

        self.logger.info(f"Exporting raw data to '{prefix}'")
        options = "\n".join(options)

        # TODO: use role for S3 auth
        query = f"""
            COPY INTO '{self.get_full_s3_path(prefix)}/'
            FROM ({query})
            {options}
            CREDENTIALS = (
                AWS_KEY_ID = '{self.aws_credentials.access_key}'
                AWS_SECRET_KEY = '{self.aws_credentials.secret_key}'
                AWS_TOKEN = '{self.aws_credentials.token}'
            )
            FILE_FORMAT = (
                TYPE = CSV
                COMPRESSION = NONE
                FIELD_DELIMITER = '{CSV_DELIMITER}'
                FIELD_OPTIONALLY_ENCLOSED_BY = '{CSV_QUOTE_CHAR}'
                null_if = ('{CSV_NULL_VALUE}')
            )
            HEADER = TRUE
            DETAILED_OUTPUT = FALSE;
        """
        return self.run_raw_query(query, is_async=is_async)

    def create_empty_table(self, table: str) -> Result:
        self.logger.info(f"Creating table '{table}'")
        query = f"CREATE TABLE IF NOT EXISTS {self.table_location}{table} (T INT);"
        return cast(Result, self.run_raw_query(query, is_async=False))

    def create_table_from_query(self, table: str, query: str, query_params: Optional[Dict[str, Any]] = None) -> Result:
        self.logger.info(f"Creating table '{table}' from query")
        query = f"CREATE TABLE {self.table_location}{table} AS ({query});"
        return cast(Result, self.run_raw_query(query, query_params, is_async=False))

    def swap_tables(self, table: str, table_alt: str) -> Result:
        self.logger.info(f"Swapping tables '{table_alt}' <-> '{table}'")
        query = f"ALTER TABLE IF EXISTS {self.table_location}{table_alt} SWAP WITH {self.table_location}{table};"
        return cast(Result, self.run_raw_query(query, is_async=False))

    def drop_table(self, table: str) -> Result:
        self.logger.info(f"Deleting '{table}'")
        query = f"DROP TABLE IF EXISTS {self.table_location}{table};"
        return cast(Result, self.run_raw_query(query, is_async=False))

    @staticmethod
    def _filter_results(results: List[Dict[str, Any]]) -> Result:
        return [{k.lower(): v for k, v in item.items()} for item in results]

    @staticmethod
    def _prepare_params(params: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        return {v.upper(): k for v, k in params.items()} if params else {}

    @classmethod
    def _prepare_query(cls, query: str) -> str:
        return query.replace(cls.table_location_replace, cls.table_location)