from concurrent.futures import ProcessPoolExecutor

import pandas as pd
import numpy as np

from tqdm import tqdm

from graph_based_intrusion_detection.utils.state_merging import merge_states


def multi_process_packets(packets: pd.DataFrame,
                          global_parameterless_processing_function: callable,
                          merging_function: callable,
                          num_processes: int = 4,
                          chunk_size_factor: float = 1,
                          chunk_size: int = None) -> dict:
    """
    Process a set of packets using multiple processes
    :param packets: packets to process
    :param global_parameterless_processing_function: function to apply to the packets, must be parameterless and global
    :param merging_function: function to merge the states
    :param num_processes: number of processes to use
    :param chunk_size_factor: factor to use to determine the chunk size, cannot be used with chunk_size, can reduce memory usage
    :param chunk_size: size of chunks to split the packets into, if None, the packets are split into num_processes chunks, cannot be used with chunk_size_factor
    :return: updated state
    """

    # chunk packets
    if chunk_size is not None:
        packet_chunks = np.array_split(packets, len(packets) // chunk_size)
    else:
        packet_chunks = np.array_split(packets, int(num_processes * (1 / chunk_size_factor)))

    # create executor with tqdm
    with ProcessPoolExecutor(max_workers=num_processes) as executor:
        states = list(
            tqdm(executor.map(global_parameterless_processing_function, packet_chunks), total=len(packet_chunks)))

    # merge states
    merged_state = states[0]
    for state in states[1:]:
        merged_state = merging_function(merged_state, state)
    return merged_state


def process_packet(packet: dict | pd.Series,
                   state: dict,
                   processing_functions: list[callable]) -> dict:
    """
    Digest a packet and update the state
    :param packet: packet to digest
    :param state: state of the model
    :param processing_functions: functions to apply to the packet
    :return: updated state
    """

    for processing_function in processing_functions:
        state = processing_function(packet, state)

    return state


def process_packets(packets: pd.DataFrame | list,
                    processing_functions: list[callable],
                    state=None,
                    verbose: bool = False) -> dict:
    """
    Process a set of packets
    :param packets: packets to process
    :param processing_functions: functions to apply to the packets
    :param state: state of the model
    :param verbose: if True, print progress
    :return: updated state
    """

    iterator = packets.to_dict("records") if isinstance(packets, pd.DataFrame) else packets
    if verbose:
        for packet in tqdm(iterator, total=len(packets)):
            state = process_packet(packet, state, processing_functions)
    else:
        for packet in iterator:
            state = process_packet(packet, state, processing_functions)

    return state
