import multiprocessing as mp
import itertools as it

from functools import partial
from os import PathLike
from typing import Iterable, Callable, Union, TextIO, Any


def write(file_handle: TextIO, uid_list: list[list[str]]) -> None:
    """
    writes the uids in uid_list to file_handle and flushes it afterwards
    
    :param file_handle:   IO stream generated by functions like open
    :param uid_list:      list of lists of accession uid pairs as strings
    
    :return:              None
    """
    for uid_map in uid_list:
        file_handle.write(
            '\t'.join(uid_map) + '\n'
        )

    file_handle.flush()


def write_to_file_multiprocess(filename: Union[PathLike, str], uid_list: list[list[str]], filelock: mp.Manager().Lock) -> None:
    """
    writes the given list of uids to a file that is accesses from a
    concurrent process.
    
    :param filename:   path to the file to write the list to
    :param uid_list:   list of lists of accession uid pairs as strings
    :param filelock:   Lock object that can be used with multiprocessing imap
    
    :return:           None
    """
    with filelock:
        with open(filename, 'a') as outfile:
            write(outfile, uid_list)
            
            
def write_to_file_singleprocess(filename: Union[PathLike, str], uid_list: list[list[str]], *args):
    """
    writes the given list of uids to a file. single threaded equivalent to write_to_file_multiprocess
    
    :param
    """
    with open(filename, 'a') as outfile:
        write(outfile, uid_list)
        

def multiprocess_map(chunks: Iterable, func: Callable, n_processes: int, **kwargs) -> None:
    """
    takes an iterable containing chunks of a whole and a function to process these chunks and uses n_processes to do it
    
    :param chunks:        Iterable containing chunks of a whole that need to be processed concurrently
    :param func:          function to process the chunks with
    :param n_processes:   number of concurrent processes to use for processing
    :param **kwargs:      any keyword arguments to pass to func
    
    :return:              None
    """
    with mp.Pool(n_processes) as p:
        # lock needs to be a managed one otherwise passing it to the pool will fail
        filelock = mp.Manager().Lock()
        
        map_function = partial(func, filelock = filelock, **kwargs)

        # this needs to be invoked by iterating over it
        results = [result for result in p.imap(map_function, chunks)]
    
    return results
    
            
def singleprocess_map(chunks: Iterable, func: Callable, **kwargs) -> None:
    """
    processes chunks with func using map. single threaded equivalent to multiprocess_map
    
    :param chunks:        Iterable containing chunks of a whole that need to be processed concurrently
    :param func:          function to process the chunks with
    :param **kwargs:      any keyword arguments to pass to func
    
    :return:              None
    """
    map_function = partial(func, **kwargs)
    
    results_iterable = map(map_function, chunks)
    # this needs to be invoked by iterating over it
    return list(results_iterable)


def process_data_in_chunks(iterable: Iterable, func: Callable, chunksize: int = 5000, n_processes: int = 1, **kwargs) -> Any:
    """
    uses func to process the given iterable in chunks of size chunksize possibly concurrently
    
    :param iterable:      iterable containing the data to be processed with func
    :param func:          function that processes the contents of iterable
    :param chunksize:     size of the indiviually processed chunks of the input iterable
    :param n_processes:   number of processes to use for mapping if n_processes > 1 this will be done concurrently using multiprocessing.imap
    :param **kwargs:      any keyword arguments that need to be passed to func
    
    :return:              anything that is returned by func
    """
    chunks = it.batched(
        iterable,
        n = chunksize
    )

    if n_processes > 1:
        results = multiprocess_map(
            chunks,
            func,
            n_processes,
            **kwargs
        )

    else:
        results = singleprocess_map(
            chunks,
            func,
            **kwargs
        )

    return results