"""Process task pool executor."""

from __future__ import annotations

import asyncio
import sys
import threading
import traceback
from concurrent.futures import ProcessPoolExecutor
from functools import partial
from logging import getLogger
from multiprocessing.managers import SyncManager
from signal import SIGTERM
from signal import getsignal
from signal import signal
from typing import TYPE_CHECKING
from typing import Any
from typing import Awaitable
from typing import Callable
from typing import Iterable
from typing import Literal

import psutil

from asynccpu.process_task import ProcessTask
from asynccpu.subprocess import LoggingInitializer
from asynccpu.subprocess import Replier
from asynccpu.subprocess import run

if TYPE_CHECKING:
    import queue
    from asyncio.futures import Future
    from logging import LogRecord
    from multiprocessing.context import BaseContext
    from multiprocessing.managers import DictProxy
    from types import FrameType
    from types import TracebackType

    from asynccpu.types import ParamSpecCoroutineFunctionArguments
    from asynccpu.types import TypeVarReturnValue

__all__ = ["ProcessTaskPoolExecutor"]


class ProcessTaskFactory:
    """Factory for ProcessTask."""

    def __init__(
        self,
        process_pool_executor: ProcessPoolExecutor,
        sync_manager: SyncManager,
        *,
        logging_initializer: LoggingInitializer | None = None,
    ) -> None:
        self.process_pool_executor = process_pool_executor
        self.sync_manager = sync_manager
        # When use list, hash seems to be updated when check in parent process side (due to pickle?)
        # > assert list_process[0].id == process.id
        # > assert list_process[0].__hash__ == process.__hash__
        # E   assert <bound method...7fc75e20f880>> == <bound method...7fc75e20f850>>
        # E     +<bound method ProcessForWeakSet.__hash__ of
        # E        <asynccpu.subprocess.ProcessForWeakSet object at 0x7fc75e20f880>>
        # E     -<bound method ProcessForWeakSet.__hash__ of
        # E        <asynccpu.subprocess.ProcessForWeakSet object at 0x7fc75e20f850>>
        self.dictionary_process: DictProxy[int, ProcessTask] = sync_manager.dict()
        self.logging_initializer = logging_initializer
        self.loop = asyncio.get_event_loop()
        self.logger = getLogger(__name__)

    def create(
        self,
        function_coroutine: Callable[ParamSpecCoroutineFunctionArguments, Awaitable[TypeVarReturnValue]],
        *args: ParamSpecCoroutineFunctionArguments.args,
        **kwargs: ParamSpecCoroutineFunctionArguments.kwargs,
    ) -> Future[TypeVarReturnValue]:
        """Creates ProcessTask."""
        try:
            queue_process_id = self.sync_manager.Queue()
        # Reason: Only on Windows.
        except InterruptedError as error:  # pragma: no cover
            # On Windows, CTRL_C_EVENT raises InterruptedError during blocking operations
            # Convert to KeyboardInterrupt for consistent handling across platforms
            self.logger.debug("InterruptedError during Queue creation: %s", error)
            raise KeyboardInterrupt from error
        replier = Replier(self.dictionary_process, queue_process_id)
        partial_func = partial(run, replier, self.logging_initializer, function_coroutine, **kwargs)
        task = self.loop.run_in_executor(self.process_pool_executor, partial_func, *args)
        try:
            queue_process_id.get()
        # Reason: Only on Windows.
        except InterruptedError as error:  # pragma: no cover
            # On Windows, CTRL_C_EVENT raises InterruptedError during blocking operations
            # Convert to KeyboardInterrupt for consistent handling across platforms
            self.logger.debug("InterruptedError during queue get: %s", error)
            self.logger.debug(task)
            cancel = task.cancel()
            self.logger.debug("Cancel succeed?: %s", cancel)
            raise KeyboardInterrupt from error
        except BaseException as error:
            self.logger.debug("%s", error)
            self.logger.debug(task)
            cancelled = task.cancelled()
            self.logger.debug("Cancelled?: %s", cancelled)
            if cancelled:
                # Reason: This conditional raise is just in case.
                raise  # pragma: no cover
            cancel = task.cancel()
            self.logger.debug("Cancel succeed?: %s", cancel)
            raise
        return task

    def send_signal(self, signal_number: int) -> None:
        for process_task in self.items_dictionary_process():
            self.logger.debug("Send signal: Start")
            process_task.send_signal(signal_number)
            self.logger.debug("Send signal: Finish")

    def items_dictionary_process(self) -> Iterable[ProcessTask]:
        """DictProxy often returns not list but single instance when call items()"""
        items_dictionary_process = self.dictionary_process.values()
        self.logger.debug("type(items_dictionary_process)=%s", type(items_dictionary_process))
        if isinstance(items_dictionary_process, ProcessTask):
            return [items_dictionary_process]
        return items_dictionary_process


class ProcessTaskManager:
    """Process task manager."""

    # Reason: Assuming a general use case, method create_process_task() will not take so much than 1 second.
    TIMEOUT_IN_EMERGENCY_DEFAULT = 1

    def __init__(
        self,
        process_task_factory: ProcessTaskFactory,
        *,
        timeout_in_emergency: int = TIMEOUT_IN_EMERGENCY_DEFAULT,
    ) -> None:
        """timeout_in_emergency: The time to force breaking lock in emergency."""
        self.lock_for_dictionary_process_task = threading.Lock()
        self.process_task_factory = process_task_factory
        self.timeout_in_emergency = timeout_in_emergency
        self.logger = getLogger(__name__)

    def lock_and_send_signal(self, signal_number: int) -> None:
        """Cancels task like future if its not cancelled."""
        self.logger.debug("Lock and send signal: Start")
        # Reason: requires to set timeout. pylint: disable=consider-using-with
        is_locked = self.lock_for_dictionary_process_task.acquire(timeout=self.timeout_in_emergency)
        self.logger.debug(
            "Lock and send signal: Locked" if is_locked else "Lock and send signal: Timed out to acquire lock",
        )
        self.process_task_factory.send_signal(signal_number)
        self.logger.debug("Lock and send signal: Finish")

    def create_process_task(
        self,
        function_coroutine: Callable[ParamSpecCoroutineFunctionArguments, Awaitable[TypeVarReturnValue]],
        *args: ParamSpecCoroutineFunctionArguments.args,
        **kwargs: ParamSpecCoroutineFunctionArguments.kwargs,
    ) -> Future[TypeVarReturnValue]:
        """Creates task like future by wraping coroutine."""
        with self.lock_for_dictionary_process_task:
            return self.process_task_factory.create(function_coroutine, *args, **kwargs)


class ProcessTaskPoolExecutor(ProcessPoolExecutor):
    """see:

    - Answer: What's the correct way to clean up after an interrupted event loop?
      https://stackoverflow.com/a/30766124/12721873
      https://stackoverflow.com/a/58532304/12721873
    """

    # Reason: Arguments are required to initialize super class.
    def __init__(  # noqa: PLR0913  # pylint: disable=too-many-arguments
        self,
        max_workers: int | None = None,
        mp_context: BaseContext | None = None,
        initializer: Callable[..., Any] | None = None,
        initargs: tuple[Any, ...] = (),
        *,
        cancel_tasks_when_shutdown: bool = False,
        # Reason: This argument name is API. pylint: disable=redefined-outer-name
        queue: queue.Queue[LogRecord] | None = None,
        configurer: Callable[[], Any] | None = None,
    ) -> None:
        super().__init__(max_workers, mp_context, initializer, initargs)
        self.cancel_tasks_when_shutdown = cancel_tasks_when_shutdown
        # This SyncManager is to use Queue to receive process id from child process.
        # It's danger to instantiate by with statement around code to create child process.
        # It increases the probability hitting signal during instantiating SyncManager.
        # If it happens, Python hangs up!
        #
        # stack trace:
        # KeyboardInterrupt
        # File "/xxxx/xxxx.py", line xx, in create
        #   with SyncManager() as manager:
        # File "/usr/local/lib/python3.9/multiprocessing/managers.py", line 635, in __enter__
        #   self.start()  # noqa: ERA001
        # File "/usr/local/lib/python3.9/multiprocessing/managers.py", line 557, in start
        #   self._address = reader.recv() # noqa: ERA001
        # File "/usr/local/lib/python3.9/multiprocessing/connection.py", line 255, in recv
        #   buf = self._recv_bytes()  # noqa: ERA001
        # File "/usr/local/lib/python3.9/multiprocessing/connection.py", line 419, in _recv_bytes
        #   buf = self._recv(4)  # noqa: ERA001
        # File "/usr/local/lib/python3.9/multiprocessing/connection.py", line 384, in _recv
        #   chunk = read(handle, remaining)  # noqa: ERA001
        self.sync_manager = SyncManager()
        # Reason:
        # We can't use with statement without contextlib.
        # It's better to divide __enter__() and __exit__() for simplicity and testability.
        # pylint: disable=consider-using-with
        self.sync_manager.start()
        logging_initializer = LoggingInitializer(queue, configurer)
        process_task_factory = ProcessTaskFactory(self, self.sync_manager, logging_initializer=logging_initializer)
        self.process_task_manager = ProcessTaskManager(process_task_factory)
        self.logger = getLogger(__name__)
        self.original_sigint_handler = getsignal(SIGTERM)
        signal(SIGTERM, self.sigterm_handler)

    # Reason: pylint bug. pylint: disable=unsubscriptable-object
    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> Literal[False]:
        self.logger.debug("Exit ProcessTaskPoolExecutor: Start")
        self.logger.debug("exc_type = %s", exc_type)
        self.logger.debug("exc_val = %s", exc_val)
        self.logger.debug("exc_tb = %s", "".join(traceback.format_tb(exc_tb)))
        self._shutdown(wait=True, signal_number=SIGTERM)
        # The pylint can't detect that return value is False,
        # And it might be better to not call super class __exit__().
        # return_value = super().__exit__(exc_type, exc_val, exc_tb) noqa: ERA001
        self.logger.debug("Exit ProcessTaskPoolExecutor: Finish")
        return False

    # Reason: Can't collect coverage because of termination.
    def sigterm_handler(self, _signum: int, _frame: FrameType | None) -> None:  # pragma: no cover
        self.logger.debug("SIGTERM handler ProcessTaskPoolExecutor: Start")
        self._shutdown(wait=True, signal_number=SIGTERM)
        signal(SIGTERM, self.original_sigint_handler)
        self.logger.debug("SIGTERM handler ProcessTaskPoolExecutor: Finish")
        psutil.Process().terminate()
        self.logger.debug("SIGTERM handler ProcessTaskPoolExecutor: Sent SIGTERM")

    # Reason: Following super class specification of shutdown() method.
    def shutdown(self, wait: bool = True, *, cancel_futures: bool | None = None) -> None:  # noqa: FBT001,FBT002
        self._shutdown(wait=wait, cancel_futures=cancel_futures)

    def _shutdown(
        self,
        *,
        wait: bool = True,
        cancel_futures: bool | None = None,
        signal_number: int = SIGTERM,
    ) -> None:
        """Shutdown ProcessTaskPoolExecutor.

        Although ProcessPoolExecutor won't cancel any futures that are completed or running even if *cancel_futures*
        is `True`, This class attempt to cancel also them when *cancel_futures* is `True`.
        """
        if cancel_futures is None:
            cancel_futures = self.cancel_tasks_when_shutdown
        # In Windows, it should not send SIGTERM.
        # It will propergate KeyboardInterrupt to all processes including parent process.
        # see:
        # - windows - Trying to implement `signal.CTRL_C_EVENT` in Python3.6 - Stack Overflow
        #   https://stackoverflow.com/a/44147750/12721873
        # - How to handle the signal in python on windows machine - Stack Overflow
        #   https://stackoverflow.com/a/35792192/12721873
        if cancel_futures and sys.platform != "win32":
            self.logger.debug("Shutdown ProcessTaskPoolExecutor: Cancel all process tasks")
            self.process_task_manager.lock_and_send_signal(signal_number)
        self.logger.debug("Shutdown ProcessTaskPoolExecutor: Shutdown ProcessPoolExecutor")
        self.call_super_class_shutdown(wait=wait, cancel_futures=cancel_futures)
        self.logger.debug("Shutdown ProcessTaskPoolExecutor: Shutdown sync manager")
        self.sync_manager.shutdown()

    def call_super_class_shutdown(self, *, wait: bool, cancel_futures: bool) -> None:
        """Calls shutdown() of super class."""
        if sys.version_info >= (3, 9):
            super().shutdown(wait=wait, cancel_futures=cancel_futures)
        else:
            super().shutdown(wait=wait)  # pragma: no cover

    def create_process_task(
        self,
        function_coroutine: Callable[ParamSpecCoroutineFunctionArguments, Awaitable[TypeVarReturnValue]],
        *args: ParamSpecCoroutineFunctionArguments.args,
        **kwargs: ParamSpecCoroutineFunctionArguments.kwargs,
    ) -> Future[TypeVarReturnValue]:
        """Creates task like future by wrapping coroutine."""
        return self.process_task_manager.create_process_task(function_coroutine, *args, **kwargs)
