# it's really messed up on a number of levels that doing this is necessary
from multiprocessing import Process, Pipe
from multiprocessing.connection import Connection
from threading import Thread, Lock, Condition
from queue import Queue
from typing import Any, Iterable, List, Tuple, Optional
import os

class _ExceptionContainer:
    exn: Exception
    def __init__(self, exn: Exception):
        self.exn = exn

class WorkItem:
    _callable: Any
    _arg: Any

    def __init__(self, callable: Any, arg: Any):
        self._callable = callable
        self._arg = arg

    def __call__(self) -> Any:
        return self._callable(self._arg)

class WorkList:
    _inputs: List
    _outputs: List
    _position: int
    _lock: Lock

    def __init__(self, inputs: Iterable):
        self._inputs = list(inputs)
        self._outputs = [None] * len(self._inputs)
        self._position = 0
        self._number_done = 0
        self._lock = Lock()
        self._cv = Condition(self._lock)
    
    def __iter__(self):
        return self
    
    def __next__(self) -> Tuple[int, Any]:
        with self._lock:
            if self._position < len(self._inputs):
                i = self._position
                self._position += 1
                return i, self._inputs[i]
        raise StopIteration
    
    def set_output(self, i: int, output: Any) -> None:
        with self._lock:
            self._outputs[i] = output
            self._number_done += 1
            if self._number_done == len(self._inputs):
                self._cv.notify_all()

    def cancel(self) -> None:
        with self._lock:
            self._number_done = len(self._inputs)
            self._cv.notify_all()

    def wait_until_done(self) -> None:
        with self._lock:
            while self._number_done < len(self._inputs):
                self._cv.wait()

    def get_outputs(self) -> List:
        for o in self._outputs:
            if isinstance(o, _ExceptionContainer):
                raise o.exn
        return self._outputs

def _child_entry(parent_conn: Connection, child_conn: Connection) -> None:
    parent_conn.close()
    while True:
        try:
            item = child_conn.recv()
        except EOFError:
            break
        if item is None:
            break
        i, work = item
        try:
            output, exn = work(), None
        except Exception as e:
            output, exn = None, e
        child_conn.send((i, output, exn))
    child_conn.close()

class Child:
    _proc: Process

    _parent_conn: Connection
    _child_conn: Connection

    _in_queue: Queue

    _to_courier: Thread
    _from_courier: Thread

    _current_worklist: Optional[WorkList]

    def __init__(self):
        self._in_queue = Queue()
        self._parent_conn, self._child_conn = Pipe(duplex=True)
        self._proc = Process(target=_child_entry, args=(self._parent_conn, self._child_conn,), daemon=True)
        self._proc.start()
        self._child_conn.close()
        self._to_courier = Thread(target=self._to_courier_entry, daemon=True)
        self._to_courier.start()
        self._from_courier = Thread(target=self._from_courier_entry, daemon=True)
        self._from_courier.start()
        self._current_worklist = None

    def _to_courier_entry(self) -> None:
        try:
            while True:
                worklist = self._in_queue.get()
                if worklist is None:
                    self._parent_conn.send(None)
                    break
                self._current_worklist = worklist
                for i, work in worklist:
                    self._parent_conn.send((i, work))
                worklist.wait_until_done()
                self._current_worklist = None
        except IOError:
            pass

    def _from_courier_entry(self) -> None:
        try:
            while True:
                try:
                    i, output, exn = self._parent_conn.recv()
                except EOFError:
                    break
                if exn is not None:
                    output = _ExceptionContainer(exn)
                self._current_worklist.set_output(i, output)
        except IOError:
            pass

    def consume_worklist(self, worklist: WorkList) -> None:
        self._in_queue.put(worklist)

    def close(self) -> None:
        self._in_queue.put(None)

    def join(self) -> None:
        self._to_courier.join()
        self._from_courier.join()
        self._proc.join()

    def terminate(self) -> None:
        self._in_queue.put(None)
        self._parent_conn.close()
        self._proc.terminate()
        if self._current_worklist is not None:
            self._current_worklist.cancel()

class Pool:
    def __init__(self, processes=None, initializer=None, initargs=None, maxtasksperchild=None, context=None):
        if processes is None:
            self.num_processes = os.cpu_count()
        else:
            if processes < 0:
                raise ValueError("processes must be a positive integer")
            self.num_processes = processes


        if initializer:
            raise NotImplementedError("initializer not implemented")

        if initargs:
            raise NotImplementedError("initargs not implemented")

        if maxtasksperchild:
            raise NotImplementedError("maxtasksperchild not implemented")

        if context:
            raise NotImplementedError("context not implemented")

        self._closed = False

        assert self.num_processes > 0, "Must have at least one process"

        self.children = [Child() for _ in range(self.num_processes)]


    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        self.join()
        self.terminate()

    def __del__(self):
        self.terminate()

    def close(self):
        if not self._closed:
            for c in self.children:
                c.close()
            self._closed |= True

    def join(self):
        assert self._closed, "Must close before joining"
        for c in self.children:
            c.join()

    def terminate(self):
        for c in self.children:
            c.terminate()
        self._closed |= True

    def map(self, func, iterable, chunksize=None, callback=None, error_callback=None) -> List:
        assert not self._closed, "Cannot submit tasks after closing"
        worklist = WorkList(WorkItem(func, x) for x in iterable)
        for c in self.children:
            c.consume_worklist(worklist)
        worklist.wait_until_done()
        return worklist.get_outputs()