stemedb/crates/stemedb-lens/src/constraints.rs

//! Constraints Lens: Pre-flight check for must_use/forbidden/prefer.
//!
//! This lens categorizes assertions by predicate pattern for agent constraint checking.
//! Instead of picking a single winner, it groups assertions into constraint categories.
//!
//! # The Problem
//!
//! AI agents need to check constraints before taking actions:
//! - "What libraries MUST I use for this project?"
//! - "What tools are FORBIDDEN?"
//! - "What patterns are PREFERRED?"
//!
//! # Predicate Patterns
//!
//! | Pattern | Category | Meaning |
//! |---------|----------|---------|
//! | `must_use:*` | must_use | Required, non-negotiable |
//! | `forbidden:*` | forbidden | Explicitly banned |
//! | `prefer:*` | prefer | Recommended but optional |
//!
//! # Example
//!
//! ```ignore
//! // Assertions:
//! // - predicate: "must_use:http_client" -> object: "axios"
//! // - predicate: "forbidden:http_client" -> object: "requests"
//! // - predicate: "prefer:language" -> object: "typescript"
//!
//! let lens = ConstraintsLens;
//! let constraints = lens.resolve_constraints(&candidates);
//! // constraints.must_use = [axios assertion]
//! // constraints.forbidden = [requests assertion]
//! // constraints.prefer = [typescript assertion]
//! ```

use crate::traits::{compute_conflict_score, Lens, Resolution};
use stemedb_core::types::Assertion;
use tracing::instrument;

/// A set of categorized constraints from assertions.
///
/// Each category contains assertions matching the corresponding predicate pattern,
/// sorted by confidence (highest first).
#[derive(Debug, Clone, Default)]
pub struct ConstraintSet {
    /// Assertions with `must_use:*` predicates. Required constraints.
    pub must_use: Vec<Assertion>,

    /// Assertions with `forbidden:*` predicates. Banned items.
    pub forbidden: Vec<Assertion>,

    /// Assertions with `prefer:*` predicates. Recommendations.
    pub prefer: Vec<Assertion>,

    /// Total candidates considered (including non-constraint predicates).
    pub candidates_count: usize,

    /// Overall conflict score across all constraint categories.
    pub conflict_score: f32,
}

impl ConstraintSet {
    /// Create an empty constraint set.
    pub fn empty() -> Self {
        Self::default()
    }

    /// Check if any constraints exist.
    pub fn has_constraints(&self) -> bool {
        !self.must_use.is_empty() || !self.forbidden.is_empty() || !self.prefer.is_empty()
    }

    /// Total number of constraint assertions.
    pub fn total_constraints(&self) -> usize {
        self.must_use.len() + self.forbidden.len() + self.prefer.len()
    }
}

/// Constraints Lens: Categorizes assertions by predicate pattern.
///
/// # Resolution Strategy (for Lens trait)
///
/// When used as a standard `Lens`:
/// 1. Returns the highest-confidence `must_use` assertion as the winner
/// 2. If no `must_use`, returns highest-confidence `forbidden`
/// 3. If neither, returns highest-confidence `prefer`
/// 4. If no constraint predicates at all, returns empty
///
/// # Rich Result
///
/// For the full constraint set, use `resolve_constraints()` instead.
///
/// # Predicate Patterns
///
/// - `must_use:*` -> must_use category
/// - `forbidden:*` -> forbidden category
/// - `prefer:*` -> prefer category
/// - Other predicates are ignored (not categorized)
#[derive(Debug, Clone, Copy, Default)]
pub struct ConstraintsLens;

impl ConstraintsLens {
    /// Resolve constraints into a categorized set.
    ///
    /// This is the rich result method that returns all constraint categories.
    /// Use this when you need the full constraint picture.
    ///
    /// # Arguments
    ///
    /// * `candidates` - Assertions to categorize by predicate pattern
    ///
    /// # Returns
    ///
    /// A `ConstraintSet` with:
    /// - `must_use`: Assertions with `must_use:*` predicates (sorted by confidence desc)
    /// - `forbidden`: Assertions with `forbidden:*` predicates (sorted by confidence desc)
    /// - `prefer`: Assertions with `prefer:*` predicates (sorted by confidence desc)
    #[instrument(skip(self, candidates), fields(candidates_count = candidates.len(), lens = "Constraints"))]
    pub fn resolve_constraints(&self, candidates: &[Assertion]) -> ConstraintSet {
        if candidates.is_empty() {
            return ConstraintSet::empty();
        }

        let mut must_use = Vec::new();
        let mut forbidden = Vec::new();
        let mut prefer = Vec::new();

        // Categorize by predicate pattern
        for assertion in candidates {
            if assertion.predicate.starts_with("must_use:") {
                must_use.push(assertion.clone());
            } else if assertion.predicate.starts_with("forbidden:") {
                forbidden.push(assertion.clone());
            } else if assertion.predicate.starts_with("prefer:") {
                prefer.push(assertion.clone());
            }
            // Other predicates are not constraint predicates, ignore them
        }

        // Sort each category by confidence (highest first), then by timestamp (newest first)
        let sort_by_confidence = |a: &Assertion, b: &Assertion| {
            b.confidence
                .partial_cmp(&a.confidence)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| b.timestamp.cmp(&a.timestamp))
        };

        must_use.sort_by(sort_by_confidence);
        forbidden.sort_by(sort_by_confidence);
        prefer.sort_by(sort_by_confidence);

        // Compute conflict score across all constraint assertions
        // Single-pass collection: avoid intermediate Vec<&Assertion> then clone
        let constraint_count = must_use.len() + forbidden.len() + prefer.len();
        let conflict = if constraint_count <= 1 {
            0.0
        } else {
            let all_constraints: Vec<Assertion> =
                must_use.iter().chain(forbidden.iter()).chain(prefer.iter()).cloned().collect();
            compute_conflict_score(&all_constraints)
        };

        ConstraintSet {
            must_use,
            forbidden,
            prefer,
            candidates_count: candidates.len(),
            conflict_score: conflict,
        }
    }
}

impl Lens for ConstraintsLens {
    /// Standard Lens resolution: returns highest-priority constraint as winner.
    ///
    /// Priority: must_use > forbidden > prefer
    ///
    /// For the full constraint set, use `resolve_constraints()` instead.
    #[instrument(skip(self, candidates), fields(candidates_count = candidates.len(), lens = "Constraints"))]
    fn resolve(&self, candidates: &[Assertion]) -> Resolution {
        let constraints = self.resolve_constraints(candidates);

        if !constraints.has_constraints() {
            return Resolution::empty();
        }

        // Priority: must_use > forbidden > prefer
        let winner = constraints
            .must_use
            .first()
            .or_else(|| constraints.forbidden.first())
            .or_else(|| constraints.prefer.first())
            .cloned();

        match winner {
            Some(w) => {
                let confidence = w.confidence;
                Resolution::with_winner(
                    w,
                    constraints.total_constraints(),
                    confidence,
                    constraints.conflict_score,
                )
            }
            None => Resolution::empty(),
        }
    }

    fn name(&self) -> &'static str {
        "Constraints"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use stemedb_core::testing::AssertionBuilder;
    use stemedb_core::types::ObjectValue;

    fn create_constraint_assertion(
        subject: &str,
        predicate: &str,
        object: &str,
        confidence: f32,
        timestamp: u64,
    ) -> Assertion {
        AssertionBuilder::new()
            .subject(subject)
            .predicate(predicate)
            .object_text(object)
            .confidence(confidence)
            .timestamp(timestamp)
            .build()
    }

    // ========================================================================
    // resolve_constraints() Tests
    // ========================================================================

    #[test]
    fn test_constraints_categorizes_by_predicate() {
        let lens = ConstraintsLens;

        let must_axios = create_constraint_assertion(
            "project_alpha",
            "must_use:http_client",
            "axios",
            0.95,
            1000,
        );
        let forbidden_requests = create_constraint_assertion(
            "project_alpha",
            "forbidden:http_client",
            "requests",
            0.9,
            1000,
        );
        let prefer_ts = create_constraint_assertion(
            "project_alpha",
            "prefer:language",
            "typescript",
            0.8,
            1000,
        );

        let constraints = lens.resolve_constraints(&[
            must_axios.clone(),
            forbidden_requests.clone(),
            prefer_ts.clone(),
        ]);

        assert_eq!(constraints.must_use.len(), 1);
        assert_eq!(constraints.forbidden.len(), 1);
        assert_eq!(constraints.prefer.len(), 1);

        assert_eq!(constraints.must_use[0].object, ObjectValue::Text("axios".to_string()));
        assert_eq!(constraints.forbidden[0].object, ObjectValue::Text("requests".to_string()));
        assert_eq!(constraints.prefer[0].object, ObjectValue::Text("typescript".to_string()));
    }

    #[test]
    fn test_constraints_empty_categories() {
        let lens = ConstraintsLens;

        // Only prefer, no must_use or forbidden
        let prefer_ts =
            create_constraint_assertion("project", "prefer:language", "typescript", 0.8, 1000);

        let constraints = lens.resolve_constraints(&[prefer_ts]);

        assert!(constraints.must_use.is_empty());
        assert!(constraints.forbidden.is_empty());
        assert_eq!(constraints.prefer.len(), 1);
    }

    #[test]
    fn test_constraints_non_constraint_predicates_ignored() {
        let lens = ConstraintsLens;

        // Regular predicates (not must_use/forbidden/prefer) should be ignored
        let regular = create_constraint_assertion("project", "uses_framework", "react", 0.9, 1000);
        let must_use =
            create_constraint_assertion("project", "must_use:testing", "jest", 0.95, 1000);

        let constraints = lens.resolve_constraints(&[regular, must_use]);

        assert_eq!(constraints.must_use.len(), 1);
        assert!(constraints.forbidden.is_empty());
        assert!(constraints.prefer.is_empty());
        assert_eq!(constraints.candidates_count, 2); // Both candidates considered
    }

    #[test]
    fn test_constraints_sorted_by_confidence() {
        let lens = ConstraintsLens;

        let low = create_constraint_assertion("project", "must_use:db", "sqlite", 0.5, 1000);
        let high = create_constraint_assertion("project", "must_use:db", "postgres", 0.95, 1000);
        let medium = create_constraint_assertion("project", "must_use:db", "mysql", 0.75, 1000);

        let constraints = lens.resolve_constraints(&[low, high, medium]);

        assert_eq!(constraints.must_use.len(), 3);
        // Should be sorted by confidence descending
        assert_eq!(constraints.must_use[0].object, ObjectValue::Text("postgres".to_string()));
        assert_eq!(constraints.must_use[1].object, ObjectValue::Text("mysql".to_string()));
        assert_eq!(constraints.must_use[2].object, ObjectValue::Text("sqlite".to_string()));
    }

    #[test]
    fn test_constraints_empty_candidates() {
        let lens = ConstraintsLens;
        let constraints = lens.resolve_constraints(&[]);

        assert!(!constraints.has_constraints());
        assert_eq!(constraints.total_constraints(), 0);
        assert_eq!(constraints.candidates_count, 0);
    }

    #[test]
    fn test_constraints_has_constraints_true() {
        let lens = ConstraintsLens;
        let must =
            create_constraint_assertion("project", "must_use:formatter", "prettier", 0.9, 1000);

        let constraints = lens.resolve_constraints(&[must]);

        assert!(constraints.has_constraints());
        assert_eq!(constraints.total_constraints(), 1);
    }

    #[test]
    fn test_constraints_all_regular_predicates() {
        let lens = ConstraintsLens;

        // No constraint predicates at all
        let regular1 = create_constraint_assertion("project", "uses", "react", 0.9, 1000);
        let regular2 = create_constraint_assertion("project", "depends_on", "webpack", 0.8, 1000);

        let constraints = lens.resolve_constraints(&[regular1, regular2]);

        assert!(!constraints.has_constraints());
        assert_eq!(constraints.total_constraints(), 0);
        assert_eq!(constraints.candidates_count, 2);
    }

    // ========================================================================
    // Lens trait Tests
    // ========================================================================

    #[test]
    fn test_lens_trait_picks_must_use_winner() {
        let lens = ConstraintsLens;

        let must =
            create_constraint_assertion("project", "must_use:formatter", "prettier", 0.95, 1000);
        let forbidden =
            create_constraint_assertion("project", "forbidden:formatter", "tslint", 0.9, 1000);
        let prefer = create_constraint_assertion("project", "prefer:style", "airbnb", 0.85, 1000);

        let resolution = lens.resolve(&[must.clone(), forbidden, prefer]);

        assert!(resolution.winner.is_some());
        // must_use has priority
        assert_eq!(
            resolution.winner.as_ref().map(|a| &a.predicate),
            Some(&"must_use:formatter".to_string())
        );
    }

    #[test]
    fn test_lens_trait_falls_back_to_forbidden() {
        let lens = ConstraintsLens;

        // No must_use, should pick forbidden
        let forbidden =
            create_constraint_assertion("project", "forbidden:db", "mongodb", 0.9, 1000);
        let prefer = create_constraint_assertion("project", "prefer:db", "postgres", 0.85, 1000);

        let resolution = lens.resolve(&[forbidden.clone(), prefer]);

        assert!(resolution.winner.is_some());
        assert_eq!(
            resolution.winner.as_ref().map(|a| &a.predicate),
            Some(&"forbidden:db".to_string())
        );
    }

    #[test]
    fn test_lens_trait_falls_back_to_prefer() {
        let lens = ConstraintsLens;

        // Only prefer
        let prefer = create_constraint_assertion("project", "prefer:runtime", "bun", 0.8, 1000);

        let resolution = lens.resolve(std::slice::from_ref(&prefer));

        assert!(resolution.winner.is_some());
        assert_eq!(
            resolution.winner.as_ref().map(|a| &a.predicate),
            Some(&"prefer:runtime".to_string())
        );
    }

    #[test]
    fn test_lens_trait_empty_for_no_constraints() {
        let lens = ConstraintsLens;

        // No constraint predicates
        let regular = create_constraint_assertion("project", "uses", "react", 0.9, 1000);

        let resolution = lens.resolve(&[regular]);

        assert!(resolution.winner.is_none());
    }

    #[test]
    fn test_lens_name() {
        let lens = ConstraintsLens;
        assert_eq!(lens.name(), "Constraints");
    }

    #[test]
    fn test_lens_empty_candidates() {
        let lens = ConstraintsLens;
        let resolution = lens.resolve(&[]);

        assert!(resolution.winner.is_none());
        assert_eq!(resolution.candidates_count, 0);
    }

    // ========================================================================
    // Edge Case Tests
    // ========================================================================

    #[test]
    fn test_multiple_must_use_picks_highest_confidence() {
        let lens = ConstraintsLens;

        let low_conf =
            create_constraint_assertion("project", "must_use:bundler", "webpack", 0.5, 1000);
        let high_conf =
            create_constraint_assertion("project", "must_use:bundler", "vite", 0.95, 1000);

        let resolution = lens.resolve(&[low_conf, high_conf]);

        assert!(resolution.winner.is_some());
        assert_eq!(
            resolution.winner.as_ref().map(|a| &a.object),
            Some(&ObjectValue::Text("vite".to_string()))
        );
    }

    #[test]
    fn test_confidence_tiebreaker_uses_timestamp() {
        let lens = ConstraintsLens;

        let older = create_constraint_assertion("project", "must_use:test", "jest", 0.9, 1000);
        let newer = create_constraint_assertion("project", "must_use:test", "vitest", 0.9, 2000);

        let constraints = lens.resolve_constraints(&[older, newer]);

        // Same confidence, newer timestamp wins (sorted first)
        assert_eq!(constraints.must_use[0].object, ObjectValue::Text("vitest".to_string()));
    }

    #[test]
    fn test_predicate_pattern_exact_prefix() {
        let lens = ConstraintsLens;
        // "must_use_something" should NOT match (not "must_use:*")
        let wrong_prefix =
            create_constraint_assertion("project", "must_use_something", "foo", 0.9, 1000);
        let correct =
            create_constraint_assertion("project", "must_use:something", "bar", 0.9, 1000);
        let constraints = lens.resolve_constraints(&[wrong_prefix, correct]);
        assert_eq!(constraints.must_use.len(), 1);
        assert_eq!(constraints.must_use[0].object, ObjectValue::Text("bar".to_string()));
    }
}