"""Tasks for handling YouTube specific operations."""
import logging
import time
import uuid

import boto3

from feed_ingestion.conf.config import BOTO3_CONFIG
from feed_ingestion.util.aws.assume_role import assumed_session

logger = logging.getLogger(__name__)

DEFAULT_EXECUTION_TIMEOUT_SECONDS = 30 * 60
SLEEP_TIME_SECONDS = 10

PARQUET_FILE_PATTERN = r'\d{8}_\d{6}_\d{5}_\w+_\w+\-\w+\-\w+\-\w+\-\w+'

STATE_SUCCEEDED = 'SUCCEEDED'
STATE_FAILED = 'FAILED'
STATE_CANCELLED = 'CANCELLED'
FINAL_STATUSES = [STATE_SUCCEEDED, STATE_FAILED, STATE_CANCELLED]

QUERY_CREATE_TEMP_TABLE = """
CREATE TABLE {athena_temp_db}.{athena_temp_table}
WITH (
    format='PARQUET',
    parquet_compression='GZIP',
    external_location = 's3://{target_s3_bucket}/{target_s3_key}'
) AS
{athena_query}
"""

QUERY_DROP_TEMP_TABLE = """
DROP TABLE IF EXISTS {athena_temp_db}.{athena_temp_table}
"""


QUERY_UNLOAD_TO_S3 = """
UNLOAD (
{athena_query}
)
TO 's3://{target_s3_bucket}/{target_s3_key}'
WITH (
  format = 'PARQUET',
  compression = 'GZIP'
);
"""


class StatusError(RuntimeError):
    """Status error exception class."""

    pass


def run_query(
        athena_query,
        athena_temp_database,
        athena_workgroup,
        destination_s3_bucket, destination_s3_path,
        timeout=DEFAULT_EXECUTION_TIMEOUT_SECONDS,
        parameters=None,
        use_unload_query=False,
        assume_role=None):
    """Extract Athena query as PARQUET files in S3 location.

    For new code it is suggested to use UNLOAD query by use_unload_query=True.
    Unload method doesn't create temporary Athena table.
    It unloads to S3 directly.

    When use_unload_query=False,
    it creates and then drops temporary Athena table.

    Provided destination_s3_path should NOT exist.

    It raises StatusError when query timeout or not success.

    Args:
        athena_query: (str) SQL query to execute
        athena_temp_database (str): name of writable database in Athena
        athena_workgroup (str): Athena Workgroup
        destination_s3_bucket (str): bucket where to put the result PARQUET
        destination_s3_path (str): path in destination_s3_bucket
            (should end with "/") Provided path should NOT exist.
        use_unload_query (bool): optional modern way to extract data.
        parameters (list): optional parameters to replace "?" in query.
        assume_role (str): optional IAM role ARN to assume for Athena.
    """
    if not destination_s3_path.endswith('/'):
        raise ValueError('destination_s3_path should ends with /')

    def _wait_athena_query_final_status(execution_id):
        started_at = time.time()
        current_time = started_at
        deadline = started_at + timeout
        while current_time <= deadline:
            response = athena_client.get_query_execution(
                QueryExecutionId=execution_id
            )
            status_dict = response['QueryExecution']['Status']
            state = status_dict['State']
            change_reason = status_dict.get('StateChangeReason')
            if state in FINAL_STATUSES:
                break
            time.sleep(SLEEP_TIME_SECONDS)
            current_time = time.time()
        else:
            raise TimeoutError(f'Query timed out after {timeout} seconds')

        if state != STATE_SUCCEEDED:
            message = f'Error final status "{state}" ' \
                      f'for Athena executionId="{execution_id}"' \
                      f'ChangeReason: {change_reason}'
            raise StatusError(message)

    def _execute_athena_query(sql_query, parameters=None):
        # For SQL execution parameters to be treated as strings,
        # they must be enclosed in single quotes
        # https://docs.aws.amazon.com/athena/latest/ug/querying-with-prepared-statements.html
        # Note: sql_templater does it automatically
        if parameters:
            kwargs = {
                'ExecutionParameters': parameters,
            }
        else:
            kwargs = {}
        logger.info(f'Execution query: {athena_query}')
        logger.info(f'Parameters: {parameters}')
        response = athena_client.start_query_execution(
            QueryString=sql_query,
            ClientRequestToken=str(uuid.uuid4()),
            ResultConfiguration={
                'OutputLocation': athena_output_location,
            },
            WorkGroup=athena_workgroup,
            **kwargs
        )
        execution_id = response['QueryExecutionId']
        return execution_id

    if assume_role:
        session = assumed_session(assume_role)
    else:
        session = boto3.Session()
    athena_client = session.client('athena', config=BOTO3_CONFIG)

    athena_output_location = \
        f's3://{destination_s3_bucket}/{destination_s3_path}'
    temp_table = f'temp_athena_{uuid.uuid4()}'.replace('-', '_')
    sql_params = {
        'athena_temp_db': athena_temp_database,
        'athena_temp_table': temp_table,
        'target_s3_bucket': destination_s3_bucket,
        'target_s3_key': destination_s3_path,
        'athena_query': athena_query,
    }

    if use_unload_query:
        sql_query = QUERY_UNLOAD_TO_S3.format(
            **sql_params
        )
        execution_id = _execute_athena_query(
            sql_query=sql_query,
            parameters=parameters)
        _wait_athena_query_final_status(execution_id)
    else:
        try:
            sql_query = QUERY_CREATE_TEMP_TABLE.format(
                **sql_params
            )
            execution_id = _execute_athena_query(
                sql_query=sql_query,
                parameters=parameters)
            _wait_athena_query_final_status(execution_id)
        finally:
            sql_query = QUERY_DROP_TEMP_TABLE.format(
                **sql_params
            )
            execution_id = _execute_athena_query(
                sql_query=sql_query)
            _wait_athena_query_final_status(execution_id)