"""MR-Snowflake tasks."""

import datetime

from garcon import task
from garcon_contrib.dynamo_feed_status import garcon_feed_status

from flows import config
from flows.snowflake_executor import SnowflakeSQLExecutorMR


@task.decorate(timeout=900)
def bootstrap(activity, date):
    """Bootstrap the MR-Snowflake flow.

    Args:
        activity (ActivityWorker): the activity worker.
        date (str): import date (YYYY-MM-DD)
    """
    # if no date is passed from context, take yesterdays date.
    run_date = date or (
        datetime.date.today() - datetime.timedelta(days=1)).strftime(
        '%Y-%m-%d')
    activity.logger.info('Starting MR-Snowflake flow, date=%s', run_date)
    return dict(
        date=run_date
    )


# TODO: replace with a real snowflake import task
@task.decorate(timeout=1800)
def test_snowflake(activity, date):
    """Test activity.

    Args:
        activity (ActivityWorker): the activity worker.
        date (str): import date (YYYY-MM-DD)
    """
    print(date)
    activity.logger.info('Activity test executing. Start date was: ' + date)


@task.decorate(timeout=1800)
def drop_create_temp_table(activity):
    """Drop and create a temporary table.

    Args:
        activity (ActivityWorker): the activity worker.
    """
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.drop_tmp_table()
        sf_executor.create_tmp_table()
    activity.logger.info('Temporary table created.')


@task.decorate(timeout=1800)
def import_from_s3(activity, date):
    """Import the date from S3 bucket into a temporary table.

    Args:
        activity (ActivityWorker): the activity worker.
        date (str): import date (YYYY-MM-DD)
    """
    s3_path = 's3://{bucket_name}/{directory_name}'.format(
        bucket_name=config.S3_BUCKET_PATH,
        directory_name=date
    )
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.create_import_stage(
            s3_path, config.SF_IMPORT_STAGE_NAME)
        sf_executor.import_from_s3()
    activity.logger.info(
        'MR data loaded from S3 bucket into a temporary table.')


@task.decorate(timeout=1800)
def merge_into_mr_table(activity):
    """Drop and create a temporary table.

    Args:
        activity (ActivityWorker): the activity worker.
    """
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.merge_into_mr_table()
    activity.logger.info('MR data merged into the destination table.')


@task.decorate(timeout=1800)
def create_flattened_mr_table(activity):
    """Create a table with flattened data from MR Snowflake table.

    Args:
        activity (ActivityWorker): the activity worker.
    """
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.create_flattened_mr_table()
    activity.logger.info('Flattened registry table created.')


@task.decorate(timeout=1800)
def fill_flattened_mr_table(activity):
    """Fill the table with flattened data from MR Snowflake table.

    Truncate it before filling the data.

    Args:
        activity (ActivityWorker): the activity worker.
    """
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.truncate_table(config.SF_MR_FLATTENED_TABLE_NAME)
        sf_executor.fill_flattened_mr_table()
    activity.logger.info('Flattened registry table filled.')


@task.decorate(timeout=1800)
def fill_flattened_mr_locked_table(activity):
    """Fill the table with flattened data for locked territories.

    Truncate it before filling the data.

    Args:
        activity (ActivityWorker): the activity worker.
    """
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.truncate_table(config.SF_MR_LOCKED_FLATTENED_TABLE_NAME)
        sf_executor.fill_flattened_mr_locked_table()
    activity.logger.info('Flattened locked registry table filled.')


@task.decorate(timeout=1800)
def import_audit_from_s3(activity, date):
    """Import the audit data from S3 bucket into a temporary table.

    Args:
        activity (ActivityWorker): the activity worker.
        date (str): import date (YYYY-MM-DD)
    """
    s3_path = 's3://{bucket_name}/{directory_name}'.format(
        bucket_name=config.S3_AUDIT_BUCKET_PATH,
        directory_name=date
    )
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.create_import_stage(
            s3_path, config.SF_AUDIT_IMPORT_STAGE_NAME)
        sf_executor.import_audit_from_s3()
    activity.logger.info(
        'MR_AUDIT data loaded from S3 bucket into Snowflake.')


@task.decorate(timeout=1800)
def fill_flattened_mr_audit_table(activity):
    """Fill the table with flattened data for locked territories.

    Truncate it before filling the data.

    Args:
        activity (ActivityWorker): the activity worker.
    """
    with SnowflakeSQLExecutorMR() as sf_executor:
        sf_executor.truncate_table(config.SF_MR_AUDIT_FLATTENED_TABLE_NAME)
        sf_executor.fill_flattened_mr_audit_table()
    activity.logger.info('Flattened audit registry table filled.')


@task.decorate(timeout=1800)
def update_dynamo_status(activity, date):
    """Update the final flow status in DynamoDB.

    Args:
        activity (ActivityWorker): the activity worker.
        date (str): import date
    """
    garcon_feed_status.set_overall_status(
        config.SWF_NAME, date, garcon_feed_status.STATUS_INGESTED)
    activity.logger.info('Final status updated in DynamoDB.')