// SPDX-License-Identifier: EUPL-1.2+
// SPDX-FileCopyrightText: 2026 Demi Marie Obenour <demiobenour@gmail.com>

use std::fs::File;
use std::io::{Read as _, Seek as _, Write as _};
use std::os::unix::prelude::*;

use std::path::{Path, PathBuf};

use rustix::fs::{AtFlags, XattrFlags};
use rustix::path::Arg;
use rustix::{
    fs::{Mode, OFlags, ResolveFlags},
    io::Errno,
};

pub(crate) struct LeafCgroup {
    path: PathBuf,
    fd: OwnedFd,
}

pub enum Access {
    Read,
    Write,
}

impl LeafCgroup {
    pub fn enable_delegation(&self) -> Result<(), Errno> {
        rustix::fs::fsetxattr(self.fd.as_fd(), c"user.delegate", b"1", XattrFlags::empty())
    }

    pub fn open_subtree(&self, path: &std::path::Path, access: Access) -> Result<OwnedFd, Errno> {
        rustix::fs::openat2(
            self.fd.as_fd(),
            path,
            OFlags::NOATIME
                | OFlags::CLOEXEC
                | OFlags::NOFOLLOW
                | match access {
                    Access::Write => OFlags::WRONLY,
                    Access::Read => OFlags::RDONLY,
                },
            Mode::empty(),
            ResolveFlags::NO_SYMLINKS | ResolveFlags::BENEATH | ResolveFlags::NO_XDEV,
        )
    }

    pub fn wait_for_empty(&self, kill: bool) -> std::io::Result<()> {
        let kill_fd = if kill {
            Some(self.open_subtree(std::path::Path::new("cgroup.kill"), Access::Write)?)
        } else {
            None
        };
        let wait_file = self.open_subtree(std::path::Path::new("cgroup.events"), Access::Read)?;
        let poll_fd = wait_file.as_raw_fd();
        let mut wait_fd = File::from(wait_file);
        if let Some(kill_fd) = kill_fd {
            assert_eq!(rustix::io::write(kill_fd.as_fd(), b"1")?, 1, "kernel bug");
        }
        let mut fds = libc::pollfd {
            fd: poll_fd,
            events: libc::POLLPRI | libc::POLLERR,
            revents: 0,
        };
        let mut v = vec![];
        loop {
            v.clear();
            // SAFETY: FFI call, valid arguments, fds contains 1 element
            if unsafe { libc::poll(&raw mut fds, 1, -1) } != 1 {
                panic!("poll failed");
            }
            wait_fd
                .seek(std::io::SeekFrom::Start(0))
                .expect("Seek on control group file should succeed");
            wait_fd
                .read_to_end(&mut v)
                .expect("reading from control group should work");
            if v.split(|&c| c == b'\n').any(|line| line == b"populated 0") {
                break;
            }
        }
        Ok(())
    }
}

impl AsFd for LeafCgroup {
    fn as_fd(&self) -> std::os::fd::BorrowedFd<'_> {
        self.fd.as_fd()
    }
}

impl LeafCgroup {
    pub(crate) fn open_cgroup_at(&self, p: &Path) -> Result<Self, Errno> {
        let fd = rustix::fs::openat2(
            self.fd.as_fd(),
            p,
            OFlags::NOATIME
                | OFlags::CLOEXEC
                | OFlags::NOFOLLOW
                | OFlags::RDONLY
                | OFlags::DIRECTORY,
            Mode::empty(),
            ResolveFlags::NO_SYMLINKS | ResolveFlags::BENEATH | ResolveFlags::NO_XDEV,
        )?;
        Ok(Self {
            fd,
            path: self.path.join(p),
        })
    }
    pub(crate) fn open_cgroup(fd: BorrowedFd, p: &Path) -> Result<Self, Errno> {
        let fd = rustix::fs::openat2(
            fd,
            p,
            OFlags::NOATIME
                | OFlags::CLOEXEC
                | OFlags::NOFOLLOW
                | OFlags::RDONLY
                | OFlags::DIRECTORY,
            Mode::empty(),
            ResolveFlags::NO_SYMLINKS,
        )?;
        Ok(Self {
            fd,
            path: p.to_owned(),
        })
    }
    pub(crate) fn make_child(&self, p: &Path) -> Result<(), Errno> {
        let r = rustix::fs::mkdirat(
            self.fd.as_fd(),
            p,
            Mode::RUSR
                | Mode::WUSR
                | Mode::XUSR
                | Mode::RGRP
                | Mode::XGRP
                | Mode::ROTH
                | Mode::XOTH,
        );
        if r == Err(Errno::EXIST) { Ok(()) } else { r }
    }

    pub(crate) fn delete_child(&self, path: &Path) -> Result<(), Errno> {
        remove_all(100, self.as_fd(), &path.as_cow_c_str().unwrap())
    }

    pub(crate) fn write_cgroup_value(&self, name: &str, value: &str) -> Result<(), String> {
        let path = Path::new(name);
        let fd = rustix::fs::openat2(
            self.as_fd(),
            path,
            OFlags::NOATIME | OFlags::CLOEXEC | OFlags::NOFOLLOW | OFlags::WRONLY,
            Mode::empty(),
            ResolveFlags::NO_SYMLINKS | ResolveFlags::BENEATH | ResolveFlags::NO_XDEV,
        )
        .map_err(|e| format!("Cannot open {:?}: {}", path, e))?;
        File::from(fd)
            .write_all(value.as_bytes())
            .map_err(|e| format!("Cannot write {:?} to {:?}: {}", value, path, e))
    }
}

fn remove_recursively(fd: OwnedFd, remaining_depth: usize) -> Result<(), Errno> {
    if remaining_depth < 1 {
        panic!("control groups too deeply nested");
    }
    let mut d = rustix::fs::Dir::new(fd).expect("cannot start iterating");
    while let Some(element) = d.next() {
        let element = element.expect("Iterating through a cgroup directory failed?");
        if element.file_type() != rustix::fs::FileType::Directory {
            continue;
        }

        let remaining_depth = remaining_depth - 1;
        let d: &rustix::fs::Dir = &d;
        let dirfd = d.fd().unwrap();
        let path = element.file_name();
        remove_all(remaining_depth, dirfd, path)?;
    }
    Ok(())
}

fn remove_all(
    remaining_depth: usize,
    dirfd: BorrowedFd<'_>,
    path: &std::ffi::CStr,
) -> Result<(), Errno> {
    if path == c"." || path == c".." {
        return Ok(());
    }
    if rustix::fs::unlinkat(dirfd, path, AtFlags::REMOVEDIR).is_ok() {
        return Ok(());
    }
    let fd = rustix::fs::openat2(
        dirfd,
        path,
        OFlags::NOATIME | OFlags::CLOEXEC | OFlags::NOFOLLOW | OFlags::RDONLY | OFlags::DIRECTORY,
        Mode::empty(),
        ResolveFlags::NO_SYMLINKS | ResolveFlags::BENEATH | ResolveFlags::NO_XDEV,
    )?;
    remove_recursively(fd, remaining_depth)?;
    rustix::fs::unlinkat(dirfd, path, AtFlags::REMOVEDIR)?;
    Ok(())
}