# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: LGPL-2.0-or-later

"""A type tag for integers that represent file descriptors."""

import collections.abc
import contextlib
import fcntl
import os
import typing


# pylint: disable=too-few-public-methods  # It's that one method we describe.
@typing.runtime_checkable
class HasFileno(typing.Protocol):
    """A typing protocol representing a file-like object and looking up the
    underlying file descriptor.
    """

    def fileno(self) -> int:
        """Return the underlying file descriptor."""


FileDescriptorLike = int | HasFileno


@contextlib.contextmanager
def _oserror_function(function: str) -> collections.abc.Iterator[None]:
    try:
        yield
    except OSError as err:
        if not hasattr(err, "function"):
            setattr(err, "function", function)
        raise


class FileDescriptor(int):
    """Type tag for integers that represent file descriptors. It also provides
    a few very generic file descriptor methods.
    """

    def __new__(cls, value: FileDescriptorLike) -> typing.Self:
        """Construct a FileDescriptor from an int or HasFileno."""
        if isinstance(value, cls):
            return value  # No need to copy, it's immutable.
        if not isinstance(value, int):
            value = value.fileno()
        return super(FileDescriptor, cls).__new__(cls, value)

    def __enter__(self) -> "FileDescriptor":
        """When used as a context manager, close the file descriptor on scope
        exit.
        """
        return self

    def __exit__(self, *_: object) -> None:
        """When used as a context manager, close the file descriptor on scope
        exit.
        """
        self.close()

    def close(self) -> None:
        """Close the file descriptor. Since int is immutable, the caller is
        responsible for not closing twice.
        """
        with _oserror_function("close"):
            os.close(self)

    def dup(self, inheritable: bool = True) -> "FileDescriptor":
        """Return a duplicate of the file descriptor."""
        if inheritable:
            with _oserror_function("dup"):
                return FileDescriptor(os.dup(self))
        with _oserror_function("fcntl"):
            return FileDescriptor(fcntl.fcntl(self, fcntl.F_DUPFD_CLOEXEC, 0))

    def dup2(
        self, fd2: FileDescriptorLike, inheritable: bool = True
    ) -> "FileDescriptor":
        """Duplicate the file to the given file descriptor number."""
        with _oserror_function("dup2"):
            return FileDescriptor(
                os.dup2(self, FileDescriptor(fd2), inheritable)
            )

    def fstat(self) -> os.stat_result:
        """Return information about the file referenced by the file descriptor.
        """
        with _oserror_function("fstat"):
            return os.fstat(self)

    @classmethod
    def pidfd_open(cls, pid: int, flags: int = 0) -> typing.Self:
        """Convenience wrapper for os.pidfd_open."""
        with _oserror_function("pidfd_open"):
            return cls(os.pidfd_open(pid, flags))

    def fileno(self) -> "FileDescriptor":
        """Return self such that it satisfies the HasFileno protocol."""
        return self

    def get_blocking(self) -> bool:
        """Get the blocking mode of the file descriptor."""
        with _oserror_function("fcntl"):
            return os.get_blocking(self)

    def get_inheritable(self) -> bool:
        """Get the close-on-exec flag of the file descriptor."""
        with _oserror_function("fcntl"):
            return os.get_inheritable(self)

    @classmethod
    def pipe(
        cls, blocking: bool = True, inheritable: bool = True
    ) -> tuple["FileDescriptor", "FileDescriptor"]:
        """Create a pipe with flags set atomically. This actually corresponds
        to the pipe2 syscall, but skipping flags is equivalent to calling pipe.
        """
        with _oserror_function("pipe2"):
            rfd, wfd = os.pipe2(
                (0 if blocking else os.O_NONBLOCK)
                | (0 if inheritable else os.O_CLOEXEC),
            )
        return (cls(rfd), cls(wfd))

    def set_blocking(self, blocking: bool) -> None:
        """Set the blocking mode of the file descriptor."""
        with _oserror_function("fcntl"):
            os.set_blocking(self, blocking)

    def set_inheritable(self, inheritable: bool) -> None:
        """Set the close-on-exec flag of the file descriptor."""
        with _oserror_function("fcntl"):
            os.set_inheritable(self, inheritable)
