"""Update Smithsonian Metadata via GraphQL."""
import argparse
import os

from sys import exit as sys_exit
from signal import signal, SIGINT, SIGTERM

from connectors.logging import log
from constants.backoff import TASK_TIMEOUT
from constants.fields import FILTER_FIELDS
from logic.concurrency_logic import (
    process_contrib_data,
    process_data_async,
    get_participant_id_data,
    get_product_data_from_graph,
    set_release_localization_data,
    set_track_localizations_data,
)
from logic.graphql_logic import (
    transform_participations,
    transform_performers,
    transform_publishers,
)
from utils.file_utils import (
    parse_output_data_to_report,
    preprocess_cached_data,
    write_var_to_json,
    zip_files,
)
from utils.general_utils import (
    signal_handler,
    log_runtime,
    check_env,
)
from utils.processing_utils import (
    get_participants_from_datafile,
    pivot_list_of_graph_products,
    process_localization_errors,
    process_localization_responses,
    process_release_updates,
    process_track_updates,
    preprocess_localization_data,
    preprocess_smithsonian_contributor_data,
)

from config import (
    ABORT_ON_FAIL,
    ASYNC_CHUNK_SIZE,
    ASYNC_POOL_SIZE,
    CACHED_DATA_PATH,
    CREATE_ZIP_FILES,
    INPUT_FILE_NAME,
    INPUT_FILE_STARTING_ROW_OFFSET,
    JITTER_FACTOR,
    LOG_FREQUENCY,
    LOGGER_LEVEL,
    MAP_TASK,
    NROWS,
    PERFORMER_COUNT,
    REQUEST_DELAY,
    TASK_DELAY,
    WRITE_MASTER_OUTPUT,
)


@log_runtime
def process_cached_smithsonian_contribs(filename: str):
    """Run set_track_contributions on cached data.

    Args:
        filename (str): The filename to process.
    Returns:
        dict: The processed contributor data.
        dict: The errors encountered during processing.
    """
    contributor_data = preprocess_cached_data(filename)
    # Multi-processes
    collected_data, collected_errors = process_data_async(
        process_contrib_data,
        contributor_data,
        pool_size=ASYNC_POOL_SIZE,
        chunk_size=ASYNC_CHUNK_SIZE
    )

    return collected_data, collected_errors


@log_runtime
def process_cached_localizations(filename: str):
    """Run set_localizations on cached data.

    Args:
        filename (str): The filename to process.
    Returns:
        dict: The processed contributor data.
        dict: The errors encountered during processing.
    """
    collected_errors = {}
    collected_data = []

    localization_data = []

    localization_data = preprocess_cached_data(filename)

    # Multi-processes some method
    collected_data, collected_errors = process_data_async(
        set_track_localizations_data,
        localization_data,
        pool_size=ASYNC_POOL_SIZE,
        chunk_size=ASYNC_CHUNK_SIZE
    )

    return collected_data, collected_errors


@log_runtime
def process_localizations(filename: str = None):
    """Process localizations for a file

    Args:
        filename (str, optional): The filename to process. Defaults to None.

    Returns:
        tuple (dict, dict, dict): The processed data, errors, and replay data.
    """
    collected_errors = {}
    collected_data = []

    # Preprocess a file
    track_localization_data, release_localization_data = \
        preprocess_localization_data(filename)

    # Log the dimensions of the data
    log.info(
        f'Track-level localizations has {len(track_localization_data)} rows'
    )
    log.info(
        f'Release-level localizations has {len(release_localization_data)} '
        'rows'
    )

    # TODO: Replace ITUNES_LANGUAGES constant with call to Snowflake here
    # SELECT * FROM "ITUNES_LANGUAGES"

    # Process track localization updates
    log.info('Collecting localizations for all tracks...')

    track_updates_by_vendor = {}

    products_by_vendor_subaccount_upc = {}

    # Get all vendors and subaccounts in the track input
    track_vendors = list(set([
        (track['Vendor ID'], track.get('Subaccount ID', 0))
        for track in track_localization_data
    ]))

    # Get all vendors and subaccounts in the release input
    release_vendors = list(set([
        (release['Vendor ID'], release.get('Subaccount ID', 0))
        for release in release_localization_data
    ]))

    # Get all UPCs by vendor in the track input
    track_upcs_by_vendor = {}
    for vend_sub_pair in track_vendors:
        track_upcs_by_vendor[vend_sub_pair] = \
            list(set([
                track['UPC'] for track in track_localization_data
                if track['Vendor ID'] == vend_sub_pair[0]
            ]))

    # Make Vendor, sub, upc tuples for track-level
    all_track_upcs_vendors_subs = list(set(
        [
            (vend_sub_pair[0], vend_sub_pair[1], upc)
            for vend_sub_pair, upc_list
            in track_upcs_by_vendor.items()
            for upc in upc_list
        ]
    ))

    # Get all UPCs by vendor in the release input
    release_upcs_by_vendor = {}
    for vend_sub_pair in release_vendors:
        release_upcs_by_vendor[vend_sub_pair] = \
            list(set([
                release['UPC'] for release in release_localization_data
                if release['Vendor ID'] == vend_sub_pair[0]
            ]))

    all_release_upcs_vendors_subs = list(set(
        [
            (vend_sub_pair[0], vend_sub_pair[1], upc)
            for vend_sub_pair, upc_list
            in release_upcs_by_vendor.items()
            for upc in upc_list
        ]
    ))

    # Note: We want all unique UPC's across both inputs, as we only want to
    #   make one GraphQL query per UPC, so get all UPC/Vendor pairs as tuples
    #   to leverage set logic
    all_upcs_vendors_subs = list(set(
        all_track_upcs_vendors_subs +
        all_release_upcs_vendors_subs
    ))

    # Get all track data for all UPC's, and store by UPC - Make the Graph calls
    products_from_graph, _ = process_data_async(
        get_product_data_from_graph,
        all_upcs_vendors_subs,
        pool_size=ASYNC_POOL_SIZE,
        chunk_size=ASYNC_CHUNK_SIZE
    )

    products_by_vendor_subaccount_upc = \
        pivot_list_of_graph_products(products_from_graph)

    # NOTE: -------------------------------------------------------------------
    # NOTE: At this point, we have all the product-level data we need for
    # NOTE: track and release localization updates.
    # NOTE: -------------------------------------------------------------------

    # Make the track-level update payloads and store them by vendor
    _, track_updates_by_vendor = process_track_updates(
        track_localization_data,
        products_by_vendor_subaccount_upc,
    )

    # Flatten the track updates for async processing
    flattened_track_payload = [
        (vendor_id, subaccount_id, upc, isrc, payload)
        for vendor_id, subaccounts in track_updates_by_vendor.items()
        for subaccount_id, upcs in subaccounts.items()
        for upc, isrcs in upcs.items()
        for isrc, payload in isrcs.items()
    ]

    # Make the track-level update payloads and store them by vendor
    _, track_updates_by_vendor = process_release_updates(
        release_localization_data,
        products_by_vendor_subaccount_upc,
    )

    flattened_release_payload = [
        (vendor_id, subaccount_id, upc, payload)
        for vendor_id, subaccounts in track_updates_by_vendor.items()
        for subaccount_id, upcs in subaccounts.items()
        for upc, payload in upcs.items()
    ]

    # NOTE: -------------------------------------------------------------------
    # NOTE: At this point, we have all the track-level and release-level
    # NOTE:  updates prepared
    # NOTE: -------------------------------------------------------------------

    # Attempt to make track updates
    log.info('Making Track Localization Updates...')
    # Call the track localization jobs
    track_localization_data, track_localization_errors = process_data_async(
        set_track_localizations_data,
        flattened_track_payload,
        pool_size=ASYNC_POOL_SIZE,
        chunk_size=ASYNC_CHUNK_SIZE
    )

    # Attempt to make release updates
    log.info('Making Release Localization Updates...')
    release_localization_data, release_localization_errors = \
        process_data_async(
            set_release_localization_data,
            flattened_release_payload,
            pool_size=ASYNC_POOL_SIZE,
            chunk_size=ASYNC_CHUNK_SIZE
        )

    # Merge the responses to coherent success and fails...
    collected_data = process_localization_responses(
        track_localization_data, release_localization_data)
    collected_errors = process_localization_errors(
        track_localization_errors, release_localization_errors)


    # Return the data
    return collected_data, collected_errors


@log_runtime
def process_smithsonian_contribs(filename: str = None):
    """Attach contributors to tracks via GraphQL."""
    # Load the data froom the Smithsonian file
    participations, performers = \
        preprocess_smithsonian_contributor_data(filename)

    # Log the dimensions of the data
    log.info(
        f'Participations has {len(participations)} rows'
    )

    # Log the dimensions of the data
    log.info(
        f'Performers has {len(performers)} rows'
    )

    log.info('Collecting Participations...')
    # Get all the participants from the datafile keyed by vendor
    all_artists_by_vendor = \
        get_participants_from_datafile(participations)

    log.info(f'Collecting {len(all_artists_by_vendor)} Participant IDs...')

    # Collect all the participant IDs
    participant_ids, errors = process_data_async(
        get_participant_id_data,
        all_artists_by_vendor,
        pool_size=ASYNC_POOL_SIZE,
        chunk_size=ASYNC_CHUNK_SIZE
    )

    if WRITE_MASTER_OUTPUT:
        write_var_to_json(participant_ids)

    if errors:
        # Count the number of errors
        error_count = sum(
            len(errors['vendor'][vendor_id][subaccount_id])
            for vendor_id in errors['vendor'].keys()
            for subaccount_id in errors['vendor'][vendor_id].keys()
        )
        write_var_to_json(errors, 'get_participant_errors.json')
        msg = f'Failed to process {error_count} rows for ' \
            f'Participant IDs. Check the input file for errors.'
        log.error(msg)

        if ABORT_ON_FAIL:
            raise ValueError(msg)

    msg = 'Transforming Participations...'
    log.info(msg)

    participant_data, errors = \
        transform_participations(participations, participant_ids)

    if WRITE_MASTER_OUTPUT:
        write_var_to_json(participant_data)

    if errors:
        write_var_to_json(errors, 'transform_participations_errors.json')
        msg = f'Failed to process {len(errors["tuid"].keys())} rows for ' \
            f'Participations. Check the input file for errors.'
        log.error(msg)

        if ABORT_ON_FAIL:
            raise ValueError(msg)

    msg = 'Transforming Publishers...'
    log.info(msg)

    publisher_data, errors = \
        transform_publishers(participations)

    if WRITE_MASTER_OUTPUT:
        write_var_to_json(publisher_data)

    if errors:
        write_var_to_json(errors, 'transform_publishers_errors.json')
        msg = f'Failed to process {len(errors.keys())} rows for ' \
            f'Publishers. Check the input file for errors.'
        log.error(msg)

        if ABORT_ON_FAIL:
            raise ValueError(msg)

    msg = 'Transforming Performers...'
    log.info(msg)

    # Feed participant data into contributor data
    performer_data, errors = \
        transform_performers(performers)

    if WRITE_MASTER_OUTPUT:
        write_var_to_json(performer_data)

    if errors:
        write_var_to_json(errors, 'transform_performers_errors.json')
        msg = f'Failed to process {len(errors.keys())} rows for ' \
            f'Performers. Check the input file for errors.'
        log.error(msg)

        if ABORT_ON_FAIL:
            raise ValueError(msg)

    # If you want to sanity check vendor_ids and subaccount_ids
    # vendors_match, subaccounts_match = check_vendor_subaccount_match(
    #     contributor_data)

    # Get all unique tuids from participations and performers
    all_tuids = list(set(participant_data.keys()).union(performer_data.keys()))

    # Recompose the three data variables to a single variable
    contributor_data = {
        tuid: {
            'vendor_id': participant_data.get(tuid, [])['vendor_id'],
            'subaccount_id': participant_data.get(tuid, [])['subaccount_id'],
            'participations': participant_data.get(tuid, [])['participations'],
            'performers': performer_data.get(tuid, [])['performers'],
            'publishers': publisher_data.get(tuid, [])['publishers']
        }
        for tuid in all_tuids
    }

    if WRITE_MASTER_OUTPUT:
        # Write the data to a JSON file
        write_var_to_json(contributor_data, CACHED_DATA_PATH)

    if errors:
        msg = f'Failed to process {len(errors)} rows for Performers. ' \
            f'Check the input file for errors.'
        log.error(msg)

        if ABORT_ON_FAIL:
            raise ValueError(msg)

    # Run the final processing
    # N.B Now that this is async, the data is returned encapsulated in a
    # Future, and must be grabbed with .result()

    # Multi-processes
    collected_data, collected_errors = process_data_async(
        process_contrib_data,
        contributor_data,
        pool_size=ASYNC_POOL_SIZE,
        chunk_size=ASYNC_CHUNK_SIZE
    )

    return collected_data, collected_errors, contributor_data


@log_runtime
def main():
    """Main function to run the script."""
    # Parse command-line arguments
    parser = argparse.ArgumentParser(
        description='Attach Metadata to Tracks and Releases via GraphQL.')
    parser.add_argument(
        '--filename',
        type=str,
        help='The filename to process')
    parser.add_argument(
        '--cached_filename',
        type=str,
        help='The filename of cached data to process')

    args = parser.parse_args()

    # Check the environment for required variables
    check_env(args.filename, args.cached_filename)

    # Init some loop vars
    cached_replay_data = None
    cached_replay_output_file = None
    cached_filename = None

    # Check for cached data
    if CACHED_DATA_PATH or args.cached_filename:
        # Make a log msg
        # Determine the filename source to use - prefer CLI
        if args.cached_filename:
            cached_filename = args.cached_filename
            msg = f'Using Cached Data Path from CLI: {cached_filename}'
        elif CACHED_DATA_PATH:
            msg = f'Using Cached Data Path from Config: {CACHED_DATA_PATH}'
            cached_filename = CACHED_DATA_PATH

        if not cached_filename:
            log.error('No cached data file specified. Exiting.')
            sys_exit(1)

        # Check the cached data file exists
        if not os.path.exists(cached_filename):
            log.error(
                f'Cached data file not found: {cached_filename}. Exiting.')
            sys_exit(1)  # Exit with error

        # Finally report the log msg
        log.info(msg)

        if MAP_TASK == 'SMITHSONIAN':
            # Process the cached data
            collected_data, collected_errors = \
                process_cached_smithsonian_contribs(cached_filename)
        if MAP_TASK == 'LOCALIZATIONS':
            # Process the cached data
            collected_data, collected_errors = \
                process_cached_localizations(cached_filename)
    else:
        if args.filename:
            input_file_name = args.filename
            log.info(f'Using filename from commandline: {input_file_name}')
        elif INPUT_FILE_NAME:
            input_file_name = INPUT_FILE_NAME
            log.info(f'Using config input filename: {input_file_name}')

        if not input_file_name:
            log.error('No input file specified. Exiting.')
            sys_exit(1)

        if MAP_TASK == 'SMITHSONIAN':
            # Process fresh data
            collected_data, collected_errors, cached_replay_data = \
                process_smithsonian_contribs(input_file_name)
        if MAP_TASK == 'LOCALIZATIONS':
            # Process fresh data
            collected_data, collected_errors = \
                process_localizations(input_file_name)

        if cached_replay_data:
            cached_replay_output_file = write_var_to_json(cached_replay_data)

    # Write output files to JSON
    data_output_file = write_var_to_json(collected_data)
    error_output_file = write_var_to_json(collected_errors)

    # Parse JSON data to a report
    report_output_file = \
        parse_output_data_to_report(data_output_file, MAP_TASK)

    if CREATE_ZIP_FILES:
        file_list = [
            data_output_file,
            error_output_file,
            report_output_file
        ]

        if cached_replay_data:
            file_list.append(cached_replay_output_file)

        if cached_filename:
            file_list.append(cached_filename)

        zipfile_name = zip_files(*file_list)

        log.info(f'Output files zipped to: {zipfile_name}')

        return zipfile_name
    return


def log_env():
    log.debug('Extending Signal Handlers: SIGINT, SIGTERM')
    log.info(f'Task Timeout: {TASK_TIMEOUT} seconds')
    log.info(f'Async Pool Size: {ASYNC_POOL_SIZE}')
    log.info(f'Async Chunk Size: {ASYNC_CHUNK_SIZE}')
    log.info(f'Logging Level: {LOGGER_LEVEL}')
    log.info(f'Log Frequency: <ITEM_LENGTH>/{LOG_FREQUENCY}')
    log.info(f'Request Delay: {REQUEST_DELAY} seconds')
    log.info(f'Task Delay: {TASK_DELAY} seconds')
    log.info(f'Jitter Factor: {JITTER_FACTOR}')
    log.info(f'Abort on Fail: {ABORT_ON_FAIL}')
    log.info(f'Write Master Output: {WRITE_MASTER_OUTPUT}')
    log.info(f'Input File Name: {INPUT_FILE_NAME}')
    log.info(
        f'Input File Starting Row Offset: {INPUT_FILE_STARTING_ROW_OFFSET}')
    log.info(f'Number of Rows: {NROWS}')
    log.info(f'Performers Count: {PERFORMER_COUNT}')
    log.info(f'Cached Data Path: {CACHED_DATA_PATH}')

    for field, value in FILTER_FIELDS.items():
        log.info(f'Filter Field: {field} = {value}')


if __name__ == '__main__':
    # Check that a task was specified
    if not MAP_TASK:
        log.error('No MAP_TASK specified. Exiting.')
        sys_exit(1)

    # Spit out the environment variables
    log_env()

    # Set up signal handling for aggressively graceful shutdown
    signal(SIGINT, signal_handler)
    signal(SIGTERM, signal_handler)

    # Run the main function
    main()

    # Exit explicitly
    sys_exit(0)