"""
Generic activities to load data to staging_raw table.

Used by FlowLoadRawMixinSF mixin.
"""

from garcon import task

from feed_ingestion.conf.config import merge_configs
from feed_ingestion.flows import registered_executors
from feed_ingestion.flows.apple_music_streams import config
from feed_ingestion.flows.helpers import get_sf_config
from feed_ingestion.util import task_status

TASK_ID = 'staging_raw_table_tasks'


@task.decorate(timeout=300)
def create_temp_staging_raw_table(
        activity, date, feed_name, temp_staging_raw_table, sfdb_params,
        secrets_path=None, kwargs=None):
    """Create temporary table in Snowflake.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed to get executor class.
        temp_staging_raw_table (str): Name of the temp staging_raw table.
        sfdb_params (dict): Dict with params to optionally override default
            ones (Snowflake db and schema name).
        secrets_path (str): Secrets manager path of the flow.
        kwargs (dict): Custom activity params.
    """
    kwargs = kwargs or {}
    has_contexts = False
    contexts = task_status.get_report_contexts(feed_name, date)
    report_name = kwargs.get('report_name', None)
    completed_report = True
    if contexts:
        has_contexts = True
        completed_report = (
            all(val[task_status.FIELD_CONTEXT_STATUS] ==
                task_status.CONTEXT_STATUS_PROCESSED
                for val in contexts.values()))

    if (task_status.is_completed_task(
            feed_name, date, 'staging_raw_table_tasks') and
            (not has_contexts or has_contexts
             and completed_report
             # need to load common reports, because other reports depend on it
             and report_name not in config.common_reports)):
        activity.logger.info(
            'create_temp_staging_raw_table for {feed_name} {date} '
            'already complete, and '
            '"reload" flag was not passed, skipping...'
            .format(feed_name=feed_name, date=date))
        return

    sf_config = get_sf_config(secrets_path)
    sf_config_custom = merge_configs(sf_config, sfdb_params)
    ExecutorSR = registered_executors.get(feed_name)
    with ExecutorSR(sf_config_custom) as sf_executor:
        sf_executor.drop_table(temp_staging_raw_table)
        activity.logger.info('{table} was dropped, if existed'.format(
            table=temp_staging_raw_table))
        sf_executor.create_temp_staging_raw_table(
            temp_staging_raw_table, **kwargs)
    activity.logger.info(
        '{table} was created'.format(table=temp_staging_raw_table))


@task.decorate(timeout=12600)
def load_temp_staging_raw_table(
        activity, date, feed_name, aws, key_dir, temp_staging_raw_table,
        sfdb_params, secrets_path=None, kwargs=None):
    """Load staged data from bucket to the temp datestamped table in Snowflake.

    Args:
        activity (ActivityWorker): The activity worker.
        aws (dict): AWS credentials to parametrize COPY INTO statement.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed to get executor class.
        key_dir (str): S3 directory to load files from.
        temp_staging_raw_table (str): Name of the temp staging_raw table.
        sfdb_params (dict): Dict with params to optionally override default
            ones (Snowflake db and schema name).
        secrets_path (str): Secrets manager path of the flow.
    """
    if task_status.is_completed_task(feed_name, date, TASK_ID):
        activity.logger.info(
            'load_temp_staging_raw_table for {date} already complete, and '
            '"reload" flag was not passed, skipping...'.format(date=date))
        return

    kwargs = kwargs or {}

    sf_config = get_sf_config(secrets_path)
    sf_config_custom = merge_configs(sf_config, sfdb_params)
    ExecutorSR = registered_executors.get(feed_name)

    with ExecutorSR(sf_config_custom) as sf_executor:
        result = sf_executor.load_temp_staging_raw_table(
            temp_staging_raw_table, aws, key_dir, **kwargs)

        if result and result.get('info_message'):
            activity.logger.info(result['info_message'])


@task.decorate(timeout=7200)
def load_staging_raw_table(
        activity, date, feed_name, temp_staging_raw_table, staging_raw_table,
        sfdb_params, clean=None, set_complete=None,
        secrets_path=None, kwargs=None):
    """Copy the temp_staging_raw data to the feed's staging_raw table.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed to get executor class.
        temp_staging_raw_table (str): Name of the temp staging_raw table.
        staging_raw_table (str): Name of the staging_raw table.
        sfdb_params (dict): Dict with params to optionally override default
            ones (Snowflake db and schema name).
        clean (str): If set to 'False' cleanup step will be skipped.
            It allows to append data instead of full refresh.
        set_complete (str): If set to 'False' task doesn't set completed
            status. It allows to load data using generator.
        secrets_path (str): Secrets manager path of the flow.
        kwargs (dict): Custom activity params.
    """
    if task_status.is_completed_task(feed_name, date, TASK_ID):
        activity.logger.info(
            'load_staging_raw_table for {date} already complete, and '
            '"reload" flag was not passed, skipping...'.format(date=date))
        return

    kwargs = kwargs or {}
    clean = clean or 'True'
    set_complete = set_complete or 'True'

    sf_config = get_sf_config(secrets_path)
    sf_config_custom = merge_configs(sf_config, sfdb_params)
    ExecutorSR = registered_executors.get(feed_name)
    with ExecutorSR(sf_config_custom) as sf_executor:
        if clean == 'True':
            sf_executor.clean_staging_raw_table(
                staging_raw_table, date, **kwargs)
            activity.logger.info(
                '{table} was cleaned from '
                'the rows with the date (or date range) '
                'of the current workflow run'.format(table=staging_raw_table))
        sf_executor.load_staging_raw_table(
            temp_staging_raw_table, staging_raw_table, date, **kwargs)

        if temp_staging_raw_table is None:
            activity.logger.info(
                'Temp staging table was not dropped '
                'because its name was not specified')
        else:
            sf_executor.drop_table(temp_staging_raw_table)
            activity.logger.info('{table} was dropped, if existed'.format(
                table=temp_staging_raw_table))

    # this is the only place where we mark staging_raw_table_tasks as done
    if set_complete == 'True':
        task_status.mark_completed_task(feed_name, date, TASK_ID)
        activity.logger.info(
            'All {task_id} of {feed_name} for {date} completed'.format(
                task_id=TASK_ID, feed_name=feed_name, date=date))


@task.decorate(timeout=600)
def drop_temp_staging_raw_table(
        activity, feed_name, temp_staging_raw_table,
        sfdb_params, secrets_path=None):
    """Drop temp staging raw table.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        temp_staging_raw_table (str): Table name to drop.
        sfdb_params (dict): Dict with params to optionally override default
            ones (Snowflake db and schema name).
        secrets_path (str): Secrets manager path of the flow.
    """
    sf_config = get_sf_config(secrets_path)
    sf_config_custom = merge_configs(sf_config, sfdb_params)
    ExecutorSR = registered_executors.get(feed_name)
    with ExecutorSR(sf_config_custom) as sf_executor:
        sf_executor.drop_table(temp_staging_raw_table)
    activity.logger.info('{} was dropped'.format(temp_staging_raw_table))


@task.decorate(timeout=2500)
def drop_remaining_temp_staging_raw_tables(
        activity, feed_name, temp_table_pattern,
        date, sfdb_params, secrets_path=None):
    """Drop temp staging raw table.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        temp_table_pattern (str): Table to drop name pattern.
        date (str): Date of the data being process (YYYY-MM-DD).
        sfdb_params (dict): Dict with params to optionally override default
            ones (Snowflake db and schema name).
        secrets_path (str): Secrets manager path of the flow.
    """
    sf_config = get_sf_config(secrets_path)
    sf_config_custom = merge_configs(sf_config, sfdb_params)
    ExecutorSR = registered_executors.get(feed_name)

    with ExecutorSR(sf_config_custom) as sf_executor:
        # table names are the second values
        tables = [t[1] for t in sf_executor.select_remaining_tables(
            temp_table_pattern, date) if 'temp_' in t[1].lower()]
        for table in tables:
            sf_executor.drop_table(table)
            activity.logger.info('{} table for {} was dropped'.format(
                table, date))


@task.decorate(timeout=300)
def clean_staging_raw_table(
        activity, date, feed_name, staging_raw_table,
        sfdb_params, secrets_path=None, kwargs=None):
    """Clean staging_raw_table for the particular date.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed to get executor class.
        staging_raw_table (str): Name of the staging_raw table.
        sfdb_params (dict): Dict with params to optionally override default
            ones (Snowflake db and schema name).
        secrets_path (str): Secrets manager path of the flow.
    """
    if task_status.is_completed_task(feed_name, date, TASK_ID):
        activity.logger.info(
            'create_temp_staging_raw_table for {date} already complete, and '
            '"reload" flag was not passed, skipping...'.format(date=date))
        return

    kwargs = kwargs or {}

    sf_config = get_sf_config(secrets_path)
    sf_config_custom = merge_configs(sf_config, sfdb_params)
    ExecutorSR = registered_executors.get(feed_name)
    with ExecutorSR(sf_config_custom) as sf_executor:
        sf_executor.clean_staging_raw_table(staging_raw_table, date, **kwargs)
        activity.logger.info(
            '{table} was cleaned from '
            'the rows with the date (or date range) '
            'of the current workflow run'.format(table=staging_raw_table))


@task.decorate(timeout=300)
def mark_staging_raw_table_tasks_complete(activity, date, feed_name):
    """Set complete task_id staging_raw_table_tasks to feed_name.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed to get executor class.
    """
    task_status.mark_completed_task(feed_name, date, TASK_ID)
    activity.logger.info(
        'All {task_id} of {feed_name} for {date} completed'.format(
            task_id=TASK_ID, feed_name=feed_name, date=date))