
import multiprocessing

import logging
log = logging.getLogger(__name__)

MAX_RUNTIME = 10000  # seconds
MAX_RETRIES = 3
MAX_PARALLEL = max(1, multiprocessing.cpu_count() - 1)

# Global variable to hold the user-provided worker function at module level
GLOBAL_WORKER = None

def _safe_wrapper(arg):
    """
    Module-level wrapper that calls GLOBAL_WORKER(arg) and returns (ok, payload).
    Needed for pickling in multiprocessing.
    """
    try:
        result = GLOBAL_WORKER(arg)
        return (True, result)
    except Exception as e:
        log.warning(f"Task {arg!r} failed: {e}")
        return (False, e)

# Module-level safe wrapper for pickling
# It expects an (index, subtask) tuple and applies GLOBAL_WORKER to subtask
def _safe_wrapper(item):
    idx, task = item
    try:
        result = GLOBAL_WORKER(task)
        return (idx, True, result)
    except Exception as e:
        log.warning(f"Task at index {idx} ({task!r}) failed: {e}")
        return (idx, False, e)

# Orchestrator with per-task timeout and retries, works with arbitrary subtask objects
def run_all_with_retries(worker_func, subtasks, max_retries=MAX_RETRIES, task_timeout=None, num_processes=None):
    """
    Run tasks in parallel, retrying failures until all succeed or max_retries is reached.
    Supports per-task timeout via task_timeout (seconds).

    Args:
        worker_func (callable): function to apply to each subtask
        subtasks (list): inputs for the worker; can be any objects
        max_retries (int|None): max attempts per task (None = infinite)
        task_timeout (float|None): seconds to wait for each individual task

    Returns:
        list: results in same order as subtasks
    """
    global GLOBAL_WORKER
    GLOBAL_WORKER = worker_func

    log.info("Starting multiprocess run for tile prediction")
    # Represent pending tasks as list of (index, task)
    pending = list(enumerate(subtasks))
    attempts = {idx: 0 for idx, _ in pending}
    successes = {}

    if num_processes is None:
        max_parallel = max(1, multiprocessing.cpu_count() - 1)
    else:
        max_parallel = max(1, num_processes)
        
    while pending:
        log.debug(f"launching {len(pending)} tasks with {max_parallel} processes")
        pool = multiprocessing.Pool(processes=max_parallel, maxtasksperchild=2)
        try:
            # launch each pending task via apply_async
            async_map = {idx: pool.apply_async(_safe_wrapper, args=(item,)) for idx, item in enumerate(pending)}
            pool.close()

            results_map = {}
            # collect results per task with timeout
            for idx, async_res in async_map.items():
                try:
                    task_idx, ok, payload = async_res.get(timeout=task_timeout)
                except multiprocessing.TimeoutError:
                    task_idx = pending[idx][0]
                    ok = False
                    payload = TimeoutError(f"Timeout on {pending[idx][1]}")
                    log.error(f"Task at index {task_idx} ({pending[idx][1]!r}) timed out after {task_timeout} seconds")
                results_map[task_idx] = (ok, payload)
        finally:
            pool.terminate()
            pool.join()

        # Determine which to retry
        new_pending = []
        for idx, task in pending:
            attempts[idx] += 1
            ok, payload = results_map[idx]
            if ok:
                successes[idx] = payload
            else:
                if max_retries is not None and attempts[idx] >= max_retries:
                    log.error("Task at index %d (%r) failed %d times, giving up", idx, task, attempts[idx])
                    raise RuntimeError(f"Task {task!r} failed {attempts[idx]} times")
                new_pending.append((idx, task))

        if new_pending:
            log.info("Retrying %d failed tasks: %s", len(new_pending), [t for _, t in new_pending])
        pending = new_pending

    # All tasks succeeded; assemble results in original order
    return [successes[i] for i in range(len(subtasks))]



# ### timeout for all tasks
# # Orchestrator: retries failed tasks until all succeed
# def run_all_with_retries(worker_func, subtasks, max_retries=None, timeout=12000):
#     """
#     Run tasks in parallel using worker_func, retrying failures until all succeed or until max_retries is reached.
# 
#     Args:
#         worker_func (callable): the function to apply to each subtask
#         subtasks (list): list of inputs for the worker
#         max_retries (int|None): maximum attempts per task (None = infinite)
# 
#     Returns:
#         list: results in the same order as `subtasks`
#     """
#     global GLOBAL_WORKER
#     GLOBAL_WORKER = worker_func
# 
#     pending = list(subtasks)
#     attempts = {arg: 0 for arg in subtasks}
#     successes = {}
# 
#     max_parallel = max(1, multiprocessing.cpu_count() - 1)
#     chunksize = 1 #max(1, len(subtasks) // (max_parallel * 4))
# 
#     while pending:
#         with multiprocessing.Pool(processes=max_parallel, maxtasksperchild=10) as pool:
#             async_res = pool.map_async(_safe_wrapper, pending, chunksize=chunksize)
#             try:
#                 raw = async_res.get(timeout=timeout)
#             except multiprocessing.TimeoutError:
#                 logging.error("Pool hung; retrying batch of %d tasks", len(pending))
#                 pool.terminate()
#                 continue
# 
#         failed = []
#         for arg, (ok, payload) in zip(pending, raw):
#             attempts[arg] += 1
#             if ok:
#                 successes[arg] = payload
#             else:
#                 if max_retries is not None and attempts[arg] >= max_retries:
#                     logging.error(
#                         "Task %r failed %d times, giving up", arg, attempts[arg]
#                     )
#                     raise RuntimeError(
#                         f"Task {arg!r} failed after {attempts[arg]} attempts"
#                     )
#                 failed.append(arg)
# 
#         if failed:
#             logging.info("Retrying %d failed tasks: %s", len(failed), failed)
#         pending = failed
# 
#     return [successes[arg] for arg in subtasks]