//! Thin RAII wrapper around a Windows Job Object configured for per-cell
//! containment: kill-on-close, an optional job-wide memory cap, and an optional
//! hard CPU-rate cap.
//!
//! All `unsafe` is confined to this module. Every call is a single Win32 API
//! invocation whose pre/post-conditions are documented inline.

use std::ffi::c_void;

use windows::core::PCWSTR;
use windows::Win32::Foundation::{CloseHandle, HANDLE};
use windows::Win32::System::JobObjects::{
    AssignProcessToJobObject, CreateJobObjectW, JobObjectCpuRateControlInformation,
    JobObjectExtendedLimitInformation, SetInformationJobObject, TerminateJobObject,
    JOBOBJECT_CPU_RATE_CONTROL_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
    JOB_OBJECT_CPU_RATE_CONTROL_ENABLE, JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP,
    JOB_OBJECT_LIMIT_JOB_MEMORY, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
};

use cellos_core::CellosError;

/// Owned handle to a Windows Job Object. Dropping it closes the handle, which —
/// because the job is created with `KILL_ON_JOB_CLOSE` — terminates any process
/// still inside the job. This is the backend's residue backstop.
pub(crate) struct JobObject {
    handle: HANDLE,
}

// SAFETY: a Job Object `HANDLE` is a process-wide kernel handle with no thread
// affinity, and the Win32 Job Object APIs are safe to call from any thread. The
// backend serialises all access through an async `Mutex`, and the handle is
// closed exactly once (in `Drop`). Moving the handle between tokio worker
// threads is therefore sound.
unsafe impl Send for JobObject {}
unsafe impl Sync for JobObject {}

impl JobObject {
    /// Create a fresh unnamed job with `KILL_ON_JOB_CLOSE` set, plus the
    /// requested memory / CPU caps.
    ///
    /// `cpu_rate_hundredths_percent` is the hard CPU cap in units of 1/100 of a
    /// percent of total CPU across all cores (the unit Windows' CPU-rate control
    /// uses), in the range `(0, 10000]`. `None` leaves CPU uncapped.
    pub(crate) fn create(
        memory_max_bytes: Option<u64>,
        cpu_rate_hundredths_percent: Option<u32>,
    ) -> Result<Self, CellosError> {
        // SAFETY: a null name + null security attributes creates a fresh,
        // process-owned unnamed job. The returned handle is checked by `?`.
        let handle = unsafe { CreateJobObjectW(None, PCWSTR::null()) }
            .map_err(|e| CellosError::Host(format!("CreateJobObjectW failed: {e}")))?;
        let job = Self { handle };

        // Extended limits: always KILL_ON_JOB_CLOSE; add a job-wide memory cap
        // when the spec requested one.
        let mut ext = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default();
        ext.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        if let Some(bytes) = memory_max_bytes {
            ext.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_JOB_MEMORY;
            ext.JobMemoryLimit = bytes as usize;
        }
        // SAFETY: `ext` is a fully-initialised struct that outlives the call;
        // the length we pass matches its size exactly.
        unsafe {
            SetInformationJobObject(
                job.handle,
                JobObjectExtendedLimitInformation,
                &ext as *const _ as *const c_void,
                std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
            )
        }
        .map_err(|e| {
            CellosError::Host(format!(
                "SetInformationJobObject(extended limits) failed: {e}"
            ))
        })?;

        if let Some(rate) = cpu_rate_hundredths_percent {
            // JOBOBJECT_CPU_RATE_CONTROL_INFORMATION carries a union, so it has
            // no derived Default; zero-initialise then set the fields we use.
            // SAFETY: the struct is plain-old-data and all-zero is a valid bit
            // pattern (it is the "no control flags set" state before we fill it).
            let mut cpu: JOBOBJECT_CPU_RATE_CONTROL_INFORMATION = unsafe { std::mem::zeroed() };
            cpu.ControlFlags =
                JOB_OBJECT_CPU_RATE_CONTROL_ENABLE | JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP;
            cpu.Anonymous.CpuRate = rate.clamp(1, 10_000);
            // SAFETY: same contract as the extended-limits call above.
            unsafe {
                SetInformationJobObject(
                    job.handle,
                    JobObjectCpuRateControlInformation,
                    &cpu as *const _ as *const c_void,
                    std::mem::size_of::<JOBOBJECT_CPU_RATE_CONTROL_INFORMATION>() as u32,
                )
            }
            .map_err(|e| {
                CellosError::Host(format!("SetInformationJobObject(cpu rate) failed: {e}"))
            })?;
        }

        Ok(job)
    }

    /// Assign an already-spawned process to this job. Any process the child
    /// later spawns inherits job membership, so the whole tree is contained.
    pub(crate) fn assign(&self, process: HANDLE) -> Result<(), CellosError> {
        // SAFETY: `self.handle` is a live job handle owned by `self`; `process`
        // is the child's handle, valid until the child is reaped by `wait`.
        unsafe { AssignProcessToJobObject(self.handle, process) }
            .map_err(|e| CellosError::Host(format!("AssignProcessToJobObject failed: {e}")))
    }

    /// Force-terminate every process still in the job with `exit_code`.
    pub(crate) fn terminate(&self, exit_code: u32) -> Result<(), CellosError> {
        // SAFETY: `self.handle` is a live job handle owned by `self`.
        unsafe { TerminateJobObject(self.handle, exit_code) }
            .map_err(|e| CellosError::Host(format!("TerminateJobObject failed: {e}")))
    }
}

impl Drop for JobObject {
    fn drop(&mut self) {
        // Closing the last handle to a KILL_ON_JOB_CLOSE job terminates any
        // process still inside it — the residue backstop.
        // SAFETY: `self.handle` was created in `create` and is closed exactly
        // once, here.
        unsafe {
            let _ = CloseHandle(self.handle);
        }
    }
}
