"""Athena utility functions."""
import datetime
import time

import boto3
import botocore
from lambdacommon.common_config import logger

import config


class AthenaUtil:
    """Athena utility functions."""

    def __init__(self):
        """Initialize Athena client."""
        sts_client = boto3.client('sts')
        assumed_role = sts_client.assume_role(
            RoleArn=config.SHARED_AWS_ROLE_ARN,
            RoleSessionName=config.ROLE_SESSION_NAME
        )
        ACCESS_KEY = assumed_role['Credentials']['AccessKeyId']
        SECRET_KEY = assumed_role['Credentials']['SecretAccessKey']
        SESSION_TOKEN = assumed_role['Credentials']['SessionToken']
        self.athena_client = boto3.client(
            'athena',
            region_name=config.AWS_REGION,
            aws_access_key_id=ACCESS_KEY,
            aws_secret_access_key=SECRET_KEY,
            aws_session_token=SESSION_TOKEN,
        )

    def get_named_query_statements(self, account_id):
        """
        Get named query statements for execution.

        Args:
            account_id (str): The AWS Account ID

        Returns:
            str: SQL statements of named query
        """
        named_queries = self.athena_client.list_named_queries(
            WorkGroup=config.ATHENA_WORKGROUP,
        )
        named_query_name = f'{config.ATHENA_NAME_PREFIX}_{account_id}_break_glass' \
            if not config.ATHENA_NAMED_QUERY else config.ATHENA_NAMED_QUERY  # noqa: E501
        logger.info(f'Named queries for workgroup are {named_queries}')
        for query_id in named_queries['NamedQueryIds']:
            query = self.athena_client.get_named_query(
                NamedQueryId=query_id,
            )
            if query['NamedQuery']['Name'] == named_query_name:
                logger.info(f'Query ID {query_id} matches '
                            f'{named_query_name}')
                return query['NamedQuery']['QueryString']
        logger.exception(f'Query ID not found for {named_query_name}')
        return None

    def get_query_execution_details(self, query_execution_id):
        """
        Get query execution details, including status and result configuration.

        Args:
            query_execution_id (str): query execution ID

        Returns:
            str: S3 output location of query result file
        """
        try:
            polling_timeout = time.time() + config.ATHENA_POLLING_TIMEOUT
            while time.time() < polling_timeout:
                response = self.athena_client.get_query_execution(
                    QueryExecutionId=query_execution_id)
                execution_state = response['QueryExecution']['Status']['State']
                logger.info(f'Query {query_execution_id} state: {execution_state}')  # noqa: E501
                if execution_state in ['QUEUED', 'RUNNING']:
                    time.sleep(5)
                elif execution_state in ['FAILED', 'CANCELLED']:
                    logger.exception(
                        'Error occurred during Athena query execution: {}'.format(  # noqa: E501
                            response['QueryExecution']['Status']
                            ['AthenaError']['ErrorMessage']))
                    raise SystemExit('Query execution has failed for ID '
                                     f'{query_execution_id}')
                elif execution_state == 'SUCCEEDED':
                    output_location = response['QueryExecution'][
                        'ResultConfiguration']['OutputLocation']
                    return output_location
                else:
                    raise SystemExit('Unexpected execution state')
            raise SystemExit('Timeout exceeded, and Athena query execution '
                             'has not finished. Check details '
                             f'for execution ID {query_execution_id}')
        except botocore.exceptions.ClientError as error:
            logger.exception('Error getting Athena query execution: {}'.format(
                error.response['Error']['Message']))
            raise error

    def start_audit_query_execution(self, user, date, query_statement):
        """
        Start query execution that generates break-glass audit log.

        Relies on workgroup configuration for result/output configuration,
        rather than being specified in the client settings,
        which would be overridden.

        Args:
            user (str): the user that broke glass
            date (str): the date, in YYYY/MM/DD format
            query_statement (str): query statements to run

        Returns:
            str: query execution ID

        Generate the formatted time window in which to query Athena
        in order to only scan the date partitions in question.
        Dates must be quoted in order to use comparison operators.
        """
        message_date = datetime.datetime.strptime(date, '%Y/%m/%d')
        day_before = (message_date - datetime.timedelta(
            days=1)).strftime('%Y/%m/%d')
        day_after = (message_date + datetime.timedelta(
            days=1)).strftime('%Y/%m/%d')

        try:
            response = self.athena_client.start_query_execution(
                QueryString=query_statement,
                QueryExecutionContext={
                    'Database': config.ATHENA_DATABASE,
                    'Catalog': config.ATHENA_CATALOG,
                },
                WorkGroup=config.ATHENA_WORKGROUP,
                ExecutionParameters=[
                    f"'{day_before}'",
                    f"'{day_after}'",
                    config.BREAK_GLASS_ROLE_NAME,
                    f'%{user}',
                ]
            )
        except botocore.exceptions.ClientError as error:
            logger.exception('Error starting Athena execution: {}'.format(
                error.response['Error']['Message']))
            raise error

        return response['QueryExecutionId']