"""Logic for concurrency."""
import queue
from time import sleep

from typing import Callable, Dict, Iterable, List, Tuple, Union
from functools import partial

from pebble import concurrent

from connectors.logging import log
from constants.backoff import TASK_TIMEOUT
from constants.localization_definitions import (
    PRODUCT_LOCALIZATIONS_FIELD,
    TRACK_LOCALIZATIONS_FIELD,
    PARTICIPATIONS_FIELD,
)
from tasks.task_get_participant_id import task_get_participant_id
from tasks.task_set_track_contributions import task_set_track_contributions
from tasks.task_set_localizations import (
    task_set_track_localizations,
    task_set_product_localizations,
)
from tasks.task_get_product_from_graph_by_upc import \
    task_get_product_from_graph_by_upc
from utils.concurrency_utils import (
    chunk_data,
    prepare_errors_dict,
    process_response_data,
    init_async_process,
)
from utils.jitter import jitter

from config import (
    LOGGER_LEVEL,
    REQUEST_DELAY,
)


# @track_future
@concurrent.process(timeout=TASK_TIMEOUT)
def process_contrib_data(
        contributor_data: dict, task_id: int = None) -> tuple:
    """Run set_track_contributions on cached data. Errors are logged and
    returned.

    Args:
        contributor_data (dict): Formatted contributor data.
        task_id (int, optional): The task ID. Defaults to None.

    Returns:
        tuple (list, list): The processed data and any errors encountered.
    """
    log_frequency = init_async_process(contributor_data, task_id)

    # Continue standard processing
    contrib_len = len(contributor_data)
    # Loop through data df and execute saveTrack request in GQL
    count = 0

    collected_errors = []
    collected_data = []

    for tuid, contributors in contributor_data.items():
        count += 1

        # Reset data and errors
        track_data = {}
        errors = {}

        msg = f'Processing Track {count} of {contrib_len} ' \
            f'({(count/contrib_len)*100:.2f}%) ' \
            # f'Elaspse Time: {pd.Timestamp.now()-starttime}.....'

        if task_id:
            msg = f'Task {task_id}: {msg}'

        log.debug(msg)

        if LOGGER_LEVEL != 'DEBUG' and count % log_frequency == 0:
            log.info(msg)

        # Call the Updates!
        track_data, errors = task_set_track_contributions(
            tuid=tuid,
            vendor_id=contributors['vendor_id'],
            subaccount_id=contributors['subaccount_id'],
            participations=contributors['participations'],
            performers=contributors['performers'],
            publishers=contributors['publishers'],
            task_id=task_id or None
        )

        if errors:
            collected_errors.append(errors)
        if track_data:
            if track_data:
                for t in track_data:
                    collected_data.append(t)

        # Anti-hammer
        if REQUEST_DELAY:
            sleep_delay = \
                    REQUEST_DELAY + jitter(do_sleep=False, small_jitter=True)
            msg = f' (Sleeping {sleep_delay}s).'

            sleep(REQUEST_DELAY)

            if task_id:
                msg = f'Task {task_id}: {msg}'
            log.debug(msg)

    return collected_data, collected_errors


@concurrent.process(timeout=TASK_TIMEOUT)
def get_product_data_from_graph(
        upcs_vends_subs: list, task_id: int = None) -> dict:
    """Run get_product_by_upc() on product data. Errors are logged and
    returned.

    Args:
        upcs_by_vendor (list): A list of tuples containing the UPC, vendor ID,
            and subaccount ID.
        task_id (int, optional): The task ID. Defaults to None.

    Returns:
        list: The processed data.
        list: The collected errors.
    """
    log_frequency = init_async_process(upcs_vends_subs, task_id)

    # Init vars
    count = 0
    product = {}
    errors = {}

    collected_errors = []
    collected_data = []

    contrib_len = len(upcs_vends_subs)

    # graph_products_by_vendor_upc = {}
    for vendor_id, subaccount_id, upc in upcs_vends_subs:
        count += 1

        # graph_products_by_vendor_upc[vendor_id] = {}

        msg = f'Processing {count} of {contrib_len} ' \
            f'({(count/contrib_len)*100:.2f}%): {upc} on {vendor_id}'
        msg = subaccount_id and f'{msg} on {subaccount_id}' or msg
        msg = task_id and f'Task {task_id}: {msg}' or msg
        log.debug(msg)

        if LOGGER_LEVEL != 'DEBUG' and count % log_frequency == 0:
            log.info(msg)

        # log.info(
        #     f'Getting product data for UPC {upc} on Vendor {vendor_id}')

        # Get the product!
        product, errors = task_get_product_from_graph_by_upc(
            upc=upc,
            vendor_id=vendor_id,
            subaccount_id=subaccount_id,
            task_id=task_id
        )

        log.debug(f'Processed {upc} on {vendor_id}.')

        if errors:
            collected_errors.append(errors)
            errors = {}
        # if data:
        if product:
            collected_data.append(product)
            product = {}

        # Anti-hammer
        if REQUEST_DELAY:
            sleep_delay = \
                REQUEST_DELAY + jitter(do_sleep=False, small_jitter=True)
            msg = f' (Sleeping {sleep_delay}s).'

            sleep(REQUEST_DELAY)

            if task_id:
                msg = f'Task {task_id}: {msg}'
            log.debug(msg)

    return collected_data, collected_errors


@concurrent.process(timeout=TASK_TIMEOUT)
def get_participant_id_data(
        participant_data: list, task_id: int = None) -> dict:
    """Run get_or_create_label_participant() on contributor data. Errors are
    logged and returned.

    Args:
        participant_data (list): Formatted participant data. A list of dicts.
        task_id (int, optional): The task ID. Defaults to None.

    Returns:
        dict: The processed data.
        dict: The collected errors.
    """
    log_frequency = init_async_process(participant_data, task_id)

    # Continue standard processing
    contrib_len = len(participant_data)
    # Loop through data df and execute saveTrack request in GQL
    count = 0

    collected_errors = []
    collected_data = []

    for participant in participant_data:
        artist_name = participant['name']
        vendor_id = participant['vendor_id']
        subaccount_id = participant['subaccount_id']

        count += 1

        msg = f'Processing Participant {count} of {contrib_len} ' \
            f'({(count/contrib_len)*100:.2f}%) '
        if task_id:
            msg = f'Task {task_id}: {msg}'
        log.debug(msg)

        if LOGGER_LEVEL != 'DEBUG' and count % log_frequency == 0:
            log.info(msg)

        # Get the id!
        data, errors = task_get_participant_id(
            artist_name=artist_name,
            vendor_id=vendor_id,
            subaccount_id=subaccount_id,
            task_id=task_id or None
        )

        log.debug(f'Processed {artist_name} on {vendor_id}.')
        log.debug('Adding results to collected responses')

        if errors:
            collected_errors.append(errors)
            errors = {}
        if data:
            collected_data.append(data)
            data = {}

        # Anti-hammer
        if REQUEST_DELAY:
            sleep_delay = \
                REQUEST_DELAY + jitter(do_sleep=False, small_jitter=True)
            msg = f' (Sleeping {sleep_delay}s).'

            sleep(REQUEST_DELAY)

            if task_id:
                msg = f'Task {task_id}: {msg}'
            log.debug(msg)

    return collected_data, collected_errors


@concurrent.process(timeout=TASK_TIMEOUT)
def set_track_localizations_data(
        localization_data: dict, task_id: int = None) -> dict:
    """Run set_localizations() using localizaton data. Errors are
    logged and returned.

    Args:
        localization_data (dict): Formatted localization data.
        level (str): 'track' or 'release'.
        task_id (int, optional): The task ID. Defaults to None.

    Returns:
        dict: The processed data.
        dict: The collected errors.
    """
    log_frequency = init_async_process(localization_data, task_id)

    # Continue standard processing
    localization_len = len(localization_data)
    # Loop through data df and execute saveTrack request in GQL
    count = 0

    collected_errors = []
    collected_data = []

    for track_localization in localization_data:
        count += 1

        vendor_id = track_localization[0]
        subaccount_id = track_localization[1]
        tuid = track_localization[4]['tuid']
        localizations = track_localization[4][TRACK_LOCALIZATIONS_FIELD]
        participations = track_localization[4][PARTICIPATIONS_FIELD]

        msg = f'Processing Localization {count} of {localization_len} ' \
            f'({(count/localization_len)*100:.2f}%) '
        msg = task_id and f'Task {task_id}: {msg}' or msg
        log.debug(msg)

        if LOGGER_LEVEL != 'DEBUG' and count % log_frequency == 0:
            log.info(msg)

        # Get the id!
        localized_result, errors = task_set_track_localizations(
            tuid=tuid,
            localizations=localizations,
            participations=participations,
            vendor_id=vendor_id,
            subaccount_id=subaccount_id,
            task_id=task_id or None
        )

        msg = f'Processed localization for {tuid} on {vendor_id}.'
        if task_id:
            msg = f'Task {task_id}: {msg}'
        log.debug(msg)

        if errors:
            collected_errors.append(errors)
            errors = None
        if localized_result:
            collected_data.append(localized_result)
            localized_result = None

        # Anti-hammer
        if REQUEST_DELAY:
            sleep_delay = \
                REQUEST_DELAY + jitter(do_sleep=False, small_jitter=True)
            msg = f' (Sleeping {sleep_delay}s).'
            msg = task_id and f'Task {task_id}: {msg}' or msg
            log.debug(msg)

            sleep(REQUEST_DELAY)

    return collected_data, collected_errors


@concurrent.process(timeout=TASK_TIMEOUT)
def set_release_localization_data(
        localization_data: dict, task_id: int = None) -> dict:
    """Run set_localizations() using localizaton data. Errors are
    logged and returned.

    Args:
        localization_data (dict): Formatted localization data.
        task_id (int, optional): The task ID. Defaults to None.

    Returns:
        dict: The processed data.
        dict: The collected errors.
    """
    log_frequency = init_async_process(localization_data, task_id)

    # Continue standard processing
    localization_len = len(localization_data)
    # Loop through data df and execute saveTrack request in GQL
    count = 0

    collected_errors = []
    collected_data = []

    for release_localization in localization_data:
        count += 1

        vendor_id = release_localization[0]
        subaccount_id = release_localization[1]
        upc = release_localization[2]
        localizations = \
            release_localization[3][upc][PRODUCT_LOCALIZATIONS_FIELD]
        product_id = release_localization[3]['product_id']

        msg = f'Processing Localization {count} of {localization_len} ' \
            f'({(count/localization_len)*100:.2f}%) '
        msg = task_id and f'Task {task_id}: {msg}' or msg
        log.debug(msg)

        if LOGGER_LEVEL != 'DEBUG' and count % log_frequency == 0:
            log.info(msg)

        # Get the id!
        localized_result, errors = task_set_product_localizations(
            product_id=product_id,
            localizations=localizations,
            vendor_id=vendor_id,
            subaccount_id=subaccount_id,
            task_id=task_id or None
        )

        msg = f'Processed localization for {upc} on {vendor_id}.'
        if task_id:
            msg = f'Task {task_id}: {msg}'
        log.debug(msg)

        if errors:
            collected_errors.append(errors)
            errors = None
        if localized_result:
            collected_data.append(localized_result)
            localized_result = None

        # Anti-hammer
        if REQUEST_DELAY:
            sleep_delay = \
                REQUEST_DELAY + jitter(do_sleep=False, small_jitter=True)
            msg = f' (Sleeping {sleep_delay}s).'
            msg = task_id and f'Task {task_id}: {msg}' or msg
            log.debug(msg)

            sleep(REQUEST_DELAY)

    return collected_data, collected_errors


def process_data_async(
        func: Callable[[Dict, int], Tuple[Dict, Dict]],
        data: Iterable[Union[Dict, List]],
        pool_size: int = 1,
        chunk_size: int = 1) -> Tuple[Dict, Dict]:
    """Run data processing asynchronously using the provided function.

    Args:
        func (Callable[[Dict, int], Tuple[Dict, Dict]]): The function to
            process data.
        data (Iterable[Union[Dict, List]]): The data to process.
        pool_size (int, optional): The number of concurrent tasks to run.
            Defaults to 1.
        chunk_size (int, optional): The number of items to process per
        task. Defaults to 1.

    Returns:
        Tuple[Dict, Dict]: The processed data and any errors encountered.
    """
    task_queue = queue.Queue()

    finished_tasks = set()
    active_tasks = set()
    completed_task_count = 0

    errors = []
    collected_data = []
    collected_errors = []

    log.debug(f'Processing {len(data)} items.')
    log.debug(f'Chunking data into {chunk_size} items per task.')

    # Chunk the data efficiently using an iterator for list or dict
    outer_task_id = 0
    for chunk in chunk_data(data, chunk_size):
        outer_task_id += 1
        task_queue.put((outer_task_id, chunk))

    task_count = task_queue.qsize()

    # Define our queue processing callback inline, so it has outer scope access
    def on_task_done(task_id, future):
        """Handles task completion: retrieves the result and starts a new task
        if applicable.

        Args:
            task_id (int): The task ID.
            future (concurrent.Future): The completed future.

        Returns:
            None
        """
        log.debug(f'Task {task_id}: `on_task_done()` callback called .')

        if future.exception():
            log.error(f'Task {task_id}: Exception: {repr(future.exception())}')
            try:
                exception = future.exception()
                log.error(f'Task {task_id}: {exception.traceback}')
            except AttributeError:
                log.error(f'Task {task_id}: Traceback not available.')
        elif future.cancelled():
            log.error(f'Task {task_id}: Cancelled')
        elif future.done():
            log.debug(
                f'Task {task_id}: -------------------------------------- Done')
            log.debug(
                f'Task {task_id}: '
                f'Removing task {task_id} from active tasks')
        else:
            log.error(f'Task {task_id}: Done, but not finished! This is bad.')

        finished_tasks.add((task_id, future))
        active_tasks.discard((task_id, future))

        try:
            new_task_id, next_data_chunk = task_queue.get(block=False)
        except queue.Empty:
            return future

        if new_task_id == task_id:
            return future

        # Start a new task if items are left
        if new_task_id not in [t_id for t_id, _ in active_tasks]:
            log.debug(
                f'Task {task_id}: Starting task {new_task_id} of {task_count} '
                f'with payload size: {len(next_data_chunk)}')

            # Jitter to avoid hammering the server
            jitter()

            # Start a new task
            new_future = func(next_data_chunk, new_task_id)

            # Create and add a callback with the task_id bound
            callback_with_param = partial(on_task_done, new_task_id)
            new_future.add_done_callback(callback_with_param)

            # Add task and id to the active set
            active_tasks.add((new_task_id, new_future))  # Add to active tasks
        else:
            log.warning(
                f'Task {task_id}: Task {new_task_id} is already active')

        return future

    if task_count == 0:
        log.warning(f'No tasks to process for {func.__name__}.')
        return {}, {}

    log.info(f'Queue size: {task_count}')
    log.info(
        f'Starting initial batch of {pool_size} tasks with chunk size: '
        f'{chunk_size}')

    for _ in range(min(pool_size, task_count)):
        try:
            # Get the next item from the queue
            outer_task_id, next_data_chunk = task_queue.get(block=False)

        except queue.Empty:
            break

        msg = f'Starting task {outer_task_id} of {task_count} with payload ' \
            f'size: {len(next_data_chunk)}'
        log.debug(msg)

        # Start a task
        future = func(next_data_chunk, outer_task_id)

        # Create and add a callback with the task_id bound
        callback_with_param = partial(on_task_done, outer_task_id)
        future.add_done_callback(callback_with_param)  # This will always run

        # Add task and id to the active set
        active_tasks.add((outer_task_id, future))

    log.info('Task watchdog is starting.')
    log.info(f'Waiting for {task_count} tasks to complete.')
    log.info(f'{completed_task_count} tasks completed.')

    # Wait for all tasks to complete
    while task_count > 0:
        # NOTE: `finished_tasks` is kept filled from `on_task_done()` callback.
        loop_tasks = finished_tasks.copy()

        # Loop over the finished tasks
        for outer_task_id, future in loop_tasks:
            response = None
            error = None
            if future.exception():
                log.error(
                    f'Task {outer_task_id}: Exception: '
                    f'{future.exception()}')
                traceback = None
                try:
                    traceback = future.exception().traceback

                except AttributeError:
                    pass  # no op
                error = {
                        'outer_task_id': outer_task_id,
                        'error': repr(future.exception()),
                        'calling_function': 'process_data_async()',
                    }
                if traceback:
                    error['traceback'] = traceback

            elif future.cancelled():
                log.error(f'Task {outer_task_id}: Cancelled')
                error = {
                    'outer_task_id': outer_task_id,
                    'error': repr(future.exception()),
                    'calling_function': 'process_data_async()',
                }
            else:
                log.info(
                    f'Task {outer_task_id}: Completed. Processing '
                    'results.')
                try:
                    response = future.result(timeout=0)
                except TimeoutError as t:
                    log.warning(
                        f'Task {outer_task_id}: TimeoutError.')
                    error = {
                        'outer_task_id': outer_task_id,
                        'error': repr(t),
                        'calling_function': 'process_data_async()',
                    }

            log.debug(f'Collecting results for task {outer_task_id}')

            # Merge the data and errors dicts
            if error:  # If there was an error in this task runner
                # Create an error dict for the task
                errors = prepare_errors_dict(
                    errors={},
                    category='task',
                    variable=outer_task_id)
                errors['task'][outer_task_id].append(error)

                # Add the error to the full run's collected errors
                collected_errors.append(errors)
            if response:
                if isinstance(response[0], list):
                    collected_data.extend(response[0])
                elif isinstance(response[0], dict):
                    collected_data.append(response[0])
                else:
                    log.warning(
                        f'Task {outer_task_id}: Unexpected result type: '
                        f'{type(response[0])}. Expected list or dict.')
                    collected_data.append(response[0])

                if isinstance(response[1], list):
                    collected_errors.extend(response[1])
                elif isinstance(response[1], dict):
                    collected_errors.append(response[1])
                else:
                    log.warning(
                        f'Task {outer_task_id}: Unexpected error type: '
                        f'{type(response[1])}. Expected list or dict.')
                    collected_errors.append(response[1])

                response = None

            task_count -= 1
            completed_task_count += 1

            finished_tasks.discard((outer_task_id, future))

    return process_response_data(collected_data, collected_errors)