"""Athena utils."""

import time

import boto3

import config
import constants.athena as athena_consts
import constants.common as consts
import s3


class AthenaClient:
    """Athena logic container.

    """

    def __init__(self):
        """Init client.

        """
        self.boto_session = boto3.Session()
        self.athena_client = self.boto_session.client('athena')

    def _execute_query(self, query: str, output_s3_path: str or None = None, replacements: dict or None = None) -> str:
        """Execute Athena query.

        Args:
            query (str): Athena query.
            output_s3_path (str): Output S3 path.
            replacements (dict or None): Query formatting.

        Returns:
            str: Query ID.

        """
        if not output_s3_path:
            output_s3_path = s3.S3Client.get_temp_path()
        if replacements:
            query = query.format(**replacements)
        result = self.athena_client.start_query_execution(
            QueryString=query,
            ResultConfiguration={'OutputLocation': output_s3_path},
            QueryExecutionContext={'Database': config.ATHENA_DATABASE}
        )
        return result['QueryExecutionId']

    def _wait_query(self, query_id: str, sleep_interval: int = 1):
        """Wait Athena query to finish.

        Args:
            query_id (str): Query ID.
            sleep_interval (int): Check state interval in seconds.

        """
        query_status = None
        while query_status == 'QUEUED' or query_status == 'RUNNING' or query_status is None:
            result = self.athena_client.get_query_execution(QueryExecutionId=query_id)
            query_status = result['QueryExecution']['Status']['State']
            if query_status == 'FAILED' or query_status == 'CANCELLED':
                raise Exception('Athena query with the string "{}" failed or was cancelled'.format(query_status))
            time.sleep(sleep_interval)

    def _execute_and_wait(
            self, query: str, output_s3_path: str or None = None,
            replacements: dict or None = None, sleep_interval: int = 1):
        """Execute and wait Athena query.

        Args:
            query (str): Athena query.
            output_s3_path (str): Output S3 path.
            replacements (dict or None): Query formatting.
            sleep_interval (int): Check state interval in seconds.

        """
        query_id = self._execute_query(query, output_s3_path, replacements)
        self._wait_query(query_id, sleep_interval)

    def create_database(self):
        """Create Athena database.

        """
        self._execute_and_wait(
            athena_consts.SQL_ATHENA_CREATE_DATABASE, replacements=dict(database=config.ATHENA_DATABASE))

    def drop_database(self):
        """Drop Athena database.

        """
        self._execute_and_wait(
            athena_consts.SQL_ATHENA_DROP_DATABASE, replacements=dict(database=config.ATHENA_DATABASE))

    def create_table(self, vendor: str, playlist_id: str):
        """Create Athena table.

        Args:
            vendor (str): Vendor name.
            playlist_id (str): Playlist ID.

        """
        table_s3_path = s3.S3Client.get_full_path(vendor, playlist_id, consts.FOLDER_UPLOAD)
        self._execute_and_wait(
            athena_consts.SQL_ATHENA_CREATE_TABLE[vendor],
            replacements=dict(s3_path=table_s3_path, database=config.ATHENA_DATABASE)
        )

    def drop_table(self, vendor: str):
        """Drop Athena table.

        Args:
            vendor (str): Vendor name.

        """
        self._execute_and_wait(
            athena_consts.SQL_ATHENA_DROP_TABLE,
            replacements=dict(database=config.ATHENA_DATABASE, table=athena_consts.TABLE_HISTORY[vendor]),
        )

    def update_table(self, vendor: str, playlist_id: str):
        """Update Athena table source.

        Args:
            vendor (str): Vendor name.
            playlist_id (str): Playlist ID.

        """
        table_s3_path = s3.S3Client.get_full_path(vendor, playlist_id, consts.FOLDER_UPLOAD)
        self._execute_and_wait(
            athena_consts.SQL_ATHENA_UPDATE_TABLE,
            replacements=dict(s3_path=table_s3_path, table=athena_consts.TABLE_HISTORY[vendor])
        )

    def execute_aggregate_query(self, vendor: str, playlist_id: str):
        """Execute Athena aggregating query.

        Args:
            vendor (str): Vendor name.
            playlist_id (str): Playlist ID.

        """
        output_s3_path = s3.S3Client.get_full_path(vendor, playlist_id, consts.FOLDER_DOWNLOAD)
        self._execute_and_wait(athena_consts.SQL_ATHENA_AGG[vendor], output_s3_path, sleep_interval=3)