"""Athena utils."""

import time

import boto3

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

athena_client = boto3.client("athena", region_name=config.AWS_DEFAULT_REGION)


def _execute_query(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.get_temp_path()
    if replacements:
        query = query.format(**replacements)
    result = athena_client.start_query_execution(
        QueryString=query,
        ResultConfiguration={"OutputLocation": output_s3_path},
        QueryExecutionContext={"Database": config.ATHENA_DATABASE}
    )
    return result["QueryExecutionId"]


def _wait_query(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 = 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 failed or cancelled: '{}'".format(
                result["QueryExecution"]["Status"].get("StateChangeReason", query_status))
            )
        time.sleep(sleep_interval)


def _execute_and_wait(
    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 = _execute_query(query, output_s3_path, replacements)
    _wait_query(query_id, sleep_interval)


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

    """
    _execute_and_wait(
        athena_consts.AthenaSQL.CREATE_DATABASE, replacements=dict(database=config.ATHENA_DATABASE))


def create_table(table_type: str):
    """Create Athena table.

    Args:
        table_type: Create table type.

    """
    table_s3_path = s3.get_full_path(table_type, consts.Folder.UPLOAD)
    _execute_and_wait(
        athena_consts.AthenaSQL.CREATE_TABLE,
        replacements=dict(
            s3_path=table_s3_path,
            database=config.ATHENA_DATABASE,
            table=athena_consts.TABLE_HISTORY,
            columns=athena_consts.AthenaSQL.CREATE_TABLE_COLUMNS[table_type],
        )
    )


def drop_table():
    """Drop Athena table.

    """
    _execute_and_wait(
        athena_consts.AthenaSQL.DROP_TABLE,
        replacements=dict(database=config.ATHENA_DATABASE, table=athena_consts.TABLE_HISTORY),
    )


def execute_aggregate_query(table_type: str):
    """Execute Athena aggregating query.

    Args:
        table_type: Aggregate table type.

    """
    output_s3_path = s3.get_full_path(table_type, consts.Folder.DOWNLOAD)
    _execute_and_wait(
        athena_consts.AthenaSQL.AGGREGATE[table_type],
        output_s3_path,
        sleep_interval=3,
        replacements=dict(table=athena_consts.TABLE_HISTORY)
    )