"""Amazon Music generators."""
import re

from garcon_contrib.aws.utils import garcon_s3
from garcon_contrib.dynamo_feed_status import garcon_feed_status

from feed_ingestion.flows.amazon_music import config
from feed_ingestion.flows.amazon_music.countries_config \
    import org_countries_for_date


def drop_files_generator(context):
    """Generate raw files to archive.

    Used by the grab_drop_files activity.

    Args:
        context (dict): The current context.
    Yields:
        dict: file to archive and archive destination.
    """
    feed_name = context['bootstrap_feed.feed_name']
    date = context['bootstrap_feed.date']
    expected_files = context['bootstrap_feed.expected_files']
    s3_drop_location = context['bootstrap_feed.s3_drop_location']
    archive_bucket = context['bootstrap_feed.archive_bucket']

    for expected_file in expected_files:
        if _is_downloaded(feed_name, date, expected_file):
            continue
        yield dict(
            source_key_name=_get_key_name(s3_drop_location, expected_file),
            destination_key_name=_get_key_name(archive_bucket, expected_file))


def awal_s3_files_generator(context):
    """Generate raw files to archive.

    Used by the awal_grab_drop_files_s3 activity.

    Args:
        context (dict): The current context.
    Yields:
        dict: file to archive and archive destination.
    """
    archive_bucket = context['bootstrap_feed.archive_bucket']
    filenames = context['bootstrap_feed.expected_files']

    for expected_file, awal_file_path in filenames.items():
        yield dict(
            filename=expected_file,
            source_key_name=awal_file_path,
            destination_key_name=_get_key_name(archive_bucket, expected_file)
        )


def altafonte_s3_files_generator(context):
    """Generate raw files to archive.

    Used by the altafonte_grab_drop_files_s3 activity.

    Args:
        context (dict): The current context.
    Yields:
        dict: file to archive and archive destination.
    """
    archive_bucket = context['bootstrap_feed.archive_bucket']
    filenames = context['bootstrap_feed.expected_files']

    for expected_file, altafonte_file_path in filenames.items():
        yield dict(
            filename=expected_file,
            source_key_name=altafonte_file_path,
            destination_key_name=_get_key_name(archive_bucket, expected_file)
        )


def zip_to_gzip_generator(context):
    """Generate zip and gz filenames.

    Used by the zip_to_gzip_generator activity.
    Files will be renamed -
    ZQAWA_Daily_Unlimited_Activity_20220711_AT.txt.zip became
    ZQAWA_AT_20220711_Daily_Activity_Report.txt.gz.

    Args:
        context (dict): The current context.
    Yields:
        dict: file to archive and archive destination.
    """
    expected_files = context['bootstrap_feed.expected_files']
    archive_bucket = context['bootstrap_feed.archive_bucket']
    clean_path = context['bootstrap_feed.clean_path']
    licensor = context['bootstrap_feed.licensor']

    for expected_file in expected_files:
        match = config.licensors_config[licensor]['file_name_regexp'].match(
            expected_file)
        target_file_name = config.gz_file_template.format(
            org=match.group('org'),
            country=match.group('country'),
            date=match.group('file_date'),
            report_type='Activity',
        )
        yield dict(
            zip_s3_path=f'{archive_bucket}{expected_file}',
            gz_s3_path=f'{clean_path}{target_file_name}')


def sme_drop_files_generator(context):
    """Generate raw files to archive for SME.

    Used by the sme_grab_and_clean_files activity.

    Args:
        context (dict): The current context.
    Yields:
        dict: file to archive and archive destination.
    """
    # feed_name = context['bootstrap_feed.feed_name']
    # date = context['bootstrap_feed.date']
    map_expected_files_to_sme = context['map_expected_files_to_sme.map']
    archive_bucket = context['bootstrap_feed.archive_bucket']

    for expected_file, sme_file_path in map_expected_files_to_sme.items():
        if not sme_file_path:
            # skip missing files in source bucket
            continue
        yield dict(
            source_key_name=sme_file_path,
            destination_key_name=_get_key_name(archive_bucket, expected_file)
        )


def org_country_generator(report_name, only_countries,
                          licensor, date):
    """Generate tuples for report_type, org, country."""
    org_countries = org_countries_for_date(licensor, report_name, date)
    for org, countries in org_countries.items():
        for country in countries:
            if country in only_countries:
                yield org, country


def temp_staging_tables_generator(context):
    """Generate parameters for temporary staging tables.

    Used by the create_temp_staging_raw_table,
    load_temp_staging_raw_table activity and drop_temp_staging_table.

    Args:
        context (dict): The current context.
    Yields:
        dict: Dictionary of temporary staging table names.
    """
    countries = context['bootstrap_feed.countries']
    date = context['bootstrap_feed.date'].replace('-', '')
    feed_name = context['bootstrap_feed.feed_name']
    clean_path = context['bootstrap_feed.clean_path']
    archive_bucket = context['bootstrap_feed.archive_bucket']
    licensor = context['bootstrap_feed.licensor']
    report_name = context['bootstrap_feed.report_name']
    expected_files = context['bootstrap_feed.expected_files']
    date_obj = config.parse_date(context['bootstrap_feed.date']).date()
    if licensor == 'altafonte':
        countries = re.sub(r'(NA)', 'ROW_NA',
                           re.sub(r'(EU)', 'ROE_EU', countries))
    generator = org_country_generator(report_name, countries,
                                      licensor, date_obj)
    if licensor == 'awal' and context['bootstrap_feed.use_s3']:
        for file in expected_files.keys():
            key_dir = archive_bucket + file
            report_type = config.licensors_config[licensor]['report_types'][0]
            org = config.licensors_config[licensor]['org']
            temp_staging_raw_table = _get_table_name(
                feed_name, date, report_type, file.split('_')[5].split('.')[0],
                org, licensor)
            yield dict(
                temp_staging_raw_table=temp_staging_raw_table,
                key_dir=key_dir,
                kwargs=dict(
                    date_for_sqlloader=context['bootstrap_feed.date'],
                    error_on_column_count_mismatch=context[
                        'bootstrap_feed'
                        '.snowflake_error_on_column_count_mismatch'],
                    report_type=report_type))
    else:
        for org, country in generator:
            for report_type in config.licensors_config[
                    licensor]['report_types']:
                temp_staging_raw_table = _get_table_name(
                    feed_name, date, report_type, country, org, licensor)
                key_dir = '{}{}'.format(
                    clean_path, _get_filename(date, report_type, country, org))
                yield dict(
                    temp_staging_raw_table=temp_staging_raw_table,
                    key_dir=key_dir,
                    kwargs=dict(
                        date_for_sqlloader=context['bootstrap_feed.date'],
                        error_on_column_count_mismatch=context[
                            'bootstrap_feed'
                            '.snowflake_error_on_column_count_mismatch'],
                        report_type=report_type))


def load_staging_raw_table_generator(context):
    """Generate parameters for load_staging_raw_table.

    Used by the load_staging_raw_table.

    Args:
        context (dict): The current context.
    Yields:
        dict: Dictionary of temporary staging table names,
            staging_raw_table.
    """
    countries = context['bootstrap_feed.countries']
    date = context['bootstrap_feed.date'].replace('-', '')
    feed_name = context['bootstrap_feed.feed_name']
    licensor = context['bootstrap_feed.licensor']
    report_name = context['bootstrap_feed.report_name']
    expected_files = context['bootstrap_feed.expected_files']
    date_obj = config.parse_date(context['bootstrap_feed.date']).date()
    if licensor == 'altafonte':
        countries = re.sub(r'(NA)', 'ROW_NA',
                           re.sub(r'(EU)', 'ROE_EU', countries))
    generator = org_country_generator(report_name, countries,
                                      licensor, date_obj)
    if licensor == 'awal' and context['bootstrap_feed.use_s3']:
        report_type = config.licensors_config[licensor]['report_types'][0]
        for file in expected_files.keys():
            kwargs = {'licensor': licensor}
            report_type = report_type.lower()
            org = config.licensors_config[licensor]['org']
            kwargs['temp_staging_raw_{}_table'.format(
                report_type)] = _get_table_name(
                feed_name, date, report_type,
                file.split('_')[5].split('.')[0], org, licensor)
            kwargs['{}_report_filename'.format(
                report_type)] = file
            kwargs['date_for_sqlloader'] = context['bootstrap_feed.date']
            yield dict(kwargs=kwargs)
    else:
        for org, country in generator:
            kwargs = {'licensor': licensor}
            for report_type in config.licensors_config[
                    licensor]['report_types']:
                report_type = report_type.lower()
                kwargs['temp_staging_raw_{}_table'.format(
                    report_type)] = _get_table_name(
                    feed_name, date, report_type, country, org, licensor)
                kwargs['{}_report_filename'.format(
                    report_type)] = _get_filename(date, report_type,
                                                  country, org)
                kwargs['date_for_sqlloader'] = context['bootstrap_feed.date']
            yield dict(kwargs=kwargs)


def _get_filename(date, report_type, country, org):
    """Return formatted filename with gz extension."""
    return config.gz_file_template.format(
        report_type=report_type.capitalize(),
        country=country,
        org=org,
        date=date.replace('-', ''))


def _get_table_name(feed_name, date, report_type, country, org, licensor):
    """Return formatted temporary table name."""
    return config.table_name_template.format(
        feed_name=feed_name, report_type=report_type,
        org=org, country=country, date=date, licensor=licensor)


def _get_key_name(bucket, expected_file):
    """Return s3 path without bucket name."""
    path = '{bucket}{expected_file}'.format(
        bucket=bucket,
        expected_file=expected_file)
    return garcon_s3.extract_bucket_path(path)[1]


def _is_downloaded(feed_name, date, expected_file):
    file_status = garcon_feed_status.get_status(
        feed_name,
        date,
        expected_file)
    return file_status == garcon_feed_status.STATUS_DOWNLOADED