alma 0.1.1

//! Workspace I/O limits and per-update request accounting.

use crate::plugin::PluginIdentity;
use std::{
    fmt::{Debug, Display, Formatter},
    num::{NonZeroU64, NonZeroUsize},
};

use super::PluginWorkspaceIoCommitError;

/// Default accepted workspace requests per plugin update.
const DEFAULT_WORKSPACE_IO_REQUESTS_PER_UPDATE: usize = 32;

/// Workspace I/O budget field names with stable diagnostic spellings.
#[derive(Clone, Copy, Eq, PartialEq)]
pub enum PluginWorkspaceIoBudgetField {
    /// Maximum accepted workspace requests per update.
    MaxRequests,
    /// Maximum bytes returned by one workspace read.
    MaxReadBytes,
    /// Maximum bytes accepted by one workspace write.
    MaxWriteBytes,
}

impl PluginWorkspaceIoBudgetField {
    /// Stable field name used in diagnostics.
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::MaxRequests => "max_requests",
            Self::MaxReadBytes => "max_read_bytes",
            Self::MaxWriteBytes => "max_write_bytes",
        }
    }
}

impl Display for PluginWorkspaceIoBudgetField {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(self.as_str())
    }
}

impl Debug for PluginWorkspaceIoBudgetField {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(self.as_str())
    }
}

/// Non-zero workspace read byte cap.
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct PluginWorkspaceReadByteLimit(NonZeroU64);

impl PluginWorkspaceReadByteLimit {
    /// Creates a read byte cap.
    ///
    /// # Errors
    ///
    /// Returns [`PluginWorkspaceIoCommitError::ZeroLimit`] when `max_bytes` is zero.
    pub const fn try_new(max_bytes: u64) -> Result<Self, PluginWorkspaceIoCommitError> {
        let Some(max_bytes) = NonZeroU64::new(max_bytes) else {
            return Err(PluginWorkspaceIoCommitError::ZeroLimit {
                field: PluginWorkspaceIoBudgetField::MaxReadBytes,
            });
        };
        Ok(Self(max_bytes))
    }

    /// Returns the retained byte cap.
    #[must_use]
    pub const fn get(self) -> u64 {
        self.0.get()
    }
}

impl Debug for PluginWorkspaceReadByteLimit {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_tuple("PluginWorkspaceReadByteLimit")
            .field(&self.get())
            .finish()
    }
}

/// Non-zero workspace write byte cap.
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct PluginWorkspaceWriteByteLimit(NonZeroU64);

impl PluginWorkspaceWriteByteLimit {
    /// Creates a write byte cap.
    ///
    /// # Errors
    ///
    /// Returns [`PluginWorkspaceIoCommitError::ZeroLimit`] when `max_bytes` is zero.
    pub const fn try_new(max_bytes: u64) -> Result<Self, PluginWorkspaceIoCommitError> {
        let Some(max_bytes) = NonZeroU64::new(max_bytes) else {
            return Err(PluginWorkspaceIoCommitError::ZeroLimit {
                field: PluginWorkspaceIoBudgetField::MaxWriteBytes,
            });
        };
        Ok(Self(max_bytes))
    }

    /// Returns the retained byte cap.
    #[must_use]
    pub const fn get(self) -> u64 {
        self.0.get()
    }
}

impl Debug for PluginWorkspaceWriteByteLimit {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_tuple("PluginWorkspaceWriteByteLimit")
            .field(&self.get())
            .finish()
    }
}

/// Bounded workspace I/O budget for one successful plugin update.
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct PluginWorkspaceIoBudget {
    /// Maximum requests per update.
    requests: NonZeroUsize,
    /// Maximum bytes read per file.
    read_bytes: PluginWorkspaceReadByteLimit,
    /// Maximum write payload bytes.
    write_bytes: PluginWorkspaceWriteByteLimit,
}

impl PluginWorkspaceIoBudget {
    /// Creates a workspace I/O budget.
    ///
    /// # Errors
    ///
    /// Returns [`PluginWorkspaceIoCommitError`] when a limit is zero.
    pub const fn try_new(
        max_requests: usize,
        max_read_bytes: u64,
        max_write_bytes: u64,
    ) -> Result<Self, PluginWorkspaceIoCommitError> {
        let Some(requests) = NonZeroUsize::new(max_requests) else {
            return Err(PluginWorkspaceIoCommitError::ZeroLimit {
                field: PluginWorkspaceIoBudgetField::MaxRequests,
            });
        };
        let Some(read_bytes) = NonZeroU64::new(max_read_bytes) else {
            return Err(PluginWorkspaceIoCommitError::ZeroLimit {
                field: PluginWorkspaceIoBudgetField::MaxReadBytes,
            });
        };
        let Some(write_bytes) = NonZeroU64::new(max_write_bytes) else {
            return Err(PluginWorkspaceIoCommitError::ZeroLimit {
                field: PluginWorkspaceIoBudgetField::MaxWriteBytes,
            });
        };
        Ok(Self {
            requests,
            read_bytes: PluginWorkspaceReadByteLimit(read_bytes),
            write_bytes: PluginWorkspaceWriteByteLimit(write_bytes),
        })
    }

    /// Derives a conservative v1 budget from the runtime message cap.
    ///
    /// # Errors
    ///
    /// Returns [`PluginWorkspaceIoCommitError`] when the cap is zero.
    pub const fn from_max_message_bytes(
        max_message_bytes: u64,
    ) -> Result<Self, PluginWorkspaceIoCommitError> {
        Self::try_new(
            DEFAULT_WORKSPACE_IO_REQUESTS_PER_UPDATE,
            max_message_bytes,
            max_message_bytes,
        )
    }

    /// Derives a conservative v1 budget from an already-validated runtime message cap.
    #[must_use]
    pub(in crate::plugin) const fn from_max_message_byte_limit(
        max_message_bytes: NonZeroU64,
    ) -> Self {
        let requests = match NonZeroUsize::new(DEFAULT_WORKSPACE_IO_REQUESTS_PER_UPDATE) {
            Some(requests) => requests,
            None => NonZeroUsize::MIN,
        };
        Self {
            requests,
            read_bytes: PluginWorkspaceReadByteLimit(max_message_bytes),
            write_bytes: PluginWorkspaceWriteByteLimit(max_message_bytes),
        }
    }

    /// Returns the maximum requests accepted from one update.
    #[must_use]
    pub const fn max_requests(self) -> usize {
        self.requests.get()
    }

    /// Returns the request cap as a non-zero proof for internal diagnostics.
    #[must_use]
    pub(in crate::plugin::host) const fn max_requests_limit(self) -> NonZeroUsize {
        self.requests
    }

    /// Returns the maximum bytes read from one file.
    #[must_use]
    pub const fn max_read_bytes(self) -> u64 {
        self.read_bytes.get()
    }

    /// Returns the read byte cap as a non-zero proof for internal task ownership.
    #[must_use]
    pub(in crate::plugin::host) const fn max_read_bytes_limit(
        self,
    ) -> PluginWorkspaceReadByteLimit {
        self.read_bytes
    }

    /// Returns the maximum bytes accepted for one write payload.
    #[must_use]
    pub const fn max_write_bytes(self) -> u64 {
        self.write_bytes.get()
    }

    /// Returns the write byte cap as a non-zero proof for internal filesystem writes.
    #[must_use]
    pub(in crate::plugin::host) const fn max_write_bytes_limit(
        self,
    ) -> PluginWorkspaceWriteByteLimit {
        self.write_bytes
    }
}

impl Debug for PluginWorkspaceIoBudget {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("PluginWorkspaceIoBudget")
            .field("requests", &self.max_requests())
            .field("read_bytes", &self.max_read_bytes())
            .field("write_bytes", &self.max_write_bytes())
            .finish()
    }
}

/// Request accounting shared by deferred I/O and non-deferred workspace reads.
#[derive(Eq, PartialEq)]
pub(in crate::plugin::host) struct PluginWorkspaceIoRequestLedger {
    /// Stable plugin identity for diagnostics.
    identity: PluginIdentity,
    /// Per-update workspace I/O budget.
    budget: PluginWorkspaceIoBudget,
    /// Read requests charged outside the deferred I/O batch.
    non_deferred_requests: usize,
}

/// Request ledger state split at the successful-return or discard boundary.
pub(in crate::plugin::host) struct PluginWorkspaceIoRequestLedgerParts {
    /// Stable plugin identity for diagnostics.
    pub identity: PluginIdentity,
    /// Per-update workspace I/O budget.
    pub budget: PluginWorkspaceIoBudget,
}

impl PluginWorkspaceIoRequestLedger {
    /// Creates an empty request ledger.
    pub(in crate::plugin::host) const fn new(
        identity: PluginIdentity,
        budget: PluginWorkspaceIoBudget,
    ) -> Self {
        Self {
            identity,
            budget,
            non_deferred_requests: 0,
        }
    }

    /// Returns the validated identity that owns this accounting boundary.
    #[must_use]
    pub(in crate::plugin::host) const fn identity_proof(&self) -> &PluginIdentity {
        &self.identity
    }

    /// Returns non-deferred read requests charged against the cap.
    #[must_use]
    pub(in crate::plugin::host) const fn non_deferred_request_count(&self) -> usize {
        self.non_deferred_requests
    }

    /// Returns all workspace requests accepted during this update session.
    #[must_use]
    pub(in crate::plugin::host) const fn accepted_request_count(
        &self,
        deferred_requests: usize,
    ) -> usize {
        match deferred_requests.checked_add(self.non_deferred_requests) {
            Some(count) => count,
            None => usize::MAX,
        }
    }

    /// Returns the read byte cap shared by deferred reads and direct guest observations.
    #[must_use]
    pub(in crate::plugin::host) const fn max_read_bytes(&self) -> u64 {
        self.budget.max_read_bytes()
    }

    /// Returns the read byte cap as a non-zero proof.
    #[must_use]
    pub(in crate::plugin::host) const fn max_read_bytes_limit(
        &self,
    ) -> PluginWorkspaceReadByteLimit {
        self.budget.max_read_bytes_limit()
    }

    /// Returns the write byte cap as a non-zero proof.
    #[must_use]
    pub(in crate::plugin::host) const fn max_write_bytes_limit(
        &self,
    ) -> PluginWorkspaceWriteByteLimit {
        self.budget.max_write_bytes_limit()
    }

    /// Checks whether one more deferred workspace operation fits the shared cap.
    ///
    /// # Errors
    ///
    /// Returns [`PluginWorkspaceIoCommitError`] when the session has already spent its request
    /// budget.
    pub(in crate::plugin::host) fn ensure_deferred_request_budget(
        &self,
        deferred_requests: usize,
    ) -> Result<(), PluginWorkspaceIoCommitError> {
        self.ensure_request_budget(deferred_requests)
    }

    /// Prepares a read request whose durable work is tracked by another pending batch.
    ///
    /// # Errors
    ///
    /// Returns [`PluginWorkspaceIoCommitError`] when the session has already spent its request
    /// budget.
    pub(in crate::plugin::host) fn prepare_non_deferred_read(
        &mut self,
        deferred_requests: usize,
    ) -> Result<PluginWorkspaceNonDeferredReadPermit<'_>, PluginWorkspaceIoCommitError> {
        self.ensure_request_budget(deferred_requests)?;
        Ok(PluginWorkspaceNonDeferredReadPermit { ledger: self })
    }

    /// Consumes the ledger into the identity and limit proof needed by the next boundary.
    #[must_use]
    pub(in crate::plugin::host) fn into_parts(self) -> PluginWorkspaceIoRequestLedgerParts {
        PluginWorkspaceIoRequestLedgerParts {
            identity: self.identity,
            budget: self.budget,
        }
    }

    /// Checks the shared per-update workspace request cap.
    fn ensure_request_budget(
        &self,
        deferred_requests: usize,
    ) -> Result<(), PluginWorkspaceIoCommitError> {
        if self.accepted_request_count(deferred_requests) >= self.budget.max_requests() {
            return Err(PluginWorkspaceIoCommitError::TooManyRequests {
                identity: self.identity.clone(),
                limit: self.budget.max_requests_limit(),
            });
        }
        Ok(())
    }
}

impl Debug for PluginWorkspaceIoRequestLedger {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("PluginWorkspaceIoRequestLedger")
            .field("identity", &self.identity)
            .field("budget", &self.budget)
            .field("non_deferred_requests", &self.non_deferred_requests)
            .finish()
    }
}

/// Checked but uncommitted non-deferred workspace read charge.
///
/// The permit lets callers validate the shared per-update request cap before reserving another
/// owner resource, while charging the cap only after that reservation succeeds.
#[must_use = "dropping the permit leaves the workspace read request uncharged"]
pub(in crate::plugin::host) struct PluginWorkspaceNonDeferredReadPermit<'ledger> {
    /// Ledger that issued this request permit.
    ledger: &'ledger mut PluginWorkspaceIoRequestLedger,
}

impl PluginWorkspaceNonDeferredReadPermit<'_> {
    /// Returns the read byte cap retained by this update.
    #[must_use]
    pub(in crate::plugin::host) const fn max_read_bytes_limit(
        &self,
    ) -> PluginWorkspaceReadByteLimit {
        self.ledger.max_read_bytes_limit()
    }

    /// Commits the pending request charge and returns the read byte cap.
    pub(in crate::plugin::host) const fn commit(self) -> PluginWorkspaceReadByteLimit {
        let max_read_bytes = self.max_read_bytes_limit();
        self.ledger.non_deferred_requests += 1;
        max_read_bytes
    }
}