enact-context 0.0.1

//! Context Compaction
//!
//! Strategies for reducing context size when approaching limits.
//!
//! @see packages/enact-schemas/src/context.schemas.ts

use crate::segment::{ContextPriority, ContextSegment, ContextSegmentType};
use chrono::{DateTime, Utc};
use enact_core::kernel::ExecutionId;
use serde::{Deserialize, Serialize};
use thiserror::Error;

/// Compaction errors
#[derive(Debug, Error)]
pub enum CompactionError {
    #[error("Nothing to compact")]
    NothingToCompact,

    #[error("Target token count too low: {0}")]
    TargetTooLow(usize),

    #[error("Summarization failed: {0}")]
    SummarizationFailed(String),
}

/// Available compaction strategies
///
/// Matches `compactionStrategyTypeSchema` in @enact/schemas
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CompactionStrategyType {
    /// Simple truncation (remove oldest)
    Truncate,
    /// LLM summarization
    Summarize,
    /// Extract key points only
    ExtractKeyPoints,
    /// Keep only recent N messages
    SlidingWindow,
    /// Keep based on importance scores
    ImportanceWeighted,
    /// Combination of strategies
    Hybrid,
}

/// Configuration for context compaction
///
/// Matches `compactionStrategySchema` in @enact/schemas
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionStrategy {
    /// Strategy type
    #[serde(rename = "type")]
    pub strategy_type: CompactionStrategyType,

    /// Target token count after compaction
    pub target_tokens: usize,

    /// Minimum content to preserve (percentage, 0-100)
    pub min_preserve_percent: u8,

    /// Segments to compact (in priority order)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub segments_to_compact: Option<Vec<ContextSegmentType>>,

    /// Segments to never compact
    #[serde(skip_serializing_if = "Option::is_none")]
    pub protected_segments: Option<Vec<ContextSegmentType>>,

    /// For summarize strategy: max summary tokens
    #[serde(skip_serializing_if = "Option::is_none")]
    pub summary_max_tokens: Option<usize>,

    /// For sliding_window strategy: window size
    #[serde(skip_serializing_if = "Option::is_none")]
    pub window_size: Option<usize>,

    /// For importance_weighted: minimum importance score to keep
    #[serde(skip_serializing_if = "Option::is_none")]
    pub min_importance_score: Option<f64>,
}

impl CompactionStrategy {
    /// Create a truncation strategy
    pub fn truncate(target_tokens: usize) -> Self {
        Self {
            strategy_type: CompactionStrategyType::Truncate,
            target_tokens,
            min_preserve_percent: 20,
            segments_to_compact: None,
            protected_segments: Some(vec![
                ContextSegmentType::System,
                ContextSegmentType::UserInput,
            ]),
            summary_max_tokens: None,
            window_size: None,
            min_importance_score: None,
        }
    }

    /// Create a sliding window strategy
    pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
        Self {
            strategy_type: CompactionStrategyType::SlidingWindow,
            target_tokens,
            min_preserve_percent: 20,
            segments_to_compact: Some(vec![ContextSegmentType::History]),
            protected_segments: Some(vec![
                ContextSegmentType::System,
                ContextSegmentType::UserInput,
            ]),
            summary_max_tokens: None,
            window_size: Some(window_size),
            min_importance_score: None,
        }
    }

    /// Create a summarization strategy
    pub fn summarize(target_tokens: usize, summary_max_tokens: usize) -> Self {
        Self {
            strategy_type: CompactionStrategyType::Summarize,
            target_tokens,
            min_preserve_percent: 30,
            segments_to_compact: Some(vec![
                ContextSegmentType::History,
                ContextSegmentType::ToolResults,
            ]),
            protected_segments: Some(vec![
                ContextSegmentType::System,
                ContextSegmentType::UserInput,
                ContextSegmentType::Guidance,
            ]),
            summary_max_tokens: Some(summary_max_tokens),
            window_size: None,
            min_importance_score: None,
        }
    }

    /// Check if a segment type is protected
    pub fn is_protected(&self, segment_type: ContextSegmentType) -> bool {
        self.protected_segments
            .as_ref()
            .map(|p| p.contains(&segment_type))
            .unwrap_or(false)
    }

    /// Check if a segment type should be compacted
    pub fn should_compact(&self, segment_type: ContextSegmentType) -> bool {
        if self.is_protected(segment_type) {
            return false;
        }

        self.segments_to_compact
            .as_ref()
            .map(|s| s.contains(&segment_type))
            .unwrap_or(true) // If not specified, compact all non-protected
    }
}

/// Result of a compaction operation
///
/// Matches `compactionResultSchema` in @enact/schemas
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionResult {
    /// Execution ID
    pub execution_id: ExecutionId,

    /// Strategy used
    pub strategy: CompactionStrategyType,

    /// Tokens before compaction
    pub tokens_before: usize,

    /// Tokens after compaction
    pub tokens_after: usize,

    /// Tokens saved
    pub tokens_saved: usize,

    /// Compression ratio (tokensAfter / tokensBefore)
    pub compression_ratio: f64,

    /// Number of segments compacted
    pub segments_compacted: usize,

    /// Duration in milliseconds
    pub duration_ms: u64,

    /// Whether compaction was successful
    pub success: bool,

    /// Error message if failed
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,

    /// Timestamp
    pub compacted_at: DateTime<Utc>,
}

impl CompactionResult {
    /// Create a successful result
    pub fn success(
        execution_id: ExecutionId,
        strategy: CompactionStrategyType,
        tokens_before: usize,
        tokens_after: usize,
        segments_compacted: usize,
        duration_ms: u64,
    ) -> Self {
        let tokens_saved = tokens_before.saturating_sub(tokens_after);
        let compression_ratio = if tokens_before > 0 {
            tokens_after as f64 / tokens_before as f64
        } else {
            1.0
        };

        Self {
            execution_id,
            strategy,
            tokens_before,
            tokens_after,
            tokens_saved,
            compression_ratio,
            segments_compacted,
            duration_ms,
            success: true,
            error: None,
            compacted_at: Utc::now(),
        }
    }

    /// Create a failed result
    pub fn failure(
        execution_id: ExecutionId,
        strategy: CompactionStrategyType,
        tokens_before: usize,
        error: String,
        duration_ms: u64,
    ) -> Self {
        Self {
            execution_id,
            strategy,
            tokens_before,
            tokens_after: tokens_before,
            tokens_saved: 0,
            compression_ratio: 1.0,
            segments_compacted: 0,
            duration_ms,
            success: false,
            error: Some(error),
            compacted_at: Utc::now(),
        }
    }
}

/// Compactor - applies compaction strategies to segments
pub struct Compactor {
    strategy: CompactionStrategy,
}

impl Compactor {
    /// Create a new compactor with the given strategy
    pub fn new(strategy: CompactionStrategy) -> Self {
        Self { strategy }
    }

    /// Create a truncation compactor
    pub fn truncate(target_tokens: usize) -> Self {
        Self::new(CompactionStrategy::truncate(target_tokens))
    }

    /// Create a sliding window compactor
    pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
        Self::new(CompactionStrategy::sliding_window(
            target_tokens,
            window_size,
        ))
    }

    /// Get the strategy
    pub fn strategy(&self) -> &CompactionStrategy {
        &self.strategy
    }

    /// Compact segments using truncation strategy
    ///
    /// Removes oldest segments (lowest priority first) until target is reached.
    pub fn compact_truncate(
        &self,
        segments: &mut Vec<ContextSegment>,
        current_tokens: usize,
    ) -> Result<usize, CompactionError> {
        if current_tokens <= self.strategy.target_tokens {
            return Ok(0);
        }

        let tokens_to_remove = current_tokens - self.strategy.target_tokens;
        let mut removed = 0;

        // Sort by priority (ascending) then by sequence (ascending = oldest first)
        segments.sort_by(|a, b| {
            a.priority
                .cmp(&b.priority)
                .then(a.sequence.cmp(&b.sequence))
        });

        // Remove lowest priority, oldest segments first
        let mut i = 0;
        while i < segments.len() && removed < tokens_to_remove {
            let segment = &segments[i];

            // Skip protected segments
            if !segment.compressible || self.strategy.is_protected(segment.segment_type) {
                i += 1;
                continue;
            }

            // Skip critical priority
            if segment.priority == ContextPriority::Critical {
                i += 1;
                continue;
            }

            removed += segment.token_count;
            segments.remove(i);
        }

        Ok(removed)
    }

    /// Compact using sliding window strategy
    ///
    /// Keeps only the most recent N messages in the history.
    pub fn compact_sliding_window(
        &self,
        segments: &mut Vec<ContextSegment>,
    ) -> Result<usize, CompactionError> {
        let window_size = self.strategy.window_size.unwrap_or(10);

        // Find history segments
        let history_indices: Vec<usize> = segments
            .iter()
            .enumerate()
            .filter(|(_, s)| s.segment_type == ContextSegmentType::History)
            .map(|(i, _)| i)
            .collect();

        if history_indices.len() <= window_size {
            return Ok(0);
        }

        // Remove oldest history segments (keep window_size most recent)
        let to_remove = history_indices.len() - window_size;
        let mut removed_tokens = 0;

        // Remove from oldest first (indices are in ascending order)
        for &idx in history_indices.iter().take(to_remove).rev() {
            removed_tokens += segments[idx].token_count;
            segments.remove(idx);
        }

        Ok(removed_tokens)
    }
}

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

    #[test]
    fn test_truncation_strategy() {
        let strategy = CompactionStrategy::truncate(1000);
        assert_eq!(strategy.strategy_type, CompactionStrategyType::Truncate);
        assert!(strategy.is_protected(ContextSegmentType::System));
        assert!(!strategy.is_protected(ContextSegmentType::History));
    }

    #[test]
    fn test_compaction_result() {
        let exec_id = ExecutionId::new();
        let result = CompactionResult::success(
            exec_id,
            CompactionStrategyType::Truncate,
            10000,
            5000,
            5,
            100,
        );

        assert!(result.success);
        assert_eq!(result.tokens_saved, 5000);
        assert!((result.compression_ratio - 0.5).abs() < 0.01);
    }
}