pe-core 0.1.0

//! Cognitive budget — resource limits for the inner reasoning graph.
//!
//! Controls how much compute the cognitive graph gets per outer-graph step.
//! Prevents runaway costs from inner reasoning while allowing users to
//! tune the think-vs-do balance.

use serde::{Deserialize, Serialize};
use std::time::{Duration, Instant};

/// Resource limits for cognitive graph execution.
///
/// # Example
///
/// ```
/// use pe_core::cognitive_budget::CognitiveBudget;
///
/// let budget = CognitiveBudget::default();
/// assert_eq!(budget.max_tokens, 2000);
/// assert_eq!(budget.max_iterations, 5);
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct CognitiveBudget {
    /// Max tokens for all cognitive processing combined.
    pub max_tokens: u32,

    /// Max wall-clock time for cognitive processing.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        with = "crate::cognitive_budget::duration_serde"
    )]
    pub max_duration: Option<Duration>,

    /// Max number of cognitive graph iterations per outer step.
    pub max_iterations: u32,

    /// What fraction of the agent's total budget goes to cognitive (0.0-1.0).
    pub budget_fraction: f64,
}

impl Default for CognitiveBudget {
    fn default() -> Self {
        Self {
            max_tokens: 2000,
            max_duration: Some(Duration::from_secs(10)),
            max_iterations: 5,
            budget_fraction: 0.15,
        }
    }
}

/// Tracks consumption against a [`CognitiveBudget`].
///
/// Created at the start of cognitive processing, consumed as lobes execute.
///
/// # Example
///
/// ```
/// use pe_core::cognitive_budget::{CognitiveBudget, BudgetTracker};
///
/// let budget = CognitiveBudget::default();
/// let mut tracker = BudgetTracker::new(&budget);
/// assert!(!tracker.is_exhausted());
/// tracker.consume_tokens(1500);
/// assert!(!tracker.is_exhausted());
/// tracker.consume_tokens(600);
/// assert!(tracker.is_exhausted());
/// ```
#[derive(Debug)]
pub struct BudgetTracker {
    max_tokens: u32,
    max_duration: Option<Duration>,
    tokens_used: u32,
    start: Instant,
}

impl BudgetTracker {
    /// Create a new tracker from a budget. Starts the clock immediately.
    pub fn new(budget: &CognitiveBudget) -> Self {
        Self {
            max_tokens: budget.max_tokens,
            max_duration: budget.max_duration,
            tokens_used: 0,
            start: Instant::now(),
        }
    }

    /// Record token consumption.
    pub fn consume_tokens(&mut self, tokens: u32) {
        self.tokens_used = self.tokens_used.saturating_add(tokens);
    }

    /// Tokens remaining before limit.
    pub fn remaining_tokens(&self) -> u32 {
        self.max_tokens.saturating_sub(self.tokens_used)
    }

    /// Total tokens consumed so far.
    pub fn tokens_used(&self) -> u32 {
        self.tokens_used
    }

    /// Elapsed wall-clock time since tracker creation.
    pub fn elapsed(&self) -> Duration {
        self.start.elapsed()
    }

    /// Whether the budget is exhausted (tokens OR time).
    pub fn is_exhausted(&self) -> bool {
        if self.tokens_used >= self.max_tokens {
            return true;
        }
        if let Some(max_dur) = self.max_duration {
            if self.start.elapsed() >= max_dur {
                return true;
            }
        }
        false
    }
}

/// Serde support for `Option<Duration>` as seconds.
mod duration_serde {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::time::Duration;

    pub fn serialize<S: Serializer>(
        duration: &Option<Duration>,
        serializer: S,
    ) -> Result<S::Ok, S::Error> {
        duration.map(|d| d.as_secs_f64()).serialize(serializer)
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(
        deserializer: D,
    ) -> Result<Option<Duration>, D::Error> {
        let secs: Option<f64> = Option::deserialize(deserializer)?;
        Ok(secs.map(Duration::from_secs_f64))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_budget_defaults() {
        let budget = CognitiveBudget::default();
        assert_eq!(budget.max_tokens, 2000);
        assert_eq!(budget.max_iterations, 5);
        assert!((budget.budget_fraction - 0.15).abs() < f64::EPSILON);
        assert_eq!(budget.max_duration, Some(Duration::from_secs(10)));
    }

    #[test]
    fn test_tracker_token_consumption() {
        let budget = CognitiveBudget::default();
        let mut tracker = BudgetTracker::new(&budget);
        assert_eq!(tracker.remaining_tokens(), 2000);
        assert_eq!(tracker.tokens_used(), 0);

        tracker.consume_tokens(500);
        assert_eq!(tracker.remaining_tokens(), 1500);
        assert_eq!(tracker.tokens_used(), 500);
    }

    #[test]
    fn test_tracker_exhaustion_by_tokens() {
        let budget = CognitiveBudget {
            max_tokens: 100,
            max_duration: None,
            ..Default::default()
        };
        let mut tracker = BudgetTracker::new(&budget);
        assert!(!tracker.is_exhausted());

        tracker.consume_tokens(99);
        assert!(!tracker.is_exhausted());

        tracker.consume_tokens(1);
        assert!(tracker.is_exhausted());
    }

    #[test]
    fn test_tracker_token_saturation() {
        let budget = CognitiveBudget {
            max_tokens: 100,
            max_duration: None,
            ..Default::default()
        };
        let mut tracker = BudgetTracker::new(&budget);
        tracker.consume_tokens(u32::MAX);
        assert!(tracker.is_exhausted());
        assert_eq!(tracker.remaining_tokens(), 0);
    }

    #[test]
    fn test_budget_serialization() {
        let budget = CognitiveBudget::default();
        let json = serde_json::to_string(&budget).unwrap();
        let back: CognitiveBudget = serde_json::from_str(&json).unwrap();
        assert_eq!(back.max_tokens, budget.max_tokens);
        assert_eq!(back.max_iterations, budget.max_iterations);
        assert_eq!(back.max_duration, budget.max_duration);
        assert!((back.budget_fraction - budget.budget_fraction).abs() < f64::EPSILON);
    }
}