storybook/src/query.rs

//! Query interface for filtering and searching entities
//!
//! This module provides convenient methods for querying entities in a storybook
//! project. You can filter by various criteria like traits, age ranges, field
//! values, etc.

use crate::{
    syntax::ast::Value,
    types::*,
};

/// Extension methods for querying characters
pub trait CharacterQuery<'a> {
    /// Filter characters by age range
    fn with_age_range(
        self,
        min: i64,
        max: i64,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a>;

    /// Filter characters by trait value
    fn with_trait(
        self,
        trait_name: &'a str,
        min: f64,
        max: f64,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a>;

    /// Filter characters that have a specific field
    fn with_field(
        self,
        field_name: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a>;

    /// Filter characters by field value
    fn with_field_value(
        self,
        field_name: &'a str,
        value: Value,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a>;
}

impl<'a, I> CharacterQuery<'a> for I
where
    I: Iterator<Item = &'a ResolvedCharacter> + 'a,
{
    fn with_age_range(
        self,
        min: i64,
        max: i64,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a> {
        Box::new(self.filter(move |c| {
            if let Some(Value::Number(age)) = c.fields.get("age") {
                *age >= min && *age <= max
            } else {
                false
            }
        }))
    }

    fn with_trait(
        self,
        trait_name: &'a str,
        min: f64,
        max: f64,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a> {
        Box::new(self.filter(move |c| {
            if let Some(Value::Decimal(value)) = c.fields.get(trait_name) {
                *value >= min && *value <= max
            } else {
                false
            }
        }))
    }

    fn with_field(
        self,
        field_name: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a> {
        Box::new(self.filter(move |c| c.fields.contains_key(field_name)))
    }

    fn with_field_value(
        self,
        field_name: &'a str,
        value: Value,
    ) -> Box<dyn Iterator<Item = &'a ResolvedCharacter> + 'a> {
        Box::new(self.filter(move |c| c.fields.get(field_name) == Some(&value)))
    }
}

/// Extension methods for querying relationships
pub trait RelationshipQuery<'a> {
    /// Filter relationships by bond strength
    fn with_bond_range(
        self,
        min: f64,
        max: f64,
    ) -> Box<dyn Iterator<Item = &'a ResolvedRelationship> + 'a>;

    /// Filter relationships that include a specific participant
    fn with_participant(
        self,
        participant_name: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedRelationship> + 'a>;

    /// Filter relationships that have a specific field
    fn with_field(
        self,
        field_name: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedRelationship> + 'a>;
}

impl<'a, I> RelationshipQuery<'a> for I
where
    I: Iterator<Item = &'a ResolvedRelationship> + 'a,
{
    fn with_bond_range(
        self,
        min: f64,
        max: f64,
    ) -> Box<dyn Iterator<Item = &'a ResolvedRelationship> + 'a> {
        Box::new(self.filter(move |r| {
            if let Some(Value::Decimal(bond)) = r.fields.get("bond") {
                *bond >= min && *bond <= max
            } else {
                false
            }
        }))
    }

    fn with_participant(
        self,
        participant_name: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedRelationship> + 'a> {
        Box::new(self.filter(move |r| {
            r.participants
                .iter()
                .any(|p| p.name.last().is_some_and(|name| name == participant_name))
        }))
    }

    fn with_field(
        self,
        field_name: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedRelationship> + 'a> {
        Box::new(self.filter(move |r| r.fields.contains_key(field_name)))
    }
}

/// Extension methods for querying schedules
pub trait ScheduleQuery<'a> {
    /// Filter schedules that have an activity
    fn with_activity(
        self,
        activity: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedSchedule> + 'a>;
}

impl<'a, I> ScheduleQuery<'a> for I
where
    I: Iterator<Item = &'a ResolvedSchedule> + 'a,
{
    fn with_activity(
        self,
        activity: &'a str,
    ) -> Box<dyn Iterator<Item = &'a ResolvedSchedule> + 'a> {
        Box::new(self.filter(move |s| s.blocks.iter().any(|block| block.activity == activity)))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::syntax::ast::Span;

    fn make_character(name: &str, age: i64, trust: f64) -> ResolvedCharacter {
        let mut fields = HashMap::new();
        fields.insert("age".to_string(), Value::Number(age));
        fields.insert("trust".to_string(), Value::Decimal(trust));

        ResolvedCharacter {
            name: name.to_string(),
            species: None,
            fields,
            prose_blocks: HashMap::new(),
            uses_behaviors: None,
            uses_schedule: None,
            span: Span::new(0, 10),
        }
    }

    #[test]
    fn test_filter_by_age_range() {
        let characters = [
            make_character("Alice", 25, 0.8),
            make_character("Bob", 35, 0.6),
            make_character("Charlie", 45, 0.9),
        ];

        let filtered: Vec<_> = characters.iter().with_age_range(30, 50).collect();

        assert_eq!(filtered.len(), 2);
        assert_eq!(filtered[0].name, "Bob");
        assert_eq!(filtered[1].name, "Charlie");
    }

    #[test]
    fn test_filter_by_trait() {
        let characters = [
            make_character("Alice", 25, 0.8),
            make_character("Bob", 35, 0.6),
            make_character("Charlie", 45, 0.9),
        ];

        let filtered: Vec<_> = characters.iter().with_trait("trust", 0.75, 1.0).collect();

        assert_eq!(filtered.len(), 2);
        assert_eq!(filtered[0].name, "Alice");
        assert_eq!(filtered[1].name, "Charlie");
    }

    #[test]
    fn test_chain_filters() {
        let characters = [
            make_character("Alice", 25, 0.8),
            make_character("Bob", 35, 0.6),
            make_character("Charlie", 45, 0.9),
            make_character("David", 40, 0.85),
        ];

        // Find characters aged 30-50 with trust > 0.8
        let filtered: Vec<_> = characters
            .iter()
            .with_age_range(30, 50)
            .with_trait("trust", 0.8, 1.0)
            .collect();

        assert_eq!(filtered.len(), 2);
        assert_eq!(filtered[0].name, "Charlie");
        assert_eq!(filtered[1].name, "David");
    }

    #[test]
    fn test_filter_with_field() {
        let mut char1 = make_character("Alice", 25, 0.8);
        char1
            .fields
            .insert("job".to_string(), Value::Text("baker".to_string()));

        let char2 = make_character("Bob", 35, 0.6);

        let characters = [char1, char2];

        let filtered: Vec<_> = characters.iter().with_field("job").collect();

        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].name, "Alice");
    }

    #[test]
    fn test_relationship_with_bond_range() {
        let mut fields1 = HashMap::new();
        fields1.insert("bond".to_string(), Value::Decimal(0.9));

        let mut fields2 = HashMap::new();
        fields2.insert("bond".to_string(), Value::Decimal(0.5));

        let relationships = [
            ResolvedRelationship {
                name: "Strong".to_string(),
                participants: vec![],
                fields: fields1,
                span: Span::new(0, 10),
            },
            ResolvedRelationship {
                name: "Weak".to_string(),
                participants: vec![],
                fields: fields2,
                span: Span::new(0, 10),
            },
        ];

        let filtered: Vec<_> = relationships.iter().with_bond_range(0.8, 1.0).collect();

        assert_eq!(filtered.len(), 1);
        assert_eq!(filtered[0].name, "Strong");
    }
}