# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
# ruff: noqa: E731
"""Multiprocessing via Popen.

This module provides a multi-processing pool backed by Popen.
with additional timeout support.
"""

import concurrent.futures
import os
import pickle
import struct
import subprocess
import sys
import threading
from collections import namedtuple
from enum import IntEnum


def kill_child_processes(pid):
    """Kill all child processes recursively for a given pid.

    Parameters
    ----------
    pid : int
        The given parameter id.
    """
    # pylint: disable=import-outside-toplevel
    import psutil

    try:
        parent = psutil.Process(pid)
        children = parent.children(recursive=True)
    except psutil.NoSuchProcess:
        return

    for process in children:
        try:
            process.kill()
        except psutil.NoSuchProcess:
            pass


class StatusKind(IntEnum):
    """Running and return value status."""

    RUNNING = 0
    COMPLETE = 1
    EXCEPTION = 2
    TIMEOUT = 3


class MapResult(namedtuple("MapResult", ["status", "value"])):
    """Result of map_with_error_catching.

    Parameters
    ----------
    status : StatusKind
        The status of the result.

    value : Any
        The result value.
    """

    __slots__ = []


class PopenWorker:
    """A subprocess worker via Popen.

    PopenWorker provides a low-level
    API to interact with a separate process via Popen.

    Parameters
    ----------
    initializer: callable or None
        A callable initializer, or None

    initargs: Tuple[object]
        A tuple of args for the initializer

    maximum_uses: Optional[int]
        The maximum number of times a process can be used before being recycled,
        i.e. killed and restarted. If `None`, the process will be reused until
        an operation times out.

    stdout: Union[None, int, IO[Any]]
        The standard output streams handler specified for the popen process.

    stderr: Union[None, int, IO[Any]]
        The standard error streams handler specified for the popen process.
    """

    def __init__(self, initializer=None, initargs=(), maximum_uses=None, stdout=None, stderr=None):
        self._proc = None
        self._initializer = initializer
        self._initargs = initargs
        self._maximum_uses = maximum_uses
        self._remaining_uses = None
        self._stdout = stdout
        self._stderr = stderr

        if self._initializer is not None and not callable(self._initializer):
            raise TypeError("initializer must be callable for PopenWorker")

    def __del__(self):
        try:
            self.kill()
        except ImportError:
            pass

    def kill(self):
        """Kill the current running process and cleanup.

        Note
        ----
        The worker can start a new process when send is called again.
        """
        if self._proc is not None:
            # allow gracefully shutdown
            try:
                self._writer.close()
            except OSError:
                pass
            try:
                self._reader.close()
            except OSError:
                pass
            # kill all child processes recursively
            try:
                kill_child_processes(self._proc.pid)
            except TypeError:
                pass
            try:
                self._proc.kill()
            except OSError:
                pass

            # Join the child process to avoid zombie processes
            self.join(timeout=1.0)
            self._proc = None
            self._remaining_uses = None

    def _start(self):
        """Start a new subprocess if nothing is available"""
        if self._proc is not None:
            return

        # connect subprocess with a pair of pipes
        main_read, worker_write = os.pipe()
        worker_read, main_write = os.pipe()

        cmd = [sys.executable, "-m", "tvm.exec.popen_worker"]
        if sys.platform == "win32":
            # pylint: disable=import-outside-toplevel
            import msvcrt

            worker_read_handle = msvcrt.get_osfhandle(worker_read)
            worker_write_handle = msvcrt.get_osfhandle(worker_write)
            os.set_handle_inheritable(worker_read_handle, True)
            os.set_handle_inheritable(worker_write_handle, True)
            cmd += [str(worker_read_handle), str(worker_write_handle)]
            self._proc = subprocess.Popen(
                cmd, close_fds=False, stdout=self._stdout, stderr=self._stderr
            )
        else:
            cmd += [str(worker_read), str(worker_write)]
            self._proc = subprocess.Popen(
                cmd, pass_fds=(worker_read, worker_write), stdout=self._stdout, stderr=self._stderr
            )

        # close worker side of the pipe
        os.close(worker_read)
        os.close(worker_write)
        self._reader = os.fdopen(main_read, "rb")
        self._writer = os.fdopen(main_write, "wb")

    def join(self, timeout=None):
        """Join the current process worker before it terminates.

        Parameters
        ----------
        timeout: Optional[number]
            Timeout value, block at most timeout seconds if it
            is a positive number.
        """
        if self._proc:
            try:
                self._proc.wait(timeout)
            except subprocess.TimeoutExpired:
                pass

    def is_alive(self):
        """Check if the process is alive"""
        if self._proc:
            return self._proc.poll() is None
        return False

    def send(self, fn, args=(), kwargs=None, timeout=None):
        """Send a new function task ``fn(*args, **kwargs)`` to the subprocess.

        Parameters
        ----------
        fn : function
            The function to be invoked.

        args : list
            Positional argument.

        kwargs : dict
            Keyword arguments

        timeout : float
            Timeout value when executing the function

        Note
        ----
        The caller must call recv before calling the next send in
        order to make sure the timeout and child process exit
        won't affect the later requests.
        """
        # use cloud pickle
        # pylint: disable=import-outside-toplevel
        import cloudpickle

        if self._proc is not None and self._maximum_uses and self._remaining_uses == 0:
            # Time to recycle the process.
            self.kill()

        if self._proc is None:
            self._start()
            # init
            if self._initializer is not None:
                self.send(self._initializer, self._initargs)
                self.recv()

            # N.B. The initializer doesn't count as a "use"
            self._remaining_uses = self._maximum_uses
        kwargs = {} if not kwargs else kwargs
        data = cloudpickle.dumps((fn, args, kwargs, timeout), protocol=pickle.HIGHEST_PROTOCOL)
        try:
            self._writer.write(struct.pack("<i", len(data)))
            self._writer.write(data)
            self._writer.flush()
        except OSError:
            pass

        if self._remaining_uses:
            self._remaining_uses -= 1

    def _child_process_error(self):
        """Raise a child process error."""
        # kill and lazily restart the process in the next send.
        self.kill()
        return ChildProcessError("Subprocess terminated")

    def recv(self):
        """Receive the result of the last send.

        Returns
        -------
        result: object
            The result of the last send.

        Raises
        ------
        ChildProcessError: if the child process exited abnormally.
        TimeoutError: if timeout happens
        Exception: if other exception happens during the execution.
        """
        # pylint: disable=import-outside-toplevel
        import cloudpickle

        try:
            len_data = self._reader.read(4)
        except OSError:
            raise self._child_process_error()

        if len(len_data) == 0:
            raise self._child_process_error()

        try:
            recv_bytes = struct.unpack("<i", len_data)[0]
            status, value = cloudpickle.loads(self._reader.read(recv_bytes))
        except OSError:
            raise self._child_process_error()

        if status == StatusKind.COMPLETE:
            return value
        if status == StatusKind.EXCEPTION:
            raise value
        assert status == StatusKind.TIMEOUT
        # kill and lazily restart the process in the next send.
        self.kill()
        raise TimeoutError()


class PopenPoolExecutor:
    """An parallel executor backed by Popen processes.

    Parameters
    ----------
    max_worker : int
        Maximum number of workers

    timeout : float
        Timeout value for each function submit.

    initializer: callable or None
        A callable initializer, or None

    initargs: Tuple[object]
        A tuple of args for the initializer

    maximum_process_uses: Optional[int]
        The maximum number of times each process can be used before being recycled,
        i.e. killed and restarted. If `None`, processes will be reused until an
        operation times out.

    stdout: Union[None, int, IO[Any]]
        The standard output streams handler specified for the workers in the pool.

    stderr: Union[None, int, IO[Any]]
        The standard error streams handler specified for the workers in the pool.

    Note
    ----
    If max_workers is NONE then the number returned by
    os.cpu_count() is used. This method aligns with the
    behavior of multiprocessing.pool().
    """

    def __init__(
        self,
        max_workers=None,
        timeout=None,
        initializer=None,
        initargs=(),
        maximum_process_uses=None,
        stdout=None,
        stderr=None,
    ):
        if max_workers is None:
            max_workers = os.cpu_count()
        # Use an internal thread pool to send to popen workers
        self._threadpool = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
        self._timeout = timeout
        self._worker_map = {}
        self._lock = threading.Lock()
        self._initializer = initializer
        self._initargs = initargs
        self._maximum_process_uses = maximum_process_uses
        self._stdout = stdout
        self._stderr = stderr
        self._shutdown = False

        if self._initializer is not None and not callable(self._initializer):
            raise TypeError("initializer must be callable for PopenPoolExecutor")

    def __del__(self):
        """Destructor.

        Note
        ----
        Called during garbage collection. This may be called later than expected.
        Always call shutdown() explicitly to avoid deadlocks.
        """
        if not self._shutdown:
            self.shutdown(wait=True)

    def shutdown(self, wait=True):
        """Shutdown the executor and clean up resources.

        Parameters
        ----------
        wait : bool
            If True, wait for pending work to complete.

        Note
        ----
        DEADLOCK WARNING: This method can deadlock when called during garbage
        collection due to exception reference cycles. When exceptions occur,
        Python creates reference cycles that delay garbage collection. The
        deadlock happens when: exception creates reference cycle → new pool
        creates worker → GC cleans old pool → old pool's __del__ calls shutdown()
        which tries to acquire locks again.
        """
        self._lock.acquire()
        for worker in self._worker_map.values():
            try:
                worker.kill()
            except ImportError:
                pass
        self._lock.release()
        self._threadpool.shutdown(wait=wait)
        self._shutdown = True

    def _worker_run(self, fn, args, kwargs):
        """Internal thread runner."""
        self._lock.acquire()
        tid = threading.get_ident()
        if tid not in self._worker_map:
            proc = PopenWorker(
                self._initializer,
                self._initargs,
                self._maximum_process_uses,
                self._stdout,
                self._stderr,
            )
            self._worker_map[tid] = proc
        else:
            proc = self._worker_map[tid]
        self._lock.release()

        proc.send(fn, args, kwargs, self._timeout)
        return proc.recv()

    def _worker_run_with_error_catching(self, fn, args, kwargs) -> MapResult:
        # pylint: disable=broad-except
        try:
            return MapResult(status=StatusKind.COMPLETE, value=self._worker_run(fn, args, kwargs))
        except TimeoutError as exception:
            return MapResult(status=StatusKind.TIMEOUT, value=exception)
        except Exception as exception:
            return MapResult(status=StatusKind.EXCEPTION, value=exception)

    def submit(self, fn, *args, **kwargs) -> concurrent.futures.Future:
        """Submit a new function job to the pool

        Parameters
        ----------
        fn : function
            The function to be invoked.

        args : list
            Positional argument.

        kwargs : dict
            Keyword arguments

        Returns
        -------
        future : concurrent.futures.Future
            A future that can be used to access the result.
        """
        # pylint: disable=unnecessary-lambda
        worker = lambda *args: self._worker_run(*args)
        return self._threadpool.submit(worker, fn, args, kwargs)

    def map_with_error_catching(self, fn, iterator):
        """Same as map, but catches exceptions and return them instead.

        Parameters
        ----------
        fn : function
            The function to be invoked.

        iterator : Iterator
            Input iterator.

        Returns
        -------
        out_iter : Iterator[MapResult]
            The result iterator.
        """
        worker = lambda x: self._worker_run_with_error_catching(fn, (x,), None)
        return self._threadpool.map(worker, iterator)