"""
Load iTunes raw data into staging_raw_itunes table.

Tasks for ingest iTunes raw data into staging_raw_itunes table.
"""

import datetime

from garcon import task
from garcon_contrib.dynamo_feed_status import garcon_feed_status
from snowflake_connector.etl_connector import SnowflakeSQLExecutor
from snowflake_connector.etl_connector import SQLLoader

from feed_ingestion.conf.config import merge_configs
from feed_ingestion.flows.helpers import get_sf_config
from feed_ingestion.flows.itunes import config
from feed_ingestion.flows.itunes.util import itunes_helpers
from feed_ingestion.tasks.validate_raw_data_tasks_sf \
    import load_temp_staging_raw_table
from feed_ingestion.util import task_status
from feed_ingestion.util.aws import s3 as s3utils


sql_loader = SQLLoader(__file__)

TASK_ID = 'load_staging_raw_itunes'


def sql_loader_version(date):
    """Return loader for latest version of sql."""
    return SQLLoader(__file__, date=date)


@task.decorate(timeout=2000)
def populate_temp_tables(
        activity, date, feed_name, aws_config, sfdb_params, report_data,
        snowflake_error_limit=None, secrets_path=None,
        snowflake_error_on_column_count_mismatch='True', use_s3=False):
    """For each file in files_info, create and populate temp tables from S3.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        feed_name (str): Name of the feed.
        aws_config (dict): AWS related credentials.
        sfdb_params (dict): Dictionary stores Snowflake params.
        report_data (dict): Dictionary from generator that contains
            vendor_id (str), temp_table (str), destination_s3_path (str)
        snowflake_error_limit (int or None): Snowflake error limit.
        secrets_path (str): Secrets manager path of the flow.
        snowflake_error_on_column_count_mismatch:
    """
    is_completed_report = task_status.is_completed_report(feed_name, date)
    if (task_status.is_completed_task(feed_name, date, TASK_ID)
            and is_completed_report):
        return

    activity.logger.info(
        f'Populate temp table. '
        f'table_name: {report_data["temp_table"]}')

    if report_data['vendor_id'] in config.licensors['awal'] and \
            not use_s3 and '_V1_3' in report_data['filename']:
        # iTunes reporter provides awal data under V1_3 version
        # and the files contain State Province and City columns
        # instead of Postal Code, which works for other distributors
        # since 2022-09-05
        sql_loader = sql_loader_version(date='2022-09-05')
    else:
        sql_loader = sql_loader_version(date=date)

    if isinstance(snowflake_error_limit, int):
        snowflake_error_limit = snowflake_error_limit
    else:
        snowflake_error_limit = config.snowflake_error_limit

    # Recreate a fresh temp table
    activity.logger.info('Loading date into: %s', report_data['temp_table'])

    sf_config = merge_configs(get_sf_config(secrets_path), sfdb_params)
    with SnowflakeSQLExecutor(sf_config) as executor:
        executor.execute_query(
            sql_loader, 'create_temp_table', {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'table_name': report_data['temp_table'],
            }
        )
        kwargs = {'query_name': 'load_temp_table',
                  'error_on_column_count_mismatch':
                      snowflake_error_on_column_count_mismatch
                  }
        load_temp_staging_raw_table(
            activity, date, config.feed_name, aws_config,
            report_data['destination_s3_path'], report_data['temp_table'],
            sf_config, snowflake_error_limit, secrets_path, kwargs)


@task.decorate(timeout=2000)
def clear_staging_raw_itunes(
        activity, date, feed_name, sfdb_params, licensor, secrets_path=None):
    """Clear staging_raw_itunes table before loading data.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date of the data file.
        feed_name (str): Name of the feed.
        sfdb_params (dict): Dictionary stores optional Snowflake db and schema.
        licensor (str): The licensor to ingest.
        secrets_path (str): Secrets manager path of the flow.
    """
    if task_status.is_completed_task(feed_name, date, TASK_ID):
        return

    sf_config = merge_configs(get_sf_config(secrets_path), sfdb_params)
    activity.logger.info(f'Clearing staging_raw for {date}')
    with SnowflakeSQLExecutor(sf_config) as executor:
        executor.execute_query(
            sql_loader, 'delete_from_staging_raw', {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'date': date,
                'staging_raw_table': config.staging_raw_table,
                'vendor_ids': config.licensors[licensor]
            }
        )


@task.decorate(timeout=2000)
def load_staging_raw_itunes(activity, date, feed_name, sfdb_params,
                            report_data, secrets_path=None, use_s3=False):
    """Load data on staging_raw_itunes table.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date of the data file.
        feed_name (str): Name of the feed.
        sfdb_params (dict): Dictionary stores optional Snowflake db and schema.
        secrets_path (str): Secrets manager path of the flow.
    """
    is_completed_report = task_status.is_completed_report(feed_name, date)
    if (task_status.is_completed_task(feed_name, date, TASK_ID)
            and is_completed_report):
        return

    activity.logger.info(
        f'Load staging raw itunes. '
        f'table_name: {config.staging_raw_table}')

    if report_data['vendor_id'] in config.licensors['awal'] and \
            not use_s3 and '_V1_3' in report_data['filename']:
        # see the explanation at 58:8
        sql_loader = sql_loader_version(date='2022-09-05')
    else:
        sql_loader = sql_loader_version(date=date)

    sf_config = merge_configs(get_sf_config(secrets_path), sfdb_params)
    file_size = s3utils.get_file_size(
        report_data['destination_s3_path'], config.expected_bucket_owner
    )
    belongs_to_ioda = 'Y' if report_data['vendor_id'] == '80029727' else 'N'
    activity.logger.info('Loading into staging_raw: %s',
                         report_data['temp_table'])

    with SnowflakeSQLExecutor(sf_config) as executor:
        executor.execute_query(
            sql_loader, 'load_staging_raw', {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'staging_raw_table': config.staging_raw_table,
                'table_name': report_data['temp_table'],
                'ingestion_time': datetime.datetime.now(),
                'file_name': report_data['filename'],
                'file_size': file_size,
                'belongs_to_ioda': belongs_to_ioda
            }
        )
        executor.execute_query(
            sql_loader, 'drop_temp_table', {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'table_name': report_data['temp_table'],
            }
        )


@task.decorate(timeout=600)
def set_populated_status(activity, date, feed_name):
    """Set status populated staging raw table.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date of the data file.
        feed_name (str): Name of the feed.
    """
    task_status.mark_completed_task(feed_name, date, TASK_ID)
    # Update overall status to populated staging if status is smaller than
    # ingested
    if itunes_helpers.check_status(
            date, feed_name, garcon_feed_status.STATUS_INGESTED):
        garcon_feed_status.set_overall_status(
            feed_name, date, garcon_feed_status.STATUS_POPULATED_RAW_TABLE)


@task.decorate(timeout=7200)
def update_upc_isrc_on_staging_raw(
        activity, sfdb_params, licensor, date, skip_mapping,
        secrets_path=None):
    """Update upc and isrc on staging_raw_itunes from apple_id_mapping.

    Args:
        activity (ActivityWorker): The activity worker.
        sfdb_params (dict): Dictionary stores Snowflake params.
        licensor (str): The licensor to ingest.
        date (str): Reporting date of the data file.
        skip_mapping (str or None): if 'True' skip this task.
        secrets_path (str): Secrets manager path of the flow.
    """
    if skip_mapping == 'True' or skip_mapping is True:
        return {'skip_mapping': True}
    activity.logger.info('Updating upc-isrc in staging_raw')
    query = 'update_upc_isrc_staging_raw_{}'.format(licensor)
    sf_config = merge_configs(get_sf_config(secrets_path), sfdb_params)
    with SnowflakeSQLExecutor(sf_config) as executor:
        executor.execute_query(
            sql_loader, query, {
                'db': sf_config['db'],
                'schema': sf_config['schema'],
                'staging_raw_table': config.staging_raw_table,
                'vendor_ids': config.licensors[licensor],
                'date': date,
            }
        )