import json
import logging

from datetime import datetime
from psi_elt.abstract_elt_job import AbstractELTJob, TRANSFORMER_ELT_JOB_TYPE
from typeguard import typechecked
from typing import Any, List, Union


@typechecked
class CustomTransformer(AbstractELTJob):
    '''A CustomTransformer is the third of the ELT jobs that is responsible for
    reading data (from S3, Athena, or Timestream) and transforming the data for
    use in a data mart.

    This class should be treated as abstract and should be extended.
    Child classes must override the _read_data and _transform functions.
    This class includes utilties for reading from S3, Athena, and Timestream.
    '''

    def __init__(
        self,
        name: str,
        capture_time: datetime,
        aws_session: Any = None,
        logger: logging.Logger = None,
        athena_api: Any = None,
        timestream_reader: Any = None
    ):
        super().__init__(
            name=name,
            elt_job_type=TRANSFORMER_ELT_JOB_TYPE,
            capture_time=capture_time,
            aws_session=aws_session,
            logger=logger
        )
        self.__athena_api = athena_api
        self.__timestream_reader = timestream_reader

    @property
    def athena_api(self) -> Any:
        return self.__athena_api

    @property
    def timestream_reader(self) -> Any:
        return self.__timestream_reader

    def transform(
        self,
        is_test: bool,
        s3_uri: str,
        read_data_kwargs: dict = None,
        transform_kwargs: dict = None
    ) -> Any:
        read_data_kwargs = read_data_kwargs or {}
        read_data_kwargs = {key: val for key, val in read_data_kwargs.items() if val}
        if s3_uri:
            raw_data = json.loads(self.get_s3_object(s3_uri))
        else:
            raw_data = self._read_data(is_test, **read_data_kwargs)

        transform_kwargs = transform_kwargs or {}
        transform_kwargs = {key: val for key, val in transform_kwargs.items() if val}
        return self._transform(raw_data, is_test, **transform_kwargs)

    def _read_data(self, is_test: bool, **read_data_kwargs) -> Any:
        raise Exception('CustomTransformer._read_data must be overridden by the child class')

    def _transform(self, raw_data: Any, is_test: bool, **transform_kwargs) -> Any:
        raise Exception('CustomTransformer._transform must be overridden by the child class')

    def run_athena_query(
        self,
        database: str,
        table: str,
        query: str,
        data_catalog: str = None,
        workgroup: str = None,
        role_arn: str = None,
        region: str = None,
        session_name: str = None,
        reload_partitions: bool = False
    ) -> Union[dict, list]:
        assert database, 'A database is required to run an Athena query'
        assert table, 'A table is required to run an Athena query'
        assert query, 'A SQL query is required to run an Athena query'
        try:
            if role_arn:
                self.athena_api.assume_role(
                    role_arn=role_arn,
                    region=region,
                    session_name=session_name
                )
            if data_catalog:
                self.athena_api.data_catalog = data_catalog
            if workgroup:
                self.athena_api.workgroup = workgroup
            if database:
                self.athena_api.database = database

            return self.athena_api.run_athena_query(
                table=table,
                query=query,
                download_results=True,
                reload_partitions=reload_partitions
            )
        except Exception as e:
            self.logger.error(f'Exception occurred when running Athena query: {e}')
            raise e

    def run_timestream_query(
        self,
        database_name: str,
        table_name: str,
        projection_expression: str = None,
        measure_value: float = None,
        measure_value_op: float = None,
        time: datetime = None,
        min_time: datetime = None,
        max_time: datetime = None,
        limit: int = None,
        **kwargs
    ) -> List[dict]:
        assert database_name, 'A database is required to run a Timestream query'
        assert table_name, 'A table is required to run a Timestream query'
        try:
            self.timestream_reader.does_table_exist(
                database_name=database_name,
                table_name=table_name,
                raise_err=True
            )
            return self.timestream_reader.read(
                database_name=database_name,
                table_name=table_name,
                projection_expression=projection_expression,
                measure_value=measure_value,
                measure_value_op=measure_value_op,
                time=time,
                min_time=min_time,
                max_time=max_time,
                limit=limit,
                **kwargs
            )
        except Exception as e:
            self.logger.error(f'Exception occurred when running Timestream query: {e}')
            raise e

    def run_timestream_custom_query(
        self,
        database_name: str,
        table_name: str,
        sql: str = None,
        file_path: str = None
    ) -> List[dict]:
        assert database_name, 'A database is required to run a Timestream custom query'
        assert table_name, 'A table is required to run a Timestream custom query'
        assert sql or file_path, 'Either a SQL query or a SQL file path is required to run a Timestream custom query'
        try:
            self.timestream_reader.does_table_exist(
                database_name=database_name,
                table_name=table_name,
                raise_err=True
            )
            return self.timestream_reader.custom_read(
                sql=sql,
                file_path=file_path
            )
        except Exception as e:
            self.logger.error(f'Exception occurred when running Timestream custom query: {e}')
            raise e
