"""elasticsearch_export workflow tasks."""

from collections import deque
import gzip
import os
import shutil

import boto3
from botocore.exceptions import ClientError
from garcon import task
from opensearchpy import OpenSearch, \
    RequestsHttpConnection as RequestsHttpConnectionOpenSearch
from opensearchpy.helpers import parallel_bulk as opensearch_bulk
from requests_aws4auth import AWS4Auth
from yt_conflict_elasticsearch.flows.elasticsearch_export import config
from yt_conflict_elasticsearch.flows.elasticsearch_export \
    import conflict_csv_to_json
from yt_conflict_elasticsearch.flows.elasticsearch_export import generators
from yt_conflict_elasticsearch.flows.elasticsearch_export.logger \
    import app_logger
from yt_conflict_elasticsearch.flows.elasticsearch_export.snowflake_executor \
    import YouTubeConflictSFExecutor
from yt_conflict_elasticsearch.util import file_util
from yt_conflict_elasticsearch.util import task_status
from yt_conflict_elasticsearch.util import task_status as reload


@task.decorate(timeout=300)
@reload.reset_dynamodb_status_on_reload(config.SWF_FLOW_NAME)
def bootstrap(activity, date):
    """Bootstrap the flow."""
    app_logger.info('Bootstrap was started for date {}'.format(date))

    return {
        'date': date
    }


@task.decorate(timeout=600)
def task_store_conflicts_in_database_as_csv(activity, date, conflict_status):
    """Store conflicts in database to S3 bucket as a CSV file.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
        conflict_status (str): Conflict status to use.
    """
    task_name = 'store_conflicts_in_database_as_csv_{}'.format(conflict_status)
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return
    app_logger.info('Started task {}'.format(task_name))

    s3_file_key = file_util.get_s3_csv_gz_file_key(date, conflict_status)
    s3_file_url = 's3://{}/{}'.format(config.S3_BUCKET, s3_file_key)
    app_logger.info(
        'Storing {} conflicts from SF into csv: {}'.format(
            conflict_status, s3_file_url))
    with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
        snowflake_executor.store_csv_to_s3_conflicts(
            s3_file_url, conflict_status)

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=86400)
def task_transform_csv_to_json_conflicts(activity, date, conflict_status):
    """Convert csv to json and store in S3 Bucket.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
        conflict_status (str): Conflict status to use.
    """
    task_name = 'transform_csv_to_json_conflicts_{}'.format(conflict_status)
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return

    app_logger.info('Started task {}'.format(task_name))

    temp_path = _create_tmp_path(date)

    # Pull the files from s3 to staging folder
    s3_bucket = boto3.resource('s3').Bucket(config.S3_BUCKET)
    s3_csv_gz_file_key = file_util.get_s3_csv_gz_file_key(
        date, conflict_status)

    app_logger.info(
        'Downloading file key {} from bucket {}'.format(
            s3_csv_gz_file_key, s3_bucket))

    # Download gzipped CSV file
    gzipped_csv_full_path = os.path.join(
        temp_path, os.path.basename(s3_csv_gz_file_key))
    if not _download_file(
            s3_bucket, s3_csv_gz_file_key, gzipped_csv_full_path):
        # there are no conflicts to transform
        app_logger.info('No file found, aborting task {}'.format(task_name))
        task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)
        return

    # Process gzipped CSV and write results to json output file
    json_output_full_path = os.path.join(
        temp_path, file_util.get_json_filename(conflict_status))
    f = open(json_output_full_path, 'w')
    conflict_csv_to_json.process_csv_file(gzipped_csv_full_path, f)
    f.close()

    # Compress json output file
    json_gz_output_full_path = _gzip_file(json_output_full_path)

    app_logger.info('Created file {}.'.format(json_gz_output_full_path))

    # Upload json output to s3 bucket
    s3_json_gz_file_key = file_util.get_s3_json_gz_file_key(
        date, conflict_status)
    s3_bucket.upload_file(json_gz_output_full_path, s3_json_gz_file_key)

    # Remove temp files
    os.remove(gzipped_csv_full_path)
    os.remove(json_output_full_path)
    os.remove(json_gz_output_full_path)

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=600)
def task_create_elasticsearch_index(activity, date):
    """Create index in elasticsearch if it doesn't exist.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
    """
    task_name = 'create_elasticsearch_index'
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return

    app_logger.info('Started task {}'.format(task_name))

    index_name_os = config.OS_CONFLICTS_INDEX_NAME
    alias_name_os = config.OS_CONFLICTS_ALIAS_NAME
    opensearch_client = _get_opensearch_client()

    if not opensearch_client.indices.exists_alias(name=alias_name_os):
        app_logger.info(
            'OS alias {} does not exist, creating'.format(alias_name_os))
        if not opensearch_client.indices.exists(index=index_name_os):
            app_logger.info(
                'OS index {} does not exist, creating'.format(index_name_os))
            opensearch_client.indices.create(index=index_name_os)
        opensearch_client.indices.put_alias(
            name=alias_name_os, index=index_name_os)
    elif config.ENVIRONMENT == 'qa':
        app_logger.info(
            'OS qa environment detected, recreating index {}'
            .format(index_name_os))
        try:
            opensearch_client.delete_by_query(
                    index=alias_name_os, body={'query': {'match_all': {}}})
        except Exception as e:
            app_logger.error('Error deleting OS alias {}: {}'.
                             format(alias_name_os, e))
    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=86400)
def task_populate_elasticsearch_index(
        activity, date, conflict_status):
    """Populate index in elasticsearch with JSON data.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
        conflict_status (str): Conflict status to use.
    """
    task_name = 'populate_elasticsearch_index_{}'.format(conflict_status)
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return

    app_logger.info('Started task {}'.format(task_name))
    temp_path = _create_tmp_path(date)

    s3_bucket = boto3.resource('s3').Bucket(config.S3_BUCKET)
    s3_json_gz_file_key = file_util.get_s3_json_gz_file_key(
        date, conflict_status)

    # Download JSON data
    es_json_filename = file_util.get_json_filename(
        conflict_status, gzipped=True)
    json_gz_output_full_path = os.path.join(temp_path, es_json_filename)

    app_logger.info(
        'Reading data from file {} from bucket {}'.format(
            s3_json_gz_file_key, config.S3_BUCKET))

    if not _download_file(
            s3_bucket, s3_json_gz_file_key, json_gz_output_full_path):
        # there are no conflicts to add to ES
        app_logger.info('No file found, aborting task {}'.format(task_name))
        task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)
        return

    # Populate index
    opensearch_client = _get_opensearch_client()

    # Update the territories for the conflicts above
    app_logger.info(
        'Updating partially resolved conflicts in opensearch...')

    # Get an iterator for conflicts to update in place
    conflicts_to_update_iterator_os = \
        generators.get_es_ids_to_update_territories_os(
            json_gz_output_full_path)

    _update_partially_resolved_conflicts_in_os(
        opensearch_client, conflicts_to_update_iterator_os)

    # Insert the new conflicts to OpenSearch
    app_logger.info('Done. Inserting new conflicts into OS...')

    # Get an iterator for conflicts to insert into OS
    # Retrieves conflict_ids from the JSON where es_id is not present
    conflicts_to_insert_iterator_os = \
        generators.get_conflict_data_to_populate_os(
            json_gz_output_full_path, config.OS_CONFLICTS_ALIAS_NAME)
    app_logger.info('Got conflicts_to_insert_iterator_os')

    app_logger.info('Creating OpenSearch bulk iterator...')
    opensearch_iterator = opensearch_bulk(
        opensearch_client,
        conflicts_to_insert_iterator_os,
        chunk_size=config.CHUNK_SIZE,
        max_chunk_bytes=config.MAX_CHUNK_BYTES,
        thread_count=config.THREAD_COUNT,
        queue_size=config.QUEUE_SIZE
    )

    # Another run through new conflicts to fill Snowlflake with ES ids.
    # Retrieves conflict_ids from the JSON where es_id is not present (again)
    app_logger.info('Creating conflict OS IDs iterator...')
    conflict_ids_iterator_os =\
        generators.get_unindexed_conflict_ids_by_line(
            json_gz_output_full_path)

    # Get a generator of (conflict_id, es_id) pairs to update snowflake with
    # This will perform the inserts into opensearch and return the new es_id
    # for each conflict_id on success
    app_logger.info('Creating conflict OS IDs to insert into SF iterator...')
    opensearch_insert_iterator = generators.get_batch_to_insert_to_sf_temp_os(
        conflict_ids_iterator_os, opensearch_iterator)

    # Insert es_ids in snowflake
    with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
        # Create a temp table to hold the (conflict_id, es_id) pairs
        snowflake_executor.create_temp_conflict_to_es_id()
        app_logger.info('Inserting OS IDs into SF temp table...')
        for insert_operation in opensearch_insert_iterator:
            # Perform the insert operation, store results in the temp table
            snowflake_executor.fill_temp_table_with_es_ids(
                config.SF_PARAMS['temp_conflict_es_id_table'],
                insert_operation)

        app_logger.info('Done. Updating OS IDs in snowflake...')
        # Bulk update existing conflicts with the es_ids from the temp table
        snowflake_executor.populate_es_ids_to_conflicts()
        app_logger.info('Done. Cleaning up temp files...')

    # Remove temp file
    os.remove(json_gz_output_full_path)

    # Remove temp directory but only if it's empty
    try:
        os.rmdir(temp_path)
    except OSError:
        pass

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=600)
def task_mark_indexed_conflicts(activity, date):
    """Mark indexed conflicts.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
    """
    task_name = 'mark_indexed_conflicts'
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return
    app_logger.info('Started task {}'.format(task_name))

    with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
        snowflake_executor.mark_indexed_conflicts()

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=600)
def clean_up_s3(activity, date):
    """Delete the csv and json conflicts files from S3 Bucket.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
    """
    task_name = 'clean_up_s3'
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return
    app_logger.info('Started task {}'.format(task_name))

    s3_client = boto3.client('s3')

    s3_tmp_path = config.S3_FOLDER_TEMPLATE.format(date=date.replace('-', '_'))

    objects = s3_client.list_objects_v2(
        Bucket=config.S3_BUCKET, Prefix=s3_tmp_path).get('Contents', [])
    keys_to_delete = [{'Key': obj['Key']} for obj in objects]

    if keys_to_delete:
        app_logger.info(
            'Deleting {} objects from {}'.format(len(keys_to_delete),
                                                 s3_tmp_path))
        s3_client.delete_objects(
            Bucket=config.S3_BUCKET,
            Delete={'Objects': keys_to_delete})

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=600)
def remove_responded_conflicts_from_es(activity, date):
    """Remove responded conflicts from ES.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
    """
    task_name = 'remove_responded_conflicts_from_es'
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return
    app_logger.info('Started task {}'.format(task_name))

    with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
        es_ids = snowflake_executor.get_responded_conflicts()

    if es_ids:
        opensearch_client = _get_opensearch_client()

        conflicts_to_delete_iterator_os = generators.conflicts_to_remove_os(
            es_ids, config.OS_CONFLICTS_ALIAS_NAME)

        opensearch_iterator = opensearch_bulk(
            opensearch_client, conflicts_to_delete_iterator_os,
            chunk_size=config.CHUNK_SIZE,
            max_chunk_bytes=config.MAX_CHUNK_BYTES,
            raise_on_error=False,
            thread_count=config.THREAD_COUNT,
            queue_size=config.QUEUE_SIZE)

        deque(opensearch_iterator)

        # Set es_id to NULL for all conflicts we just removed from ES
        with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
            snowflake_executor.remove_selected_es_ids(es_ids)

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


@task.decorate(timeout=86400)
def remove_resolved_conflicts_from_es(activity, date):
    """Remove resolved conflicts from ES.

    Args:
        activity (ActivityWorker): The Garcon activity worker.
        date (str): Run date (YYYY-MM-DD).
    """
    task_name = 'remove_resolved_conflicts_from_es'
    if task_status.is_completed_task(config.SWF_FLOW_NAME, date, task_name):
        return
    app_logger.info('Started task {}'.format(task_name))

    with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
        es_ids = snowflake_executor.get_resolved_conflicts()

    if es_ids:
        app_logger.info(
            'Found {} conflicts that have been resolved but are still '
            'existing in ES. About to delete.'.format(len(es_ids)))

        opensearch_client = _get_opensearch_client()

        conflicts_to_delete_iterator_os = generators.conflicts_to_remove_os(
            es_ids, config.OS_CONFLICTS_INDEX_NAME)

        opensearch_iterator = opensearch_bulk(
            opensearch_client, conflicts_to_delete_iterator_os,
            chunk_size=config.CHUNK_SIZE,
            max_chunk_bytes=config.MAX_CHUNK_BYTES,
            raise_on_error=False,
            thread_count=config.THREAD_COUNT,
            queue_size=config.QUEUE_SIZE)

        deque(opensearch_iterator)

        # Set es_id to NULL for all conflicts we just removed from ES
        with YouTubeConflictSFExecutor(config.SF_PARAMS) as snowflake_executor:
            batch_size = config.SF_UPDATE_RESOLVED_BATCH_SIZE
            for i in range(0, len(es_ids), batch_size):
                snowflake_executor.remove_selected_es_ids(
                    es_ids[i:i + batch_size])

    app_logger.info('{} task finished'.format(task_name))
    task_status.mark_completed_task(config.SWF_FLOW_NAME, date, task_name)


def _download_file(s3_bucket, s3_key, destination_path):
    """Download file.

    Args:
        s3_bucket: boto3.resource('s3').Bucket
        s3_key (str): the name of the key to download from
        destination_path (str): the path to the file to download to

    Returns:
        bool: True if file was downloaded successfully
            False if file does not exist
    """
    try:
        s3_bucket.download_file(s3_key, destination_path)
    except ClientError as e:
        try:
            err_code = e.response['Error']['Code']
            if err_code == '404':
                # file does not exist
                return False
            else:
                raise e
        except (KeyError, AttributeError):
            raise e
    return True


def _get_opensearch_client():
    """Get OpenSearch client."""
    region = 'us-east-1'
    service = 'es'
    credentials = boto3.Session().get_credentials()
    awsauth = AWS4Auth(
        credentials.access_key, credentials.secret_key, region, service,
        session_token=credentials.token
    )
    app_logger.info('OpenSearch client created')
    return OpenSearch(
        [config.OPENSEARCH_HOST], use_ssl=config.OPENSEARCH_USE_SSL,
        port=config.OPENSEARCH_PORT, timeout=config.OPENSEARCH_TIMEOUT,
        max_retries=config.OPENSEARCH_MAX_RETRIES, retry_on_timeout=True,
        http_auth=awsauth, connection_class=RequestsHttpConnectionOpenSearch)


def _create_tmp_path(date):
    """Create and return tmp path."""
    date = date.replace('-', '_')
    temp_path = config.ELASTICSEARCH_EXPORT_TEMP_PATH.format(date=date)
    if not os.path.isdir(temp_path):
        os.makedirs(temp_path, exist_ok=True)
    return temp_path


def _gzip_file(file_path):
    gzip_file_path = file_path + '.gz'
    input_file = open(file_path, 'rb')

    with gzip.open(gzip_file_path, 'wb') as gzip_file:
        shutil.copyfileobj(input_file, gzip_file)

    return gzip_file_path


def _update_partially_resolved_conflicts_in_os(
        opensearch_client, conflicts_to_update_iterator):
    """Update partially resolved conflicts in place with new territories."""
    opensearch_iterator = \
        opensearch_bulk(opensearch_client, conflicts_to_update_iterator,
                        chunk_size=config.CHUNK_SIZE,
                        max_chunk_bytes=config.MAX_CHUNK_BYTES,
                        thread_count=config.THREAD_COUNT,
                        queue_size=config.QUEUE_SIZE)

    # Do the actual update
    try:
        deque(opensearch_iterator)
        app_logger.info('Finished updating partially resolved conflicts in OS')
    except Exception as e:
        app_logger.error(
            'Error updating partially resolved conflicts in OS: {}'.format(e))