# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Routines common to all posix systems."""

import enum
import errno
import glob
import os
import select
import signal
import time

from . import _ntuples as ntp
from ._common import MACOS
from ._common import TimeoutExpired
from ._common import debug
from ._common import memoize
from ._common import usage_percent

if MACOS:
    from . import _psutil_osx


__all__ = ['pid_exists', 'wait_pid', 'disk_usage', 'get_terminal_map']


def pid_exists(pid):
    """Check whether pid exists in the current process table."""
    if pid == 0:
        # According to "man 2 kill" PID 0 has a special meaning:
        # it refers to <<every process in the process group of the
        # calling process>> so we don't want to go any further.
        # If we get here it means this UNIX platform *does* have
        # a process with id 0.
        return True
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        # EPERM clearly means there's a process to deny access to
        return True
    # According to "man 2 kill" possible error values are
    # (EINVAL, EPERM, ESRCH)
    else:
        return True


Negsignal = enum.IntEnum(
    'Negsignal', {x.name: -x.value for x in signal.Signals}
)


def negsig_to_enum(num):
    """Convert a negative signal value to an enum."""
    try:
        return Negsignal(num)
    except ValueError:
        return num


def convert_exit_code(status):
    """Convert a os.waitpid() status to an exit code."""
    if os.WIFEXITED(status):
        # Process terminated normally by calling exit(3) or _exit(2),
        # or by returning from main(). The return value is the
        # positive integer passed to *exit().
        return os.WEXITSTATUS(status)
    if os.WIFSIGNALED(status):
        # Process exited due to a signal. Return the negative value
        # of that signal.
        return negsig_to_enum(-os.WTERMSIG(status))
    # if os.WIFSTOPPED(status):
    #     # Process was stopped via SIGSTOP or is being traced, and
    #     # waitpid() was called with WUNTRACED flag. PID is still
    #     # alive. From now on waitpid() will keep returning (0, 0)
    #     # until the process state doesn't change.
    #     # It may make sense to catch/enable this since stopped PIDs
    #     # ignore SIGTERM.
    #     interval = sleep(interval)
    #     continue
    # if os.WIFCONTINUED(status):
    #     # Process was resumed via SIGCONT and waitpid() was called
    #     # with WCONTINUED flag.
    #     interval = sleep(interval)
    #     continue

    # Should never happen.
    msg = f"unknown process exit status {status!r}"
    raise ValueError(msg)


def wait_pid_posix(
    pid,
    timeout=None,
    _waitpid=os.waitpid,
    _timer=getattr(time, 'monotonic', time.time),  # noqa: B008
    _min=min,
    _sleep=time.sleep,
    _pid_exists=pid_exists,
):
    """Wait for a process PID to terminate.

    If the process terminated normally by calling exit(3) or _exit(2),
    or by returning from main(), the return value is the positive integer
    passed to *exit().

    If it was terminated by a signal it returns the negated value of the
    signal which caused the termination (e.g. -SIGTERM).

    If PID is not a children of os.getpid() (current process) just
    wait until the process disappears and return None.

    If PID does not exist at all return None immediately.

    If timeout is specified and process is still alive raise
    TimeoutExpired.

    If timeout=0 either return immediately or raise TimeoutExpired
    (non-blocking).
    """
    interval = 0.0001
    max_interval = 0.04
    flags = 0
    stop_at = None

    if timeout is not None:
        flags |= os.WNOHANG
        if timeout != 0:
            stop_at = _timer() + timeout

    def sleep_or_timeout(interval):
        # Sleep for some time and return a new increased interval.
        if timeout == 0 or (stop_at is not None and _timer() >= stop_at):
            raise TimeoutExpired(timeout)
        _sleep(interval)
        return _min(interval * 2, max_interval)

    # See: https://linux.die.net/man/2/waitpid
    while True:
        try:
            retpid, status = os.waitpid(pid, flags)
        except ChildProcessError:
            # This has two meanings:
            # - PID is not a child of os.getpid() in which case
            #   we keep polling until it's gone
            # - PID never existed in the first place
            # In both cases we'll eventually return None as we
            # can't determine its exit status code.
            while _pid_exists(pid):
                interval = sleep_or_timeout(interval)
            return None
        else:
            if retpid == 0:
                # WNOHANG flag was used and PID is still running.
                interval = sleep_or_timeout(interval)
            else:
                return convert_exit_code(status)


def _waitpid(pid, timeout):
    """Wrapper around os.waitpid(). PID is supposed to be gone already,
    it just returns the exit code.
    """
    try:
        retpid, status = os.waitpid(pid, 0)
    except ChildProcessError:
        # PID is not a child of os.getpid().
        return wait_pid_posix(pid, timeout)
    else:
        assert retpid != 0
        return convert_exit_code(status)


def wait_pid_pidfd_open(pid, timeout=None):
    """Wait for PID to terminate using pidfd_open() + poll(). Linux >=
    5.3 + Python >= 3.9 only.
    """
    try:
        pidfd = os.pidfd_open(pid, 0)
    except OSError as err:
        if err.errno == errno.ESRCH:
            # No such process. os.waitpid() may still be able to return
            # the status code.
            return wait_pid_posix(pid, timeout)
        if err.errno in {errno.EMFILE, errno.ENFILE, errno.ENODEV}:
            # EMFILE, ENFILE: too many open files
            # ENODEV: anonymous inode filesystem not supported
            debug(f"pidfd_open() failed ({err!r}); use fallback")
            return wait_pid_posix(pid, timeout)
        raise

    try:
        # poll() / select() have the advantage of not requiring any
        # extra file descriptor, contrary to epoll() / kqueue().
        # select() crashes if process opens > 1024 FDs, so we use
        # poll().
        poller = select.poll()
        poller.register(pidfd, select.POLLIN)
        timeout_ms = None if timeout is None else int(timeout * 1000)
        events = poller.poll(timeout_ms)  # wait

        if not events:
            raise TimeoutExpired(timeout)
        return _waitpid(pid, timeout)
    finally:
        os.close(pidfd)


def wait_pid_kqueue(pid, timeout=None):
    """Wait for PID to terminate using kqueue(). macOS and BSD only."""
    try:
        kq = select.kqueue()
    except OSError as err:
        if err.errno in {errno.EMFILE, errno.ENFILE}:  # too many open files
            debug(f"kqueue() failed ({err!r}); use fallback")
            return wait_pid_posix(pid, timeout)
        raise

    try:
        kev = select.kevent(
            pid,
            filter=select.KQ_FILTER_PROC,
            flags=select.KQ_EV_ADD | select.KQ_EV_ONESHOT,
            fflags=select.KQ_NOTE_EXIT,
        )
        try:
            events = kq.control([kev], 1, timeout)  # wait
        except OSError as err:
            if err.errno in {errno.EACCES, errno.EPERM, errno.ESRCH}:
                debug(f"kqueue.control() failed ({err!r}); use fallback")
                return wait_pid_posix(pid, timeout)
            raise
        else:
            if not events:
                raise TimeoutExpired(timeout)
            return _waitpid(pid, timeout)
    finally:
        kq.close()


@memoize
def can_use_pidfd_open():
    # Availability: Linux >= 5.3, Python >= 3.9
    if not hasattr(os, "pidfd_open"):
        return False
    try:
        pidfd = os.pidfd_open(os.getpid(), 0)
    except OSError as err:
        if err.errno in {errno.EMFILE, errno.ENFILE}:  # noqa: SIM103
            # transitory 'too many open files'
            return True
        # likely blocked by security policy like SECCOMP (EPERM,
        # EACCES, ENOSYS)
        return False
    else:
        os.close(pidfd)
        return True


@memoize
def can_use_kqueue():
    # Availability: macOS, BSD
    names = (
        "kqueue",
        "KQ_EV_ADD",
        "KQ_EV_ONESHOT",
        "KQ_FILTER_PROC",
        "KQ_NOTE_EXIT",
    )
    if not all(hasattr(select, x) for x in names):
        return False
    kq = None
    try:
        kq = select.kqueue()
        kev = select.kevent(
            os.getpid(),
            filter=select.KQ_FILTER_PROC,
            flags=select.KQ_EV_ADD | select.KQ_EV_ONESHOT,
            fflags=select.KQ_NOTE_EXIT,
        )
        kq.control([kev], 1, 0)
        return True
    except OSError as err:
        if err.errno in {errno.EMFILE, errno.ENFILE}:  # noqa: SIM103
            # transitory 'too many open files'
            return True
        return False
    finally:
        if kq is not None:
            kq.close()


def wait_pid(pid, timeout=None):
    # PID 0 passed to waitpid() waits for any child of the current
    # process to change state.
    assert pid > 0
    if timeout is not None:
        assert timeout >= 0

    if can_use_pidfd_open():
        return wait_pid_pidfd_open(pid, timeout)
    elif can_use_kqueue():
        return wait_pid_kqueue(pid, timeout)
    else:
        return wait_pid_posix(pid, timeout)


wait_pid.__doc__ = wait_pid_posix.__doc__


def disk_usage(path):
    """Return disk usage associated with path.
    Note: UNIX usually reserves 5% disk space which is not accessible
    by user. In this function "total" and "used" values reflect the
    total and used disk space whereas "free" and "percent" represent
    the "free" and "used percent" user disk space.
    """
    st = os.statvfs(path)
    # Total space which is only available to root (unless changed
    # at system level).
    total = st.f_blocks * st.f_frsize
    # Remaining free space usable by root.
    avail_to_root = st.f_bfree * st.f_frsize
    # Remaining free space usable by user.
    avail_to_user = st.f_bavail * st.f_frsize
    # Total space being used in general.
    used = total - avail_to_root
    if MACOS:
        # see: https://github.com/giampaolo/psutil/pull/2152
        used = _psutil_osx.disk_usage_used(path, used)
    # Total space which is available to user (same as 'total' but
    # for the user).
    total_user = used + avail_to_user
    # User usage percent compared to the total amount of space
    # the user can use. This number would be higher if compared
    # to root's because the user has less space (usually -5%).
    usage_percent_user = usage_percent(used, total_user, round_=1)

    # NB: the percentage is -5% than what shown by df due to
    # reserved blocks that we are currently not considering:
    # https://github.com/giampaolo/psutil/issues/829#issuecomment-223750462
    return ntp.sdiskusage(
        total=total, used=used, free=avail_to_user, percent=usage_percent_user
    )


@memoize
def get_terminal_map():
    """Get a map of device-id -> path as a dict.
    Used by Process.terminal().
    """
    ret = {}
    ls = glob.glob('/dev/tty*') + glob.glob('/dev/pts/*')
    for name in ls:
        assert name not in ret, name
        try:
            ret[os.stat(name).st_rdev] = name
        except FileNotFoundError:
            pass
    return ret