"""
Copyright 2025 Metaist LLC.

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""

from __future__ import annotations

from dataclasses import dataclass
from operator import attrgetter
from os import cpu_count
from platform import system
from queue import Queue as ThreadSafeQueue
from threading import Thread
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Generic,
    Iterable,
    Iterator,
    Literal,
    Sequence,
    TypeVar,
    Union,
)

from multiprocess import (
    Process,  # type: ignore[reportAttributeAccessIssue]
    Queue,  # type: ignore[reportAttributeAccessIssue]
)

__all__ = (
    "END_MSG",
    "IS_MACOS",
    "NUM_CPUS",
    "NUM_THREADS",
    "Context",
    "ContextName",
    "Msg",
    "MsgQ",
    "Q",
    "Task",
    "Worker",
    "mapq",
    "run",
    "run_thread",
)

T = TypeVar("T")
R = TypeVar("R")

Task = Callable[..., R]
"""Task function signature (any `Callable`)."""

Context = Union[Process, Thread]
"""Execution contexts (`Process`, `Thread`)."""

ContextName = Literal["process", "thread"]
"""Execution context names (`"process"`, `"thread"`)."""


@dataclass
class Msg(Generic[T]):
    """Message for a queue."""

    data: T
    """Message data to be transmitted."""

    kind: str = ""
    """Optional marker of message type."""

    order: int = 0
    """Optional ordering of messages."""


# NOTE: The python `queue.Queue` is not properly a generic.
# See: https://stackoverflow.com/a/48554601
if TYPE_CHECKING:  # pragma: no cover
    MsgQ = Union[Queue[Msg], ThreadSafeQueue]  # pylint: disable=unsubscriptable-object
else:
    MsgQ = Queue

END_MSG: Msg = Msg(data=None, kind="END")
"""Message that indicates no future messages will be sent."""

## Hardware-Specific Information ##

NUM_CPUS: int = cpu_count() or 1
"""Number of CPUs on this machine."""

NUM_THREADS: int = min(32, NUM_CPUS + 4)
"""Default number of threads (up to 32).

See: [CPython's default for this value][1].

[1]: https://github.com/python/cpython/blob/a635d6386041a2971cf1d39837188ffb8139bcc7/Lib/concurrent/futures/thread.py#L142
"""

IS_MACOS: bool = system().lower().startswith("darwin")
"""`True` if we're running on MacOS.

Currently, we only use this value for testing, but there are certain features that
do not work properly on MacOS.

See: [Example of MacOS-specific issues][1].

[1]: https://github.com/python/cpython/blob/c5b670efd1e6dabc94b6308734d63f762480b80f/Lib/multiprocessing/queues.py#L125
"""


class Worker:
    """A function running in a `Process` or `Thread`."""

    _worker: Context
    """Execution context."""

    @staticmethod
    def process(task: Task, *args: Any, **kwargs: Any) -> Worker:
        """Create a `Process`-based `Worker`.

        Args:
            task (Task): function to run
            *args (Any): additional arguments to `task`
            **kwargs (Any): additional keyword arguments to `task`

        Returns:
            Worker: wrapped worker.
        """
        # NOTE: On MacOS, python 3.8 switched the default method
        # from "fork" to "spawn" because fork is considered dangerous.
        # Some posts say "forkserver" should be ok.
        # See:  https://bugs.python.org/issue?@action=redirect&bpo=33725
        #
        # if IS_MACOS:
        #     ctx = get_context("forkserver")  # noqa: ERA001
        # else:  # noqa: ERA001
        #     ctx = get_context()  # noqa: ERA001
        return Worker(Process(daemon=True, target=task, args=args, kwargs=kwargs))

    @staticmethod
    def thread(task: Task, *args: Any, **kwargs: Any) -> Worker:
        """Create a `Thread`-based `Worker`.

        Args:
            task (Task): function to run
            *args (Any): additional arguments to `task`
            **kwargs (Any): additional keyword arguments to `task`

        Returns:
            Worker: wrapped worker.
        """
        return Worker(Thread(daemon=False, target=task, args=args, kwargs=kwargs))

    def __init__(self, context: Context):
        """Construct a worker from a context.

        Args:
            context (Context): a `Process` or a `Thread`
        """
        self._worker = context
        self._worker.start()

    def __getattr__(self, name: str) -> Any:
        """Delegate properties to the underlying task.

        Args:
            name (str): attribute name

        Returns:
            Any: attribute from the task
        """
        return getattr(self._worker, name)


class Q(Generic[T]):
    """Simple message queue."""

    _q: MsgQ
    """Wrapped queue."""

    _cache: list[Msg] | None = None
    """Cache of queue messages when calling `.items(cache=True)`."""

    def __init__(self, kind: ContextName = "process"):
        """Construct a queue wrapper.

        Args:
            kind (ContextName, optional): If `"thread"`, construct a lighter-weight
                `Queue` that is thread-safe. Otherwise, construct a full
                `multiprocess.Queue`. Defaults to `"process"`.
        """
        if kind == "process":
            self._q = Queue()
        elif kind == "thread":
            self._q = ThreadSafeQueue()
        else:  # pragma: no cover
            raise ValueError(f"Unknown queue type: {kind}")

    def __getattr__(self, name: str) -> Any:
        """Delegate properties to the underlying queue.

        Args:
            name (str): name of the attribute to access

        Returns:
            Any: attribute from the queue
        """
        return getattr(self._q, name)

    def __iter__(self) -> Iterator[Msg[T]]:
        """Iterate over messages in a queue until `END_MSG` is received.

        Yields:
            Iterator[Msg]: iterate over messages in the queue
        """
        return self.iter()

    def iter(self, timeout: float | None = None) -> Iterator[Msg[T]]:
        """Iterate over messages in a queue until `END_MSG` is received.

        Yields:
            Iterator[Msg]: iterate over messages in the queue
        """
        while True:
            msg = self._q.get(block=True, timeout=timeout)
            if msg.kind == END_MSG.kind:
                # We'd really like to put the `END_MSG` back in the queue
                # to prevent reading past the end, but in practice
                # this often creates an uncatchable `BrokenPipeError`.
                # q.put(END_MSG)  # noqa: ERA001
                break
            yield msg

    def items(self, *, cache: bool = False, sort: bool = False) -> Iterator[Msg[T]]:
        """End a queue and read all the current messages.

        Args:
            cache (bool, optional): if `True`, cache the messages. This allows you
                to call this method multiple times to get the same messages.
                Defaults to `False`.

            sort (bool, optional): if `True` messages are sorted by `Msg.order`.
                Defaults to `False`.

        Yields:
            Iterator[Msg]: iterate over messages in the queue
        """
        if cache:
            if self._cache is None:  # need to build a cache
                self.end()
                self._cache = list(self.sorted() if sort else self)
            return iter(self._cache)

        # not cached
        self.end()
        return self.sorted() if sort else iter(self)

    def sorted(self, start: int = 0) -> Iterator[Msg[T]]:
        """Iterate over messages sorted by `Msg.order`.

        NOTE: `Msg.order` must be incremented by one for each message.
        If there are any gaps, messages after the gap won't be yielded
        until the end.

        Args:
            start (int, optional): initial message number. Defaults to `0`.

        Yields:
            Iterator[Msg]: message yielded in the correct order
        """
        prev = start - 1
        key = attrgetter("order")
        waiting: list[Msg] = []
        for item in self:
            if not waiting and key(item) == prev + 1:
                prev += 1
                yield item
                continue

            # items came out of order
            waiting.append(item)
            waiting.sort(key=key, reverse=True)  # sort in-place for performance
            while waiting and key(waiting[-1]) == prev + 1:
                prev += 1
                yield waiting.pop()

        # generator ended; yield any waiting items
        while waiting:
            yield waiting.pop()

    def put(self, data: T | Msg[T], *, kind: str = "", order: int = 0) -> Q:
        """Put a message on the queue.

        Args:
            data (Any, optional): message data. Defaults to `None`.

            kind (str, optional): kind of message. Defaults to `""`.

            order (int, optional): message order. Defaults to `0`.

        Returns:
            Self: self for chaining
        """
        if isinstance(data, Msg):
            self._q.put(data)
        else:
            self._q.put(Msg(data=data, kind=kind, order=order))
        return self

    def end(self) -> Q:
        """Add the `END_MSG` to indicate the end of work.

        Returns:
            Self: self for chaining
        """
        self._q.put(END_MSG)
        return self

    def stop(self, workers: Worker | Sequence[Worker]) -> Q:
        """Use this queue to notify workers to end and wait for them to join.

        Args:
            workers (Worker, Sequence[Worker]): workers to wait for

        Returns:
            Self: self for chaining
        """
        _workers = [workers] if isinstance(workers, Worker) else workers

        for _ in range(len(_workers)):
            self.end()

        for task in _workers:
            task.join()

        return self


def run(task: Task, *args: Any, **kwargs: Any) -> Worker:
    """Run a function as a subprocess.

    Args:
        task (Task): function to run in each subprocess

        *args (Any): additional positional arguments to `task`.

        **kwargs (Any): additional keyword arguments to `task`.

    Returns:
        Worker: worker started in a subprocess

    .. changed:: 2.0.4
       This function now returns a `Worker` instead of a `Process`.
    """
    return Worker.process(task, *args, **kwargs)


def run_thread(task: Task, *args: Any, **kwargs: Any) -> Worker:
    """Run a function as a thread.

    Args:
        task (Task): function to run in each thread

        *args (Any): additional positional arguments to `task`.

        **kwargs (Any): additional keyword arguments to `task`.

    Returns:
        Worker: worker started in a thread

    .. changed:: 2.0.4
       This function now returns a `Worker` instead of a `Thread`.
    """
    return Worker.thread(task, *args, **kwargs)


def mapq(
    task: Callable[..., Msg[R]],
    *args: tuple[T],
    num: int | None = None,
    kind: ContextName = "process",
) -> Iterator[R]:
    """Call a function with arguments using multiple workers.

    Args:
        func (Callable): function to call

        *args (list[Any]): arguments to `func`. If multiple lists are provided,
            they will be passed to `zip` first.

        num (int, optional): number of workers. If `None`, `NUM_CPUS` or
            `NUM_THREADS` will be used as appropriate. Defaults to `None`.

        kind (ContextName, optional): execution context to use.
            Defaults to `"process"`.

    Yields:
        Any: results from applying the function to the arguments
    """
    q = Q[Iterable[T]](kind=kind)
    out = Q[R](kind=kind)

    def worker(_q: Q[Iterable[T]], _out: Q[R]) -> None:
        """Internal call to `func`."""
        for msg in _q.sorted():
            _out.put(data=task(*msg.data), order=msg.order)

    if kind == "process":
        workers = [Worker.process(worker, q, out) for _ in range(num or NUM_CPUS)]
    elif kind == "thread":
        workers = [Worker.thread(worker, q, out) for _ in range(num or NUM_THREADS)]
    else:  # pragma: no cover
        raise ValueError(f"Unknown worker context: {kind}")

    for order, value in enumerate(zip(*args)):
        q.put(value, order=order)
    q.stop(workers)

    for msg in out.end().sorted():
        yield msg.data
