"""Amazon Marketshare Garcon tasks.

Tasks to ingest Unlimited Marketshare data.
"""

import csv
from datetime import datetime
from io import StringIO


import boto3
from boto3.exceptions import S3UploadFailedError
from garcon import task
from garcon_contrib.aws.utils import garcon_s3
from garcon_contrib.dynamo_feed_status import \
    garcon_feed_status

from feed_ingestion.flows.amazon_prime_marketshare import config
from feed_ingestion.tasks import bootstrap as reload
from feed_ingestion.util.aws import s3 as s3utils


STOP_RESPONSE = {'stop': True}


@task.decorate(timeout=1000)
@reload.reset_dynamodb_status_on_reload(config.feed_name)
def bootstrap(activity, date, dw_config=None):
    """Bootstrap workflow by injecting intial context from config.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Reporting date (YYYY-MM-DD).

    Returns:
        dict: Initial context for the workflow.
    """
    date_obj = datetime.strptime(date, '%Y-%m-%d')

    temp_staging_raw_table = config.snowflake_table_names[
        'temp_staging_raw'].format(date=date_obj)

    processed_filename = config.s3['preprocessed_filename'].format(
        date=date_obj)

    s3_temp_staging_raw_bucket = '{path}{filename}'.format(
        path=config.s3['preprocessed'], filename=processed_filename)

    return {
        'feed_name': config.feed_name,
        'secrets_path': config.secrets_path,
        'date': date,
        's3_archive_path': config.s3['archive'].format(date=date_obj),
        's3_preprocessed_path': config.s3['preprocessed'],
        'processed_filename': processed_filename,
        's3_temp_staging_raw_bucket': s3_temp_staging_raw_bucket,
        'file_pattern': config.file_pattern.format(date=date_obj),
        'temp_staging_raw_table': temp_staging_raw_table,
        'staging_raw_table': config.snowflake_table_names['staging_raw'],
        'kwargs': {'date_for_sqlloader': date}
    }


def _process_data(s3_archive_path, date, source_files_dict):
    """Process input data and add additional fields.

    Args:
        s3_archive_path (str): S3 path to the archive location.
        date (str): YYYY-MM-DD date of data to delete.
    Returns:
        StringIO: Processed csv data.
    """
    def get_supplier(filename):
        """Return supplier name by filename."""
        for supplier, supplier_name in config.map_suppliers.items():
            if supplier in filename:
                return supplier_name

    def get_fieldnames(date):
        """Return fieldnames based on the date."""
        date_obj = datetime.strptime(date, '%Y-%m-%d')
        cutoff_date = datetime.strptime(config.file_changed_date, '%Y-%m-%d')
        return config.fieldnames_since_2024_07_01 \
            if date_obj >= cutoff_date else config.fieldnames

    rows = []
    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_archive_path)
    s3 = boto3.client('s3')
    fieldnames = get_fieldnames(date)
    for file in source_files_dict['files']:
        file_path = '{path}{file_name}'.format(
            path=bucket_path, file_name=file['file_name'])
        key = s3.get_object(Bucket=bucket, Key=file_path)
        buf = key['Body'].read().decode('utf-8')
        # skip the first line
        reader = csv.DictReader(
            buf.splitlines()[1:], fieldnames=fieldnames, delimiter='\t')

        for row in reader:
            # Skip empty rows
            if not any(row.values()):
                continue
            # save additional fields
            row['territory'] = file['file_name'].split(' ')[0]
            row['supplier'] = get_supplier(file['file_name'])
            row['filename'] = '{date:%Y/%m}/{filename}'.format(
                date=datetime.strptime(date, '%Y-%m-%d'),
                filename=file['file_name'])
            try:
                row['mkt_share'] = int(
                    row['vendor_total_plays']) / int(
                    row['total_plays_for_all_vendors']) * 100
            except (ValueError, ZeroDivisionError):
                row['mkt_share'] = None
            rows.append(row)

    csv_obj = StringIO()
    writer = csv.DictWriter(
        csv_obj,
        fieldnames=fieldnames,
        extrasaction='ignore',
        delimiter='\t')
    writer.writeheader()
    writer.writerows(rows)
    return csv_obj


@task.decorate(timeout=1000)
def process_drop_files(
        activity, feed_name, date, s3_archive_path, s3_preprocessed_path,
        processed_filename, source_files_dict):
    """Process and archive drop files on s3.

    Args:
        activity (ActivityWorker): The garcon activity worker.
        feed_name (str): Name of feed being ingested.
        date (str): YYYY-MM-DD date of data to process.
        s3_archive_path (str): S3 path to the archive location.
        s3_preprocessed_path (str): S3 path to the preprocessed files location.
        processed_filename (str): Filename of processed file.
        source_files_dict (dict): Dict containing metadata of files.
    Returns:
        source_files_dict (dict): Dict containing metadata of processed file.
    """
    csv_obj = _process_data(s3_archive_path, date, source_files_dict)

    try:
        file_path = '{path}{file_name}'.format(
            path=s3_preprocessed_path, file_name=processed_filename)
        source_files_dict = s3utils.upload_processed_to_s3(
            csv_obj, file_path,
            expected_bucket_owner=config.expected_bucket_owner
        )

        activity.logger.info('Successfully processed drop files for {}'.format(
            date))
        return source_files_dict

    except S3UploadFailedError as e:
        garcon_feed_status.set_overall_status(
            feed_name, date, garcon_feed_status.STATUS_NOT_AVAILABLE)
        activity.logger.error(
            'Cannot upload files to {path}. {exception_body}'.format(
                path=s3_preprocessed_path, exception_body=e))
        raise e