# 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
"""Pipe worker for multi-processing."""

import os
import subprocess
import sys

from tvm_ffi import Shape, register_global_func


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

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

    Parameters
    ----------
    worker_id : int
        The worker id of the current worker.

    num_workers : int
        The total number of workers.

    num_groups : int
        The total number of worker groups.

    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__(  # pylint: disable=too-many-arguments
        self,
        worker_id: int,
        num_workers: int,
        num_groups: int,
        entrypoint: str = "tvm.exec.disco_worker",
        stdout=None,
        stderr=None,
    ):
        self.worker_id = worker_id
        self.num_workers = num_workers
        self.num_groups = num_groups
        self.entrypoint = entrypoint
        self._proc = None
        self._stdout = stdout
        self._stderr = stderr

    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:
            # 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

    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 start(self):
        """Start a new subprocess if nothing is available"""
        if self._proc is not None:
            return None, None

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

        cmd = [
            sys.executable,
            "-m",
            self.entrypoint,
            str(self.worker_id),
            str(self.num_workers),
            str(self.num_groups),
        ]
        if sys.platform == "win32":
            import msvcrt  # pylint: disable=import-error,import-outside-toplevel

            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(  # pylint: disable=consider-using-with
                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)
        return main_read, main_write


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

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

    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


@register_global_func("runtime.disco.create_process_pool")
def _create_process_pool(num_workers: int, num_groups: int, entrypoint: str):
    """Create a process pool where the workers' are [1, num_workers)."""
    pool = [DiscoPopenWorker(i, num_workers, num_groups, entrypoint) for i in range(1, num_workers)]

    def result_func(worker_id: int):
        nonlocal pool
        if worker_id != 0:
            read_fd, write_fd = pool[worker_id - 1].start()
            return Shape([read_fd, write_fd])
        del pool
        return None

    return result_func