"""MRC Garcon tasks.

Tasks to ingest MRC data.
"""

import copy
import csv
from datetime import datetime
from io import BytesIO
from io import StringIO
import os
import sys

import boto3
from boto3.exceptions import S3UploadFailedError
from garcon import task
from garcon_contrib.aws.garcon_s3 import remove_files_from_path
from garcon_contrib.aws.utils import garcon_s3
from garcon_contrib.dynamo_feed_status import \
    garcon_feed_status
import numpy as np
import pandas as pd

from feed_ingestion.flows.mrc import config
from feed_ingestion.util.aws import s3 as s3utils

STOP_RESPONSE = {'stop': True}


@task.decorate(timeout=1000)
def bootstrap(activity, date, reload):
    """Bootstrap workflow by getting the correct configurations.

    Args:
        activity (ActivityWorker): The activity worker.
        date (str): Reporting date (YYYY-MM-DD).
        reload (str or None): If 'True' delete all feed statuses in DynamoDB.

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

    temp_table_pattern = config.snowflake_table_names[
        'temp_staging_raw'][:-7]  # .replace('_{date}', '')
    temp_staging_raw_table = config.snowflake_table_names[
        'temp_staging_raw'].format(date=date).replace('-', '')

    mapping_temp_staging_raw_tables_list = list(map(
        lambda x:
            config.snowflake_table_names['temp_staging_raw'][:-7] +
            config.mapping_snowflake_table_suffixes[x] +
            config.snowflake_table_names['temp_staging_raw'][-7:].
            format(date=date).replace('-', ''),
            config.mapping_snowflake_table_suffixes))

    mapping_staging_raw_tables_list = list(map(
        lambda x:
            config.snowflake_table_names['staging_raw'] +
            config.mapping_snowflake_table_suffixes[x],
            config.mapping_snowflake_table_suffixes))

    mapping_temp_staging_raw_tables = dict(zip(
        config.mapping_snowflake_table_suffixes.keys(),
        mapping_temp_staging_raw_tables_list))

    mapping_staging_raw_tables = dict(zip(
        config.mapping_snowflake_table_suffixes.keys(),
        mapping_staging_raw_tables_list))

    processed_filename = config.s3['preprocessed_filename'].format(
        date=date).replace('-', '')

    activity.logger.info('Bootstrap flow: {}'.format(date))

    filename = processed_filename.replace('.xlsx', '.csv.gz').replace('-', '')
    s3_archive_bucket = config.s3['archive'].format(date=date)
    s3_mapping_archive_bucket = config.s3['mapping_archive']
    s3_temp_staging_raw_bucket = config.s3['preprocessed'].format(date=date)
    s3_mapping_temp_staging_raw_bucket = config.s3['mapping_preprocessed']

    file = {'files': [{'file_name': config.file_pattern.format(
        date=date).replace('-', '')}]}

    return dict(
        feed_name=config.feed_name,
        secrets_path=config.secrets_path,
        date=date,
        s3_archive_path=s3_archive_bucket,
        s3_mapping_archive_path=s3_mapping_archive_bucket,
        processed_filename=processed_filename,
        s3_preprocessed_path=config.s3['preprocessed'].format(date=date),
        s3_temp_staging_raw_bucket=s3_temp_staging_raw_bucket + filename,
        s3_mapping_temp_staging_raw_bucket=s3_mapping_temp_staging_raw_bucket,
        file_pattern=config.file_pattern,
        temp_staging_raw_table=temp_staging_raw_table,
        temp_table_pattern=temp_table_pattern,
        mapping_temp_staging_raw_tables=mapping_temp_staging_raw_tables,
        staging_raw_table=config.snowflake_table_names['staging_raw'],
        mapping_staging_raw_tables=mapping_staging_raw_tables,
        file=file,
        reload=reload
    )


@task.decorate(timeout=1000)
def clear_s3_folders(activity, reload, date):
    """Delete files from S3 folders if 'reload'.

    Args:
        activity (ActivityWorker): The activity worker.
        reload (str or None): If 'True' delete all feed statuses in DynamoDB.
        date (str): Reporting date (YYYY-MM-DD).
    """
    s3_temp_staging_raw_bucket = config.s3['preprocessed'].format(date=date)
    s3_archive_bucket = config.s3['archive'].format(date=date)

    if reload == 'True':
        activity.logger.info('Clear S3 files flow: {}'.format(date))

        activity.logger.info('Delete status for feed: {} {} '.format(
            config.feed_name, date))
        garcon_feed_status.delete_status(config.feed_name,
                                         date)

        preprocessed_files = remove_files_from_path(
            activity, s3_temp_staging_raw_bucket, True)
        archive_files = remove_files_from_path(activity,
                                               s3_archive_bucket, True)
        if len(archive_files.get('s3.files_removed')) > 0:
            activity.logger.info(','.
                                 join(archive_files.get('s3.files_removed')) +
                                 ' files have been removed from ' +
                                 s3_archive_bucket + '.')
        if len(preprocessed_files.get('s3.files_removed')) > 0:
            activity.logger.info(','.
                                 join(preprocessed_files.
                                      get('s3.files_removed')) +
                                 ' files have been removed from ' +
                                 s3_temp_staging_raw_bucket + '.')


def _process_data(s3_archive_path, file, date):
    """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.
    """
    csv.field_size_limit(sys.maxsize)
    rows = []

    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_archive_path)
    s3 = boto3.client('s3')

    file_path = '{path}{file_name}'.format(
        path=bucket_path, file_name=file['file_name'])
    key = s3.get_object(Bucket=bucket, Key=file_path)

    # here we read the data the file contains.
    # data has xlsx format, so to correctly write
    # it to DB we should convert it to csv-type.
    data = key['Body'].read()

    # In order to efficiently merge data at the table it's
    #  been decided to remove all the columns having the
    # last seven days (excluding today's one) as names (16th -
    # 22d columns). You could check comments on Jira.
    buf = pd.read_excel(BytesIO(data), engine='openpyxl').values
    buf = np.delete(buf, [16, 17, 18,
                          19, 20, 21], 1)

    data_list = buf.tolist()
    csv_buf = []

    for item in data_list:
        csv_buf.append('\t'.join(str(v).replace('"', '')
                                 if str(v).count('"') == 1
                                 else str(v) for v in item))
    reader = csv.DictReader(
        csv_buf[1:], fieldnames=config.fieldnames, delimiter='\t')

    for row in reader:
        row['processing_date'] = datetime.today(). \
            strftime('%Y-%m-%d %H:%M:%S')
        row['filename'] = '{}'.format(
            file['file_name'])
        row['download_date'] = '{}'.format(
            date)
        rows.append(row)

    csv_obj = StringIO()

    fieldnames = copy.deepcopy(config.fieldnames)
    fieldnames.append('processing_date')
    fieldnames.append('filename')
    fieldnames.append('download_date')

    writer = csv.DictWriter(
        csv_obj, fieldnames=fieldnames, delimiter='\t')

    writer.writeheader()
    writer.writerows(rows)

    return csv_obj


def _process_mapping_data(s3_archive_path, file, date):
    """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.
    """
    csv.field_size_limit(sys.maxsize)
    rows = []
    mapping = file['file_name'][:file['file_name'].rfind('_')]

    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_archive_path)
    s3 = boto3.client('s3')

    file_path = '{path}/{file_name}'.format(
        path=bucket_path.format(mapping=mapping, date=date),
        file_name=file['file_name'])
    key = s3.get_object(Bucket=bucket, Key=file_path)

    # here we read the data the file contains.
    # data has xlsx format, so to correctly write
    # it to DB we should convert it to csv-type.
    data = key['Body'].read().replace(b'\x00', b'')

    # buf = pd.read_excel(BytesIO(data), engine='openpyxl').values
    buf = data.decode('utf-8').splitlines()
    # csv_buf = []
    fieldnames = buf[0].split('\t')
    # for item in buf:
    #     csv_buf.append(str(v).replace('"', '')
    #                    if str(v).count('"') == 1
    #                    else str(v) for v in item)
    reader = csv.DictReader(
        buf[1:], fieldnames=fieldnames, delimiter='\t')

    for row in reader:
        row['processing_date'] = datetime.today(). \
            strftime('%Y-%m-%d %H:%M:%S')
        row['filename'] = '{}'.format(
            file['file_name'])
        row['download_date'] = '{}'.format(
            date)
        rows.append(row)

    csv_obj = StringIO()

    fieldnames.append('processing_date')
    fieldnames.append('filename')
    fieldnames.append('download_date')

    writer = csv.DictWriter(
        csv_obj, fieldnames=fieldnames, delimiter='\t')

    writer.writeheader()
    writer.writerows(rows)

    return csv_obj, file['file_name']


@task.decorate(timeout=1000)
def process_drop_files(
        activity, feed_name, date, s3_archive_path,
        s3_preprocessed_path,
        processed_filename,
        file):
    """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.
        file: File metadata.
    Returns:
        file (dict): Metadata of processed file.
    """
    file = file['files'][0]

    csv_obj = _process_data(s3_archive_path, file, date)

    processed_filename = processed_filename.replace('.xlsx', '.csv.gz')

    try:
        file_path = os.path.join(
            s3_preprocessed_path,
            processed_filename)
        processed_file = s3utils.upload_processed_to_s3(
            csv_obj, file_path,
            expected_bucket_owner=config.expected_bucket_owner
        )
        csv_obj.close()

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

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


@task.decorate(timeout=21600)
def process_drop_mapping_files(
        activity, feed_name, date, s3_archive_path,
        s3_preprocessed_bucket, file):
    """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_bucket (str): S3 path to the preprocessed files.
        file: File metadata.
    Returns:
        file (dict): Metadata of processed file.
    """
    processed_files = {'source_files_dict': {'files': []}}

    if file.get('found') and file.get('file_size') > 0:
        mapping = file['file_name'][:file['file_name'].rfind('_')]
        csv_obj, filename = _process_mapping_data(s3_archive_path, file, date)

        processed_filename = filename.replace('.txt', '.csv.gz')

        try:
            file_path = os.path.join(
                s3_preprocessed_bucket.format(mapping=mapping, date=date),
                processed_filename)
            processed_file = s3utils.upload_processed_to_s3(
                csv_obj, file_path,
                expected_bucket_owner=config.expected_bucket_owner
            )
            csv_obj.close()

            activity.logger.info(
                f'Successfully processed drop files for {date}')
            processed_files['source_files_dict']['files'].append(
                processed_file['source_files_dict']['files'][0])

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