from collections.abc import Callable
from dataclasses import dataclass, field
from functools import cached_property
from multiprocessing import get_context
from multiprocessing.context import SpawnContext, SpawnProcess
from multiprocessing.queues import Queue
from multiprocessing.synchronize import Barrier
from typing import Any, Literal, Self

from common import config
from common._config import DumpedRegistry

type JobResult[T] = T | None
type JobCallable[**P, T] = Callable[P, JobResult[T]]


@dataclass(eq=True, frozen=True)
class Job[**P, T]:
    func: JobCallable[P, T]
    config: DumpedRegistry
    args: P.args = ()
    kwargs: P.kwargs = field(default_factory=dict)

    def __call__(self) -> JobResult[T]:
        config.load_from_dump(self.config)
        return self.func(*self.args, **self.kwargs)


def _pool_runner[T](
    job_queue: Queue[Job[Any, T] | None],
    result_queue: Queue[T | None],
    barrier: Barrier,
) -> None:
    while True:
        job = job_queue.get()
        if job is None:
            barrier.wait()
            result_queue.put(None)
            break
        result = job()
        if result is not None:
            result_queue.put(result)


@dataclass
class Pool[T]:
    processes: int
    start_upfront: bool = False
    children: dict[int, SpawnProcess] = field(default_factory=dict, init=False)
    is_done: bool = field(default=False, init=False)

    def __post_init__(self) -> None:
        if self.start_upfront:
            for _ in range(self.processes):
                self._spawn_child()

    @cached_property
    def context(self) -> SpawnContext:
        return get_context("spawn")

    @cached_property
    def job_queue(self) -> Queue[Job[Any, T] | None]:
        return self.context.Queue()

    @cached_property
    def result_queue(self) -> Queue[T | None]:
        return self.context.Queue()

    @cached_property
    def barrier(self) -> Barrier:
        return self.context.Barrier(self.processes)

    def _spawn_child(self) -> None:
        process = self.context.Process(
            target=_pool_runner, args=(self.job_queue, self.result_queue, self.barrier)
        )
        process.start()
        assert process.pid is not None
        self.children[process.pid] = process

    def submit[**P](
        self, func: JobCallable[P, T], *args: P.args, **kwargs: P.kwargs
    ) -> None:
        assert not self.is_done, "Cannot submit jobs to pool after marked done"

        if len(self.children) < self.processes:
            self._spawn_child()

        self.job_queue.put(Job(func, config.dump(), args, kwargs))

    def done(self) -> None:
        for _i in range(self.processes):
            self.job_queue.put(None)

        while len(self.children) < self.processes:
            self._spawn_child()

    def get(self, timeout: float) -> T | None:
        return self.result_queue.get(timeout=timeout)

    def terminate(self) -> None:
        for p in self.children.values():
            try:
                p.terminate()
            except:
                continue

    def join(self) -> None:
        for p in self.children.values():
            try:
                p.join()
            except:
                continue

    def close(self) -> None:
        for p in self.children.values():
            try:
                p.close()
            except:
                continue

    def __enter__(self) -> Self:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any,
    ) -> Literal[False]:
        self.job_queue.cancel_join_thread()
        self.result_queue.cancel_join_thread()
        self.terminate()
        self.join()
        return False
