import asyncio
import atexit
import errno
import functools
import logging
import os
import pathlib
import pwd
import shutil
import stat
from concurrent.futures import ProcessPoolExecutor
from contextlib import contextmanager, suppress
from typing import Set, Tuple, Union

from defence360agent import utils

COPY_EXPECTED_UID = True

R_FLAGS = os.O_RDONLY
W_FLAGS = os.O_TRUNC | os.O_CREAT | os.O_WRONLY
IN_PLACE_W_FLAGS = os.O_CREAT | os.O_WRONLY

logger = logging.getLogger(__name__)

# Track active ProcessPoolExecutors so they can be cleaned up during shutdown.
# Each call to drop() permanently changes the worker process identity, so we
# must use a fresh executor per call. We track them to prevent orphaned worker
# processes from blocking agent shutdown (causing systemd SIGKILL).
_active_pools: Set[ProcessPoolExecutor] = set()


async def _run_in_fresh_executor(loop: asyncio.AbstractEventLoop, *args):
    pool = ProcessPoolExecutor(max_workers=1)
    _active_pools.add(pool)
    try:
        return await loop.run_in_executor(pool, *args)
    finally:
        try:
            pool.shutdown(wait=False)
        finally:
            _active_pools.discard(pool)


def shutdown_process_pools() -> None:
    """Shutdown all tracked ProcessPoolExecutors.

    Should be called during agent shutdown to ensure clean process termination.
    """
    for pool in list(_active_pools):
        try:
            pool.shutdown(wait=False, cancel_futures=True)
        except Exception as e:
            logger.warning("Error shutting down ProcessPoolExecutor: %s", e)
    _active_pools.clear()


# Register cleanup at exit as a fallback
atexit.register(shutdown_process_pools)


def drop(fun, uid, gid, *args, **kwargs):
    os.setgroups([])
    os.setgid(gid)
    os.setuid(uid)
    return fun(*args, **kwargs)


class UnsafeFileOperation(Exception):
    pass


def ensure_regular_file(path: str) -> None:
    """Verify path is a regular file; remove and raise FileNotFoundError if not.

    Uses os.lstat() to avoid following symlinks. If the file is a FIFO,
    symlink, socket, device, etc., it is deleted so the caller can
    recreate it as a regular file.
    """
    st = os.lstat(path)  # raises FileNotFoundError if missing
    if not stat.S_ISREG(st.st_mode):
        logger.warning(
            "Identity file %s is not a regular file (mode=%s), removing",
            path,
            stat.filemode(st.st_mode),
        )
        os.unlink(path)
        raise FileNotFoundError(f"Removed non-regular identity file: {path}")


def check_non_admin_file(file):
    st = os.stat(str(file))
    if st.st_uid < utils.get_min_uid():
        raise UnsafeFileOperation(
            "The file belongs to admin user: " + str(file)
        )
    return True


def _in_checked_dir(fun, dirname, dir_id, name, *args):
    # the worker resolves dirname again, so make sure it got the directory
    # that was checked before acting in it; O_PATH, as acting in it never
    # needed read permission on it
    with open_fd(dirname, os.O_PATH | os.O_DIRECTORY) as dir_fd:
        st = os.fstat(dir_fd)
        if (st.st_dev, st.st_ino) != dir_id:
            raise OSError(
                errno.ESTALE, "Directory changed after it was checked", dirname
            )
        return fun(dir_fd, name, *args)


def safe(fun):
    """Run fun(dir_fd, name, *args) for filename as its directory's owner."""

    @functools.wraps(fun)
    async def wrapper(filename, *args, loop=None):
        path = pathlib.Path(filename)
        uid = gid = None
        # Every component is resolved by the kernel, so a symlink
        # anywhere in the path redirects the operation: it would be
        # authorised against one inode and land on another, and the
        # link's planter picks the target. Only root can plant a
        # root-owned link, so those stay trusted; refuse the rest.
        # One lookup per component decides both, so the owner and the
        # directory handed to the worker come from what was checked.
        for p in reversed(path.parents):
            st = os.lstat(str(p))
            if stat.S_ISLNK(st.st_mode):
                if st.st_uid != 0:
                    raise UnsafeFileOperation(
                        "Refusing a path through a symlink: " + str(p)
                    )
                st = os.stat(str(p))
            if uid is None and st.st_uid != 0 and st.st_gid != 0:
                uid, gid = st.st_uid, st.st_gid
        if os.path.islink(filename):
            raise UnsafeFileOperation(
                "Refusing to follow a symlink: " + str(filename)
            )
        if uid is None:
            raise UnsafeFileOperation(
                "Unsafe file operation under root: " + str(path)
            )
        # st is the parent directory: the one the worker must act in
        dir_id = (st.st_dev, st.st_ino)

        loop = loop or asyncio.get_event_loop()

        return await _run_in_fresh_executor(
            loop,
            functools.partial(
                drop,
                _in_checked_dir,
                uid,
                gid,
                fun,
                str(path.parent),
                dir_id,
                path.name,
                *args,
            ),
        )

    return wrapper


def _touch(dir_fd: int, name: str):
    # O_NONBLOCK: an existing FIFO must fail here, not hang the worker
    flags = os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK
    fd = os.open(name, flags, 0o666, dir_fd=dir_fd)
    try:
        os.utime(fd)
    finally:
        os.close(fd)


def _write_text(dir_fd: int, name: str, data: str):
    fd = os.open(name, W_FLAGS | os.O_NOFOLLOW, 0o666, dir_fd=dir_fd)
    with open(fd, "w") as f:
        return f.write(data)


# This is the only way to make _write_text and _touch pickable.
# If we use decorator syntax instead - it's impossible
# to use them in multiprocessing
async def write_text(filename: str, data: str):
    return await safe(_write_text)(filename, data)


async def touch(filename: str):
    return await safe(_touch)(filename)


@contextmanager
def safe_open_file(filename, mode, user, respect_homedir=True):
    if "w" in mode:
        raise UnsafeFileOperation("'w' mode is not permitted")
    with open(filename, mode) as f:
        st = os.fstat(f.fileno())
        passwd = pwd.getpwnam(user)
        real_path = os.readlink(f"/proc/self/fd/{f.fileno()}")
        filename_str = str(filename)

        # Checking if no symlinks along the pathway...
        # Unfortunately, that is going to fail for hosters that mapped
        # /home dir to be e.g.
        # /home -> /mnt/sdb1/home
        if (filename_str != real_path) or (st.st_uid != passwd.pw_uid):
            raise UnsafeFileOperation(f"Unable to safely read {filename_str}")

        if (
            respect_homedir
            and pathlib.Path(passwd.pw_dir)
            not in pathlib.Path(filename_str).parents
        ):
            raise UnsafeFileOperation(
                f"Unable to safely read {filename_str}. "
                "File is not in user homedir"
            )
        yield f


@contextmanager
def open_fd(*args, **kwargs):
    """
    Context manager which wraps os.open and close file descriptor at the end

    :param args: positional arguments for os.open
    :param kwargs: keyword arguments for os.open
    """
    fd = os.open(*args, **kwargs)
    try:
        yield fd
    finally:
        with suppress(OSError):
            os.close(fd)


@contextmanager
def opendir_fd(name: str, *args, **kwargs):
    """
    Context manager to get a directory file descriptor
    It also checks if a directory doesn't contain a symlink in the path

    :param name: full directory name
    :param args: positional arguments for os.open
    :param kwargs: keyword arguments for os.open
    """
    with open_fd(name, *args, flags=os.O_DIRECTORY, **kwargs) as dir_fd:
        real = os.readlink("/proc/self/fd/{}".format(dir_fd))
        if name != real:
            raise UnsafeFileOperation("Operations on symlinks are prohibited")
        yield dir_fd


@contextmanager
def open_fobj(
    f: Union[str, int], dir_fd=None, flags=0, mode=None, identity=None
):
    """
    Context manager to open file object from file name or from file descriptor
    File object extended with 'st' attribute that contains os.stat_result of
    the opened file

    :param f: file name or file descriptor to open
    :param dir_fd: directory descriptor, ignored if 'f' is a file descriptor
    :param flags: flags for os.open, ignored if 'f' is a file descriptor
    :param mode: mode for built-in open
    """

    if not isinstance(f, str):
        st = _regular_stat(f, identity)
        with open(f, mode=mode, closefd=False) as fo:
            fo.st = st
            yield fo
        return

    # Pin the inode before changing permissions. O_PATH lets an owner copy
    # a mode-000 file without resolving its mutable name again as root.
    try:
        pin = os.open(f, os.O_PATH | os.O_NOFOLLOW, dir_fd=dir_fd)
    except FileNotFoundError:
        if not flags & os.O_CREAT:
            raise
        # A concurrent creator must not supply the file we are about to use.
        with open_fd(
            f, flags | os.O_EXCL | os.O_NOFOLLOW, dir_fd=dir_fd
        ) as fd:
            with open(fd, mode=mode, closefd=False) as fo:
                fo.st = os.fstat(fd)
                yield fo
        return

    with open_fd_from_pin(pin) as pinned:
        st = _regular_stat(pinned, identity)
        if flags & os.O_EXCL:
            raise FileExistsError(errno.EEXIST, os.strerror(errno.EEXIST), f)
        proc_path = "/proc/self/fd/{}".format(pinned)
        widened = False
        try:
            with suppress(OSError):
                os.chmod(proc_path, st.st_mode | stat.S_IRUSR | stat.S_IWUSR)
                widened = True
            # The proc link names our pinned inode, not a user-controlled
            # path. Access checks still run with the worker's credentials.
            with open_fd(proc_path, flags & ~os.O_CREAT) as fd:
                with open(fd, mode=mode, closefd=False) as fo:
                    fo.st = st
                    yield fo
        finally:
            if widened:
                with suppress(OSError):
                    os.chmod(proc_path, st.st_mode)


@contextmanager
def open_fd_from_pin(fd):
    try:
        yield fd
    finally:
        os.close(fd)


def _regular_stat(fd, identity=None):
    st = os.fstat(fd)
    if not stat.S_ISREG(st.st_mode):
        raise UnsafeFileOperation("Not a regular file")
    if identity is not None and (
        st.st_dev,
        st.st_ino,
        st.st_uid,
        st.st_gid,
    ) != (
        identity.st_dev,
        identity.st_ino,
        identity.st_uid,
        identity.st_gid,
    ):
        raise UnsafeFileOperation("Source identity changed")
    return st


@contextmanager
def _copy_source(name, dir_fd, trusted, expected_uid, allow_root_owner):
    # Hold the inode through the worker's open, preventing inode reuse.
    with open_fd(name, os.O_PATH | os.O_NOFOLLOW, dir_fd=dir_fd) as pin:
        st = _regular_stat(pin)
        if (
            expected_uid is not None
            and st.st_uid != expected_uid
            and not (trusted and allow_root_owner and st.st_uid == 0)
        ):
            raise UnsafeFileOperation("Source owner changed")
        if trusted:
            with open_fd("/proc/self/fd/{}".format(pin), R_FLAGS) as fd:
                yield (fd, None), st
        else:
            # The worker opens as the selected owner and compares its pinned
            # inode before any permission change, read, or destination write.
            yield (name, dir_fd), st


@contextmanager
def safe_tuple(name: str, dir_fd: int, flags: int, is_safe: bool):
    """Open a trusted regular file, or defer the name to the worker.

    Validate with O_PATH before a potentially blocking or destructive open.
    """
    if not is_safe:
        yield name, dir_fd
        return
    try:
        pin = os.open(name, os.O_PATH | os.O_NOFOLLOW, dir_fd=dir_fd)
    except FileNotFoundError:
        if not flags & os.O_CREAT:
            raise
        with open_fd(
            name, flags | os.O_EXCL | os.O_NOFOLLOW, dir_fd=dir_fd
        ) as fd:
            yield fd, None
        return
    with open_fd_from_pin(pin):
        _regular_stat(pin)
        if flags & os.O_EXCL:
            raise FileExistsError(
                errno.EEXIST, os.strerror(errno.EEXIST), name
            )
        with open_fd(
            "/proc/self/fd/{}".format(pin), flags & ~os.O_CREAT
        ) as fd:
            yield fd, None


def _move(
    src: Union[Tuple[str, int], Tuple[int, None]],
    dst: Union[Tuple[str, int], Tuple[int, None]],
    src_unlink,
    dst_overwrite,
    racecall,
    identity=None,
):
    src_f, src_dir_fd = src
    dst_f, dst_dir_fd = dst

    w_flags = W_FLAGS | (0 if dst_overwrite else os.O_EXCL)

    with open_fobj(
        src_f, dir_fd=src_dir_fd, flags=R_FLAGS, mode="rb", identity=identity
    ) as src_fo:
        with open_fobj(
            dst_f, dir_fd=dst_dir_fd, flags=w_flags, mode="wb"
        ) as dst_fo:
            # Trusted destinations were opened without truncation until the
            # worker validated and opened the source.
            os.ftruncate(dst_fo.fileno(), 0)
            if racecall:
                racecall[0]()
            shutil.copyfileobj(src_fo, dst_fo)

            if isinstance(dst_f, str):
                # safe_dst == False
                os.chmod(dst_fo.fileno(), mode=src_fo.st.st_mode)

        if src_unlink and isinstance(src_f, str):
            # safe_src == False
            if racecall:
                racecall[1]()
            os.unlink(src_f, dir_fd=src_dir_fd)


async def safe_move(
    src: str,
    dst: str,
    safe_src=False,
    safe_dst=False,
    src_unlink=True,
    dst_overwrite=False,
    racecall=None,
    expected_uid=None,
    allow_root_owner=False,
):
    src_dir, src_name = os.path.split(src)
    dst_dir, dst_name = os.path.split(dst)

    # Do not truncate a trusted destination until the source is accepted.
    w_flags = (W_FLAGS & ~os.O_TRUNC) | (0 if dst_overwrite else os.O_EXCL)

    with opendir_fd(src_dir) as src_dir_fd, opendir_fd(
        dst_dir
    ) as dst_dir_fd, _copy_source(
        src_name, src_dir_fd, safe_src, expected_uid, allow_root_owner
    ) as (
        src_tuple,
        src_st,
    ), safe_tuple(
        dst_name, dst_dir_fd, w_flags, safe_dst
    ) as dst_tuple:
        loop = asyncio.get_event_loop()
        await _run_in_fresh_executor(
            loop,
            drop,
            _move,
            src_st.st_uid,
            src_st.st_gid,
            src_tuple,
            dst_tuple,
            src_unlink,
            dst_overwrite,
            racecall,
            src_st,
        )

        if src_unlink and safe_src:
            if racecall:
                racecall[1]()
            os.unlink(src_name, dir_fd=src_dir_fd)

        if safe_dst:
            # by name these follow a symlink planted since the open, and
            # safe_dst means the directory is writable by its owner
            dst_fd = dst_tuple[0]
            os.chown(dst_fd, src_st.st_uid, src_st.st_gid)
            os.chmod(dst_fd, src_st.st_mode)


def _check_room(src_size: int, dst_occupied: int, dst_dir_fd: int) -> None:
    vfs = os.statvfs(dst_dir_fd)
    # a sparse or compressed destination frees what it occupies, not its size
    frsize = vfs.f_frsize or vfs.f_bsize or 4096
    needed = max(0, (src_size + frsize - 1) // frsize * frsize - dst_occupied)
    available = vfs.f_bavail * frsize
    if available < needed:
        raise OSError(
            errno.ENOSPC,
            "{} bytes needed, {} available".format(needed, available),
        )


def _copy_in_place(
    src: Union[Tuple[str, int], Tuple[int, None]],
    dst: Union[Tuple[str, int], Tuple[int, None]],
    identity=None,
):
    src_f, src_dir_fd = src
    dst_f, dst_dir_fd = dst

    with open_fobj(
        src_f, dir_fd=src_dir_fd, flags=R_FLAGS, mode="rb", identity=identity
    ) as src_fo:
        with open_fobj(
            dst_f, dir_fd=dst_dir_fd, flags=IN_PLACE_W_FLAGS, mode="wb"
        ) as dst_fo:
            shutil.copyfileobj(src_fo, dst_fo)
            dst_fo.flush()
            written = dst_fo.tell()
            if written != src_fo.st.st_size:
                raise OSError(
                    errno.EIO,
                    "copied {} of {} bytes".format(written, src_fo.st.st_size),
                )
            os.ftruncate(dst_fo.fileno(), written)

            if isinstance(dst_f, str):
                # safe_dst == False; open_fobj puts back the mode of a
                # destination that was already there, so this dresses only
                # one we had to create
                os.chmod(dst_fo.fileno(), mode=src_fo.st.st_mode)


async def safe_copy_in_place(
    src: str,
    dst: str,
    safe_src=False,
    safe_dst=False,
    expected_uid=None,
    allow_root_owner=False,
):
    """Copy src over dst without emptying it first: the destination keeps its
    inode, and a copy that fails leaves the bytes written so far followed by
    the tail of the previous content rather than an empty file.
    """
    src_dir, src_name = os.path.split(src)
    dst_dir, dst_name = os.path.split(dst)

    with opendir_fd(src_dir) as src_dir_fd, opendir_fd(
        dst_dir
    ) as dst_dir_fd, _copy_source(
        src_name, src_dir_fd, safe_src, expected_uid, allow_root_owner
    ) as (
        src_tuple,
        src_st,
    ):
        dst_occupied = 0
        dst_existed = True
        try:
            dst_st = os.stat(dst_name, dir_fd=dst_dir_fd)
        except FileNotFoundError:
            dst_existed = False
        else:
            dst_occupied = dst_st.st_blocks * 512
        # before opening the destination: opening it creates one we then keep
        _check_room(src_st.st_size, dst_occupied, dst_dir_fd)

        with safe_tuple(
            dst_name, dst_dir_fd, IN_PLACE_W_FLAGS, safe_dst
        ) as dst_tuple:
            loop = asyncio.get_event_loop()
            try:
                await _run_in_fresh_executor(
                    loop,
                    drop,
                    _copy_in_place,
                    src_st.st_uid,
                    src_st.st_gid,
                    src_tuple,
                    dst_tuple,
                    src_st,
                )
            except BaseException:
                if not dst_existed:
                    with suppress(OSError):
                        if not os.stat(dst_name, dir_fd=dst_dir_fd).st_size:
                            os.unlink(dst_name, dir_fd=dst_dir_fd)
                raise

            if safe_dst:
                dst_fd = dst_tuple[0]
                os.chown(dst_fd, src_st.st_uid, src_st.st_gid)
                os.chmod(dst_fd, src_st.st_mode)