import boto3
import csv
import logging
import time

from datetime import datetime
from psi_athena import DEFAULT_LOGGER, log_it
from typing import List, Union

DEFAULT_ATHENA_DATA_CATALOG = 'AwsDataCatalog'
DEFAULT_ATHENA_DATABASE = 'default'
DEFAULT_ATHENA_WORKGROUP = 'primary'
DEFAULT_AWS_REGION = 'us-east-1'
DEFAULT_AWS_SESSION_NAME = 'psi-athena-session'
DEFAULT_MAX_POLL_DURATION_SECONDS = 300
DEFAULT_POLL_WAIT_SECONDS = 5


class Athena:

    def __init__(
        self,
        data_catalog: str = None,
        workgroup: str = None,
        database: str = None,
        max_poll_duration_seconds: int = None,
        poll_wait_seconds: int = None,
        aws_session: boto3.Session = None,
        logger: logging.Logger = None
    ):
        self.data_catalog = data_catalog or DEFAULT_ATHENA_DATA_CATALOG
        self.workgroup = workgroup or DEFAULT_ATHENA_WORKGROUP
        self.database = database or DEFAULT_ATHENA_DATABASE
        self.max_poll_duration_seconds = max_poll_duration_seconds or DEFAULT_MAX_POLL_DURATION_SECONDS
        self.poll_wait_seconds = poll_wait_seconds or DEFAULT_POLL_WAIT_SECONDS
        self.__aws_session = aws_session or boto3.Session()
        self.__athena_client = None
        self.__logger = logger or DEFAULT_LOGGER
        self.assume_role()

    def assume_role(self, role_arn: str = None, region: str = None, session_name: str = None):
        if role_arn:
            region = region or DEFAULT_AWS_REGION
            session_name = session_name or DEFAULT_AWS_SESSION_NAME
            self.__logger.info(f'Assuming role: {role_arn} in region {region}')
            stsresponse = self.__aws_session.assume_role(
                RoleArn=role_arn,
                RoleSessionName=session_name
            )

            self.__logger.debug(f'Initializing Athena client from role: {role_arn}')
            self.__athena_client = self.__aws_session.client(
                'athena',
                region_name=region,
                aws_access_key_id=stsresponse["Credentials"]["AccessKeyId"],
                aws_secret_access_key=stsresponse["Credentials"]["SecretAccessKey"],
                aws_session_token=stsresponse["Credentials"]["SessionToken"]
            )
        else:
            self.__logger.debug('Initializing Athena client from existing AWS session')
            self.__athena_client = self.__aws_session.client('athena')

    @log_it
    def run_athena_query(
        self,
        table: str,
        query: str,
        download_results: bool = True,
        reload_partitions: bool = False
    ) -> Union[dict, list]:
        assert table and isinstance(table, str), f'Invalid table name: {table}'
        assert query and isinstance(query, str), f'Invalid Athena query sql: {query}'
        if reload_partitions:
            self.reload_partitions(table)
        query_execution_id = self.start_query(query)
        query_results = self.poll_for_query_results(query_execution_id)
        if download_results:
            return self.download_query_results_as_dict(query_results)
        return query_results

    def start_query(
        self,
        query: str,
        data_catalog: str = None,
        workgroup: str = None,
        database: str = None
    ) -> str:
        data_catalog = data_catalog or self.data_catalog
        workgroup = workgroup or self.workgroup
        database = database or self.database
        assert query and isinstance(query, str), f'Invalid Athena query: {query}'
        assert data_catalog and isinstance(data_catalog, str), f'Invalid Athena data catalog: {data_catalog}'
        assert workgroup and isinstance(workgroup, str), f'Invalid Athena workgroup: {workgroup}'
        assert database and isinstance(database, str), f'Invalid Athena database: {database}'
        self.__logger.info(f'Starting Athena query: {query}')
        try:
            query_execution_id = self.__athena_client.start_query_execution(
                QueryString=query,
                QueryExecutionContext={
                    'Database': database,
                    'Catalog': data_catalog
                },
                WorkGroup=workgroup
            )['QueryExecutionId']
        except Exception as e:
            self.__logger.error(f'Failed to start Athena query execution: {e}')
            raise e
        self.__logger.info(f'Athena query execution id: {query_execution_id}')
        return query_execution_id

    def poll_for_query_results(self, query_execution_id: str) -> dict:
        assert query_execution_id and isinstance(query_execution_id, str), f'Invalid query_execution_id to poll: {query_execution_id}'
        t0 = datetime.now()
        query_execution_status = 'QUEUED'
        while (datetime.now() - t0).total_seconds() < self.max_poll_duration_seconds:
            try:
                query_results = self.__athena_client.get_query_execution(
                    QueryExecutionId=query_execution_id
                )
                query_execution_status = query_results['QueryExecution']['Status']['State']
            except Exception as e:
                self.__logger.error(f'Failed to retrieve Athena query execution status: {e}')
                raise e

            if query_execution_status in ['SUCCEEDED', 'FAILED', 'CANCELLED']:
                break
            time.sleep(self.poll_wait_seconds)
        assert query_execution_status == 'SUCCEEDED', f'Athena query execution did not succeed, polling exited with status: {query_execution_status}'
        return query_results

    def download_query_results_as_dict(self, query_results: dict) -> List[dict]:
        data = self.download_query_results(query_results)
        try:
            csvstr = data.decode('utf-8')
            parsed = list(csv.DictReader(csvstr.split('\n')))
            for i in range(len(parsed)):
                for k, v in parsed[i].items():
                    assert k and isinstance(k, str), f'Invalid CSV column (likely that extra columns found): {k}'
                    if len(v) == 0:
                        # Parse empty strings as None
                        parsed[i][k] = None
                    elif '.' in v:
                        # Cast numbers with radix point as floats
                        try:
                            float_val = float(v)
                            parsed[i][k] = float_val
                        except ValueError:
                            pass
                    else:
                        # Cast numbers as ints
                        try:
                            int_val = int(v)
                            parsed[i][k] = int_val
                        except ValueError:
                            pass
            return parsed
        except Exception as e:
            self.__logger.error(f'Failed to parse query results into a dict: {e}')
            raise e

    def download_query_results(self, query_results: dict) -> bytes:
        assert query_results and isinstance(query_results, dict), f'Invalid query_results to download results for: {query_results}'
        s3_path = query_results.get('QueryExecution', {}).get('ResultConfiguration', {}).get('OutputLocation')
        try:
            client = self.__aws_session.client('s3')
            self.__logger.debug(f'Fetching object from S3: {s3_path}')
            obj = client.get_object(
                Bucket=s3_path[len('s3://'):].split('/', 1)[0],
                Key=s3_path[len('s3://'):].split('/', 1)[1]
            )
            return obj['Body'].read()
        except Exception as e:
            self.__logger.error(f'Failed to download results with error: {e}')
            raise e

    def reload_partitions(self, table: str):
        assert table and isinstance(table, str), f'Invalid table name: {table}'
        query = f'MSCK REPAIR TABLE `{table}`'
        query_result = self.run_athena_query(
            table,
            query,
            download_results=False,
            reload_partitions=False
        )
        assert query_result['QueryExecution']['Status']['State'] == 'SUCCEEDED'
