import copyreg
import io
import multiprocessing
import pickle
import pickletools
import queue
from functools import wraps
from types import TracebackType
from typing import Any, Callable, NamedTuple, ParamSpec, TypeVar
from typing_extensions import Self

from tblib.pickling_support import pickle_traceback

T = TypeVar("T")
P = ParamSpec("P")
R = TypeVar("R")


def _register_pickle(cls: type) -> type:
    """Registers a custom pickle function for a class to handle pickle serialization."""
    copyreg.pickle(TracebackType, pickle_traceback)
    return cls


class UnpickableError(Exception):
    """Exception raised when an unpickable error occurs during multiprocessing."""

    def __init__(self, result: "_Result") -> None:
        """Constructor."""
        super().__init__()
        self._result = result

    @property
    def is_err(self) -> bool:
        """Indicates whether the unpickable result is an error."""
        return self._result.is_err

    def __str__(self) -> str:
        builder = [
            "",
            f"Module: {self._result.module}",
            f"Is Error: {self._result.is_err}",
            f"Pickle Trace: {self._result.pickle_trace}",
        ]
        return "\n".join(builder)


@_register_pickle
class _Result(NamedTuple):
    """A named tuple to hold the result of a process execution."""

    module: str
    value: bytes
    traceback: TracebackType | None = None

    @property
    def is_err(self) -> bool:
        """Indicates whether the result is an error."""
        return self.traceback is not None

    @property
    def pickle_trace(self) -> str:
        """Provides a string representation of the pickled value for debugging."""
        with io.StringIO() as fp:
            pickletools.dis(self.value, fp)
            return fp.getvalue()

    @classmethod
    def from_err(cls, ex: Exception) -> Self:
        """Creates an instance representing an error."""
        return cls(ex.__class__.__module__, pickle.dumps(ex), ex.__traceback__)

    @classmethod
    def from_ok(cls, value: Any) -> Self:
        """Creates an instance representing a successful result."""
        return cls(value.__class__.__module__, pickle.dumps(value))


def _run(q: multiprocessing.Queue, func: Callable[[], Any]) -> None:
    """Runs the function and puts the result in the queue."""
    try:
        q.put_nowait(_Result.from_ok(func()))
    except Exception as ex:
        q.put_nowait(_Result.from_err(ex))


def run_in_process(
    func: Callable[[], T],
    timeout_sec: float | None = None,
    should_reset_timeout: Callable[[], bool] | None = None,
    kill_wait_sec: float | None = None,
    start_method: str = "fork",
    name: str | None = None,
) -> T:
    """Runs a function in a separate process with an optional timeout and graceful shutdown.

    Warnings:
        The "spawn" starting method does not support unpickable objects like lambda.

    Args:
        func: The function to be executed.
        timeout_sec: The seconds to wait for the function to complete before timing out.
        should_reset_timeout: A function that decides if the timeout should be reset.
        kill_wait_sec: Seconds to wait before sending SIGKILL; None to disable.
        start_method: The method used to start the new process.
        name: Optional name for the subprocess.
    """
    assert start_method in ("fork", "spawn")

    ctx = multiprocessing.get_context(start_method)
    q = ctx.Queue()
    p = ctx.Process(name=name, target=_run, args=(q, func))
    p.start()

    def _on_timeout() -> None:
        try:
            p.terminate()
        except ValueError:
            pass

        if p.exitcode == 0:
            if q.empty():
                raise multiprocessing.ProcessError(
                    "Subprocess exited with an empty queue, check stderr for details"
                ) from None
        elif p.exitcode is not None:
            raise multiprocessing.ProcessError(
                f"Subprocess exited with code {p.exitcode}, check stderr for details"
            ) from None

        if kill_wait_sec is not None:
            try:
                p.join(kill_wait_sec)
                if p.is_alive():
                    p.kill()
            except ValueError:
                pass

        raise multiprocessing.TimeoutError(
            f"Timed out after {timeout_sec} seconds, exited with code {p.exitcode}"
        ) from None

    while True:
        try:
            p.join(timeout_sec)
            r = q.get_nowait()
        except queue.Empty:
            if should_reset_timeout and should_reset_timeout():
                continue
            _on_timeout()
            raise NotImplementedError from None  # unreachable; _on_timeout always raises

        try:
            value = pickle.loads(r.value)
        except Exception as ex:
            raise UnpickableError(r).with_traceback(r.traceback) from ex
        if r.is_err:
            assert isinstance(value, Exception)
            raise value.with_traceback(r.traceback)
        return value  # type: ignore[return-value]


def run_in_process_decorator(func: Callable[P, R]) -> Callable[P, R]:
    """A decorator to run a function in a separate process."""

    @wraps(func)
    def _inner(*args: P.args, **kwargs: P.kwargs) -> R:
        return run_in_process(lambda: func(*args, **kwargs))

    return _inner
