"""Digital ETL Tasks."""

from os.path import dirname

from botocore.exceptions import ClientError
from garcon import task
from garcon_contrib.dynamo_feed_status import garcon_feed_status

from flows import datastore
from flows import datawarehouse
from flows import util
from flows.digital import config
from flows.digital import log
from flows.digital import queries
from flows.digital import status
from flows.digital import util as etl_util


@task.decorate(timeout=14400)
def unload_from_datawarehouse(
        activity, bucket, correlation_id, date_end, date_start, upcs):
    """Unload the data from datawarehouse to s3.

    Args:
        activity (ActivityWorker): activity worker.
        bucket (str): s3 destination bucket name.
        correlation_id (str): etl correlation ID.
        date_end (str): YYYY-MM-DD exclusive date end range for unload select.
        date_start (str): YYYY-MM-DD inclusive date start range for unload
            select.
        upcs (list): upcs for unload query.

    Returns:
        dict: unload batch count.
    """
    sql_templates = queries.get_unload_from_snowflake_sql(
        bucket, correlation_id, date_end, date_start, upcs)
    for batch, sql_template in enumerate(sql_templates):
        sql = sql_template.format(batch=batch)
        datawarehouse.execute(sql)

    unload_path_format = config.UNLOAD_DESTINATION.format(
        bucket=bucket, correlation_id=correlation_id)
    unload_path = dirname(unload_path_format) + '/'

    log.update_status(correlation_id, status.UNLOADED)

    return {'batch_count': batch + 1, 'unload_path': unload_path}


@task.decorate(timeout=120)
def create_temp_table(activity, correlation_id, create, temp_table_name):
    """Create the temporary table in the datastore.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        create (str): create temp table query.
        temp_table_name (str): temp table name format.

    Returns:
        dict: temp table name.
    """
    table_name = temp_table_name.format(
        correlation_hex=util.correlation_id_hex(correlation_id))
    create_sql = create.format(table_name=table_name)
    datastore.execute(create_sql)
    log.update_status(correlation_id, status.TEMP_TABLE_CREATED)
    return {'table_name': table_name}


@task.decorate(timeout=3600)
def insert_to_temp_table(
        activity, bucket, correlation_id, insert, batch_count,
        temp_table_name):
    """Load the data from s3 to a temp table in the datastore.

    Args:
        activity (ActivityWorker): activity worker.
        bucket (str): s3 destination bucket name.
        correlation_id (str): etl correlation ID.
        insert (str): insert query.
        batch_count (int): number of unload batches as source data.
        temp_table_name (str): temp table name.
    """
    sql = insert.format(table_name=temp_table_name)

    s3_location_format = config.UNLOAD_DESTINATION.format(
        bucket=bucket, correlation_id=correlation_id)

    with datastore.context() as (cursor, connection):
        for batch_number in range(batch_count):
            try:
                s3_path = s3_location_format.format(batch=batch_number)
                rows = etl_util.read_csv_from_s3(s3_path)
                cursor.executemany(sql, rows)
            except ClientError as err:
                # If unload this batch file is not found, it just had no data,
                # and therefore not on s3 (NoSuchKey error). Just move on and
                # try the next batch.
                if err.response['Error']['Code'] != 'NoSuchKey':
                    raise err

        log.update_status(correlation_id, status.TEMP_TABLE_INSERTED)


@task.decorate(timeout=120)
def update_etl_status(activity, correlation_id, status):
    """Update the ETL status log.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        status (str): status attribute to update to.
    """
    log.update_status(correlation_id, status)


@task.decorate(timeout=3600)
def move_temp_table_to_digital_revenue(
        activity, correlation_id, date_end, date_start, temp_table_name, upcs):
    """Insert all rows from the temp table to digital revenue.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        date_end (str): YYYY-MM-DD exclusive date end range for unload select.
        date_start (str): YYYY-MM-DD inclusive date start range for unload
            select.
        temp_table_name (str): temp table name.
        upcs (list): UPCs of existing data to delete.
    """
    drop_sql = queries.DROP_TEMP_TABLE.format(table_name=temp_table_name)
    insert_sql = queries.INSERT_FROM_TEMP_TABLE.format(
        table_name=temp_table_name)
    delete_sqls = queries.get_delete_from_digital_revenue_sql(upcs)
    with datastore.context() as (cursor, connection):
        cursor.execute('START TRANSACTION')
        for delete_sql in delete_sqls:
            cursor.execute(
                delete_sql, {'date_end': date_end, 'date_start': date_start})
        cursor.execute(insert_sql)
        cursor.execute(drop_sql)
    log.update_status(correlation_id, status.TEMP_TABLE_MADE_LIVE)


@task.decorate(timeout=120)
def send_sns(activity, correlation_id, date_end, date_start, upcs):
    """Send all listeners a message that the ETL has completed.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        date_end (str): YYYY-MM-DD exclusive date end range for unload select.
        date_start (str): YYYY-MM-DD inclusive date start range for unload
            select.
        upcs (list): upcs for unload query.
    """
    sns_correlation_id = correlation_id + '.1'
    etl_util.send_success_notification(
        sns_correlation_id, date_end, date_start, upcs)


@task.decorate(timeout=120)
def queue_build_cache(activity, correlation_id, date_end, date_start, upcs):
    """Enqueue a cache building job message.

    Args:
        activity (ActivityWorker): activity worker.
        correlation_id (str): etl correlation ID.
        date_end (str): YYYY-MM-DD exclusive date end range for unload select.
        date_start (str): YYYY-MM-DD inclusive date start range for unload
            select.
        upcs (list): upcs for unload query.
    """
    sqs_correlation_id = correlation_id + '.1'
    etl_util.queue_build_cache(
        sqs_correlation_id, date_end, date_start, upcs)


@task.decorate(timeout=60)
def set_dynamo_status(activity, correlation_id, date):
    """Set DynamoDB statuses for SWF execution.

    Args:
        activity (Activity): SWF Activity.
        correlation_id (str): correlation id.
        date (str): YYYY-MM-DD start date of ETL range.
    """
    garcon_feed_status.set_overall_status(
        config.SWF_WORKFLOW_NAME, date, garcon_feed_status.STATUS_INGESTED)
    log.update_status(correlation_id, status.DYNAMO_STATUS_UPDATED)