//! Knowledge graph backed by SQLite.
//!
//! Stores entities and relations with support for graph pattern queries.
//! Supports recursive traversal with configurable depth limits.

use chrono::Utc;
use openfang_types::error::{OpenFangError, OpenFangResult};
use openfang_types::memory::{
    Entity, EntityType, GraphMatch, GraphPattern, Relation, RelationType,
};
use rusqlite::Connection;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use uuid::Uuid;

/// Maximum allowed depth for recursive graph queries (safety limit).
const MAX_ALLOWED_DEPTH: u32 = 10;

/// Default result limit for graph queries.
const DEFAULT_RESULT_LIMIT: u32 = 1000;

/// Knowledge graph store backed by SQLite.
#[derive(Clone)]
pub struct KnowledgeStore {
    conn: Arc<Mutex<Connection>>,
}

impl KnowledgeStore {
    /// Create a new knowledge store wrapping the given connection.
    pub fn new(conn: Arc<Mutex<Connection>>) -> Self {
        Self { conn }
    }

    /// Add an entity to the knowledge graph.
    pub fn add_entity(&self, entity: Entity) -> OpenFangResult<String> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| OpenFangError::Internal(e.to_string()))?;
        let id = if entity.id.is_empty() {
            Uuid::new_v4().to_string()
        } else {
            entity.id.clone()
        };
        let entity_type_str = serde_json::to_string(&entity.entity_type)
            .map_err(|e| OpenFangError::Serialization(e.to_string()))?;
        let props_str = serde_json::to_string(&entity.properties)
            .map_err(|e| OpenFangError::Serialization(e.to_string()))?;
        let now = Utc::now().to_rfc3339();
        conn.execute(
            "INSERT INTO entities (id, entity_type, name, properties, created_at, updated_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?5)
             ON CONFLICT(id) DO UPDATE SET name = ?3, properties = ?4, updated_at = ?5",
            rusqlite::params![id, entity_type_str, entity.name, props_str, now],
        )
        .map_err(|e| OpenFangError::Memory(e.to_string()))?;
        Ok(id)
    }

    /// Add a relation between two entities.
    pub fn add_relation(&self, relation: Relation) -> OpenFangResult<String> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| OpenFangError::Internal(e.to_string()))?;
        let id = Uuid::new_v4().to_string();
        let rel_type_str = serde_json::to_string(&relation.relation)
            .map_err(|e| OpenFangError::Serialization(e.to_string()))?;
        let props_str = serde_json::to_string(&relation.properties)
            .map_err(|e| OpenFangError::Serialization(e.to_string()))?;
        let now = Utc::now().to_rfc3339();
        conn.execute(
            "INSERT INTO relations (id, source_entity, relation_type, target_entity, properties, confidence, created_at)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            rusqlite::params![
                id,
                relation.source,
                rel_type_str,
                relation.target,
                props_str,
                relation.confidence as f64,
                now,
            ],
        )
        .map_err(|e| OpenFangError::Memory(e.to_string()))?;
        Ok(id)
    }

    /// Query the knowledge graph with a pattern.
    ///
    /// Supports recursive traversal via `max_depth` parameter:
    /// - `max_depth = 1` (default): Single-hop query (source -> target)
    /// - `max_depth > 1`: Recursive traversal following relation chains
    /// - Maximum allowed depth: 10 (safety limit)
    pub fn query_graph(&self, pattern: GraphPattern) -> OpenFangResult<Vec<GraphMatch>> {
        // Normalize depth to safe range
        let max_depth = pattern.max_depth.min(MAX_ALLOWED_DEPTH).max(1);

        if max_depth == 1 {
            // Single-hop query (original behavior)
            self.query_single_hop(&pattern)
        } else {
            // Recursive traversal using iterative approach
            self.query_recursive(&pattern, max_depth)
        }
    }

    /// Single-hop query (depth = 1).
    fn query_single_hop(&self, pattern: &GraphPattern) -> OpenFangResult<Vec<GraphMatch>> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| OpenFangError::Internal(e.to_string()))?;

        let mut sql = String::from(
            "SELECT
                s.id, s.entity_type, s.name, s.properties, s.created_at, s.updated_at,
                r.id, r.source_entity, r.relation_type, r.target_entity, r.properties, r.confidence, r.created_at,
                t.id, t.entity_type, t.name, t.properties, t.created_at, t.updated_at
             FROM relations r
             JOIN entities s ON r.source_entity = s.id
             JOIN entities t ON r.target_entity = t.id
             WHERE 1=1",
        );
        let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
        let mut idx = 1;

        if let Some(ref source) = pattern.source {
            sql.push_str(&format!(" AND (s.id = ?{idx} OR s.name = ?{idx})"));
            params.push(Box::new(source.clone()));
            idx += 1;
        }
        if let Some(ref relation) = pattern.relation {
            let rel_str = serde_json::to_string(relation)
                .map_err(|e| OpenFangError::Serialization(e.to_string()))?;
            sql.push_str(&format!(" AND r.relation_type = ?{idx}"));
            params.push(Box::new(rel_str));
            idx += 1;
        }
        if let Some(ref target) = pattern.target {
            sql.push_str(&format!(" AND (t.id = ?{idx} OR t.name = ?{idx})"));
            params.push(Box::new(target.clone()));
        }

        sql.push_str(" LIMIT 100");

        let mut stmt = conn
            .prepare(&sql)
            .map_err(|e| OpenFangError::Memory(e.to_string()))?;
        let param_refs: Vec<&dyn rusqlite::types::ToSql> =
            params.iter().map(|p| p.as_ref()).collect();

        let rows = stmt
            .query_map(param_refs.as_slice(), |row| {
                Ok(RawGraphRow {
                    s_id: row.get(0)?,
                    s_type: row.get(1)?,
                    s_name: row.get(2)?,
                    s_props: row.get(3)?,
                    s_created: row.get(4)?,
                    s_updated: row.get(5)?,
                    r_id: row.get(6)?,
                    r_source: row.get(7)?,
                    r_type: row.get(8)?,
                    r_target: row.get(9)?,
                    r_props: row.get(10)?,
                    r_confidence: row.get(11)?,
                    r_created: row.get(12)?,
                    t_id: row.get(13)?,
                    t_type: row.get(14)?,
                    t_name: row.get(15)?,
                    t_props: row.get(16)?,
                    t_created: row.get(17)?,
                    t_updated: row.get(18)?,
                })
            })
            .map_err(|e| OpenFangError::Memory(e.to_string()))?;

        let mut matches = Vec::new();
        for row_result in rows {
            let r = row_result.map_err(|e| OpenFangError::Memory(e.to_string()))?;
            matches.push(GraphMatch {
                source: parse_entity(
                    &r.s_id,
                    &r.s_type,
                    &r.s_name,
                    &r.s_props,
                    &r.s_created,
                    &r.s_updated,
                ),
                relation: parse_relation(
                    &r.r_source,
                    &r.r_type,
                    &r.r_target,
                    &r.r_props,
                    r.r_confidence,
                    &r.r_created,
                ),
                target: parse_entity(
                    &r.t_id,
                    &r.t_type,
                    &r.t_name,
                    &r.t_props,
                    &r.t_created,
                    &r.t_updated,
                ),
            });
        }
        Ok(matches)
    }

    /// Recursive graph traversal using iterative BFS approach.
    ///
    /// This method iteratively traverses the graph up to the specified depth,
    /// tracking visited relations to avoid cycles.
    fn query_recursive(&self, pattern: &GraphPattern, max_depth: u32) -> OpenFangResult<Vec<GraphMatch>> {
        let conn = self
            .conn
            .lock()
            .map_err(|e| OpenFangError::Internal(e.to_string()))?;

        let mut all_matches: Vec<GraphMatch> = Vec::new();
        let mut visited_relations: HashSet<String> = HashSet::new();

        // Start with initial sources
        let mut current_sources: HashSet<String> = if let Some(ref source) = pattern.source {
            // Resolve source to entity ID(s)
            self.resolve_entity_ids(&conn, source)?
        } else {
            // No source filter - start from all entities
            self.get_all_entity_ids(&conn)?
        };

        // Iteratively traverse each depth level
        for _depth in 1..=max_depth {
            if current_sources.is_empty() {
                break;
            }

            // Query relations from current sources
            let matches = self.query_from_sources(&conn, &current_sources, pattern)?;

            // Collect new targets for next iteration
            let mut next_sources: HashSet<String> = HashSet::new();

            for m in matches {
                // Create unique key for this relation to detect cycles
                let rel_key = format!("{}->{}", m.relation.source, m.relation.target);

                if visited_relations.insert(rel_key) {
                    // New relation (not visited)
                    next_sources.insert(m.target.id.clone());
                    all_matches.push(m);
                }
            }

            // Move to next level
            current_sources = next_sources;

            // Safety: limit total results
            if all_matches.len() >= DEFAULT_RESULT_LIMIT as usize {
                all_matches.truncate(DEFAULT_RESULT_LIMIT as usize);
                break;
            }
        }

        Ok(all_matches)
    }

    /// Resolve an entity identifier (ID or name) to entity IDs.
    fn resolve_entity_ids(&self, conn: &Connection, identifier: &str) -> OpenFangResult<HashSet<String>> {
        let mut stmt = conn
            .prepare("SELECT id FROM entities WHERE id = ?1 OR name = ?1")
            .map_err(|e| OpenFangError::Memory(e.to_string()))?;

        let rows: Result<Vec<String>, _> = stmt
            .query_map([identifier], |row| row.get(0))
            .map_err(|e| OpenFangError::Memory(e.to_string()))?
            .collect();

        Ok(rows.map_err(|e| OpenFangError::Memory(e.to_string()))?.into_iter().collect())
    }

    /// Get all entity IDs in the graph.
    fn get_all_entity_ids(&self, conn: &Connection) -> OpenFangResult<HashSet<String>> {
        let mut stmt = conn
            .prepare("SELECT id FROM entities")
            .map_err(|e| OpenFangError::Memory(e.to_string()))?;

        let rows: Result<Vec<String>, _> = stmt
            .query_map([], |row| row.get(0))
            .map_err(|e| OpenFangError::Memory(e.to_string()))?
            .collect();

        Ok(rows.map_err(|e| OpenFangError::Memory(e.to_string()))?.into_iter().collect())
    }

    /// Query relations from a set of source entities.
    fn query_from_sources(
        &self,
        conn: &Connection,
        sources: &HashSet<String>,
        pattern: &GraphPattern,
    ) -> OpenFangResult<Vec<GraphMatch>> {
        if sources.is_empty() {
            return Ok(Vec::new());
        }

        // Build query with IN clause for sources
        let placeholders: Vec<String> = (0..sources.len()).map(|i| format!("?{}", i + 1)).collect();
        let placeholders_str = placeholders.join(", ");

        let mut sql = format!(
            "SELECT
                s.id, s.entity_type, s.name, s.properties, s.created_at, s.updated_at,
                r.id, r.source_entity, r.relation_type, r.target_entity, r.properties, r.confidence, r.created_at,
                t.id, t.entity_type, t.name, t.properties, t.created_at, t.updated_at
             FROM relations r
             JOIN entities s ON r.source_entity = s.id
             JOIN entities t ON r.target_entity = t.id
             WHERE s.id IN ({})",
            placeholders_str
        );

        let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = sources
            .iter()
            .map(|s| Box::new(s.clone()) as Box<dyn rusqlite::types::ToSql>)
            .collect();
        let mut idx = sources.len() + 1;

        // Add relation type filter if specified
        if let Some(ref relation) = pattern.relation {
            let rel_str = serde_json::to_string(relation)
                .map_err(|e| OpenFangError::Serialization(e.to_string()))?;
            sql.push_str(&format!(" AND r.relation_type = ?{idx}"));
            params.push(Box::new(rel_str));
            idx += 1;
        }

        // Add target filter if specified
        if let Some(ref target) = pattern.target {
            sql.push_str(&format!(" AND (t.id = ?{idx} OR t.name = ?{idx})"));
            params.push(Box::new(target.clone()));
        }

        sql.push_str(&format!(" LIMIT {}", DEFAULT_RESULT_LIMIT));

        let mut stmt = conn
            .prepare(&sql)
            .map_err(|e| OpenFangError::Memory(e.to_string()))?;
        let param_refs: Vec<&dyn rusqlite::types::ToSql> =
            params.iter().map(|p| p.as_ref()).collect();

        let rows = stmt
            .query_map(param_refs.as_slice(), |row| {
                Ok(RawGraphRow {
                    s_id: row.get(0)?,
                    s_type: row.get(1)?,
                    s_name: row.get(2)?,
                    s_props: row.get(3)?,
                    s_created: row.get(4)?,
                    s_updated: row.get(5)?,
                    r_id: row.get(6)?,
                    r_source: row.get(7)?,
                    r_type: row.get(8)?,
                    r_target: row.get(9)?,
                    r_props: row.get(10)?,
                    r_confidence: row.get(11)?,
                    r_created: row.get(12)?,
                    t_id: row.get(13)?,
                    t_type: row.get(14)?,
                    t_name: row.get(15)?,
                    t_props: row.get(16)?,
                    t_created: row.get(17)?,
                    t_updated: row.get(18)?,
                })
            })
            .map_err(|e| OpenFangError::Memory(e.to_string()))?;

        let mut matches = Vec::new();
        for row_result in rows {
            let r = row_result.map_err(|e| OpenFangError::Memory(e.to_string()))?;
            matches.push(GraphMatch {
                source: parse_entity(
                    &r.s_id,
                    &r.s_type,
                    &r.s_name,
                    &r.s_props,
                    &r.s_created,
                    &r.s_updated,
                ),
                relation: parse_relation(
                    &r.r_source,
                    &r.r_type,
                    &r.r_target,
                    &r.r_props,
                    r.r_confidence,
                    &r.r_created,
                ),
                target: parse_entity(
                    &r.t_id,
                    &r.t_type,
                    &r.t_name,
                    &r.t_props,
                    &r.t_created,
                    &r.t_updated,
                ),
            });
        }
        Ok(matches)
    }
}

/// Raw row from a graph query.
struct RawGraphRow {
    s_id: String,
    s_type: String,
    s_name: String,
    s_props: String,
    s_created: String,
    s_updated: String,
    r_id: String,
    r_source: String,
    r_type: String,
    r_target: String,
    r_props: String,
    r_confidence: f64,
    r_created: String,
    t_id: String,
    t_type: String,
    t_name: String,
    t_props: String,
    t_created: String,
    t_updated: String,
}

// Suppress the unused field warning — r_id is part of the schema
impl RawGraphRow {
    #[allow(dead_code)]
    fn relation_id(&self) -> &str {
        &self.r_id
    }
}

fn parse_entity(
    id: &str,
    etype: &str,
    name: &str,
    props: &str,
    created: &str,
    updated: &str,
) -> Entity {
    let entity_type: EntityType =
        serde_json::from_str(etype).unwrap_or(EntityType::Custom("unknown".to_string()));
    let properties: HashMap<String, serde_json::Value> =
        serde_json::from_str(props).unwrap_or_default();
    let created_at = chrono::DateTime::parse_from_rfc3339(created)
        .map(|dt| dt.with_timezone(&Utc))
        .unwrap_or_else(|_| Utc::now());
    let updated_at = chrono::DateTime::parse_from_rfc3339(updated)
        .map(|dt| dt.with_timezone(&Utc))
        .unwrap_or_else(|_| Utc::now());
    Entity {
        id: id.to_string(),
        entity_type,
        name: name.to_string(),
        properties,
        created_at,
        updated_at,
    }
}

fn parse_relation(
    source: &str,
    rtype: &str,
    target: &str,
    props: &str,
    confidence: f64,
    created: &str,
) -> Relation {
    let relation: RelationType = serde_json::from_str(rtype).unwrap_or(RelationType::RelatedTo);
    let properties: HashMap<String, serde_json::Value> =
        serde_json::from_str(props).unwrap_or_default();
    let created_at = chrono::DateTime::parse_from_rfc3339(created)
        .map(|dt| dt.with_timezone(&Utc))
        .unwrap_or_else(|_| Utc::now());
    Relation {
        source: source.to_string(),
        relation,
        target: target.to_string(),
        properties,
        confidence: confidence as f32,
        created_at,
    }
}

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

    fn setup() -> KnowledgeStore {
        let conn = Connection::open_in_memory().unwrap();
        run_migrations(&conn).unwrap();
        KnowledgeStore::new(Arc::new(Mutex::new(conn)))
    }

    #[test]
    fn test_add_and_query_entity() {
        let store = setup();
        let id = store
            .add_entity(Entity {
                id: String::new(),
                entity_type: EntityType::Person,
                name: "Alice".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();
        assert!(!id.is_empty());
    }

    #[test]
    fn test_add_relation_and_query() {
        let store = setup();
        let alice_id = store
            .add_entity(Entity {
                id: "alice".to_string(),
                entity_type: EntityType::Person,
                name: "Alice".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();
        let company_id = store
            .add_entity(Entity {
                id: "acme".to_string(),
                entity_type: EntityType::Organization,
                name: "Acme Corp".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();
        store
            .add_relation(Relation {
                source: alice_id.clone(),
                relation: RelationType::WorksAt,
                target: company_id,
                properties: HashMap::new(),
                confidence: 0.95,
                created_at: Utc::now(),
            })
            .unwrap();

        let matches = store
            .query_graph(GraphPattern {
                source: Some(alice_id),
                relation: Some(RelationType::WorksAt),
                target: None,
                max_depth: 1,
            })
            .unwrap();
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].target.name, "Acme Corp");
    }

    #[test]
    fn test_recursive_traversal_depth_2() {
        // Create a chain: Alice -> WorksAt -> Acme -> LocatedIn -> NYC
        let store = setup();

        let alice_id = store
            .add_entity(Entity {
                id: "alice".to_string(),
                entity_type: EntityType::Person,
                name: "Alice".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        let acme_id = store
            .add_entity(Entity {
                id: "acme".to_string(),
                entity_type: EntityType::Organization,
                name: "Acme Corp".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        let nyc_id = store
            .add_entity(Entity {
                id: "nyc".to_string(),
                entity_type: EntityType::Location,
                name: "New York City".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        // Alice works at Acme
        store
            .add_relation(Relation {
                source: alice_id.clone(),
                relation: RelationType::WorksAt,
                target: acme_id.clone(),
                properties: HashMap::new(),
                confidence: 0.95,
                created_at: Utc::now(),
            })
            .unwrap();

        // Acme is located in NYC (using RelatedTo as LocatedIn doesn't exist)
        store
            .add_relation(Relation {
                source: acme_id.clone(),
                relation: RelationType::RelatedTo,
                target: nyc_id.clone(),
                properties: HashMap::new(),
                confidence: 1.0,
                created_at: Utc::now(),
            })
            .unwrap();

        // Query with depth=1 should only return Alice -> Acme
        let matches_depth_1 = store
            .query_graph(GraphPattern {
                source: Some(alice_id.clone()),
                relation: None,
                target: None,
                max_depth: 1,
            })
            .unwrap();
        assert_eq!(matches_depth_1.len(), 1);
        assert_eq!(matches_depth_1[0].target.name, "Acme Corp");

        // Query with depth=2 should return both relations
        let matches_depth_2 = store
            .query_graph(GraphPattern {
                source: Some(alice_id),
                relation: None,
                target: None,
                max_depth: 2,
            })
            .unwrap();
        assert_eq!(matches_depth_2.len(), 2);

        // Verify we get both hops
        let target_names: Vec<&str> = matches_depth_2.iter().map(|m| m.target.name.as_str()).collect();
        assert!(target_names.contains(&"Acme Corp"));
        assert!(target_names.contains(&"New York City"));
    }

    #[test]
    fn test_recursive_traversal_with_cycle() {
        // Create a cycle: A -> B -> C -> A
        let store = setup();

        let a_id = store
            .add_entity(Entity {
                id: "a".to_string(),
                entity_type: EntityType::Person,
                name: "Entity A".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        let b_id = store
            .add_entity(Entity {
                id: "b".to_string(),
                entity_type: EntityType::Person,
                name: "Entity B".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        let c_id = store
            .add_entity(Entity {
                id: "c".to_string(),
                entity_type: EntityType::Person,
                name: "Entity C".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        // A -> B
        store
            .add_relation(Relation {
                source: a_id.clone(),
                relation: RelationType::RelatedTo,
                target: b_id.clone(),
                properties: HashMap::new(),
                confidence: 1.0,
                created_at: Utc::now(),
            })
            .unwrap();

        // B -> C
        store
            .add_relation(Relation {
                source: b_id.clone(),
                relation: RelationType::RelatedTo,
                target: c_id.clone(),
                properties: HashMap::new(),
                confidence: 1.0,
                created_at: Utc::now(),
            })
            .unwrap();

        // C -> A (creates cycle)
        store
            .add_relation(Relation {
                source: c_id.clone(),
                relation: RelationType::RelatedTo,
                target: a_id.clone(),
                properties: HashMap::new(),
                confidence: 1.0,
                created_at: Utc::now(),
            })
            .unwrap();

        // Query should handle cycle without infinite loop
        let matches = store
            .query_graph(GraphPattern {
                source: Some(a_id),
                relation: None,
                target: None,
                max_depth: 5,
            })
            .unwrap();

        // Should only return 3 unique relations (A->B, B->C, C->A)
        // Cycle detection prevents revisiting A->B
        assert_eq!(matches.len(), 3);
    }

    #[test]
    fn test_recursive_traversal_with_relation_filter() {
        // Create chain with different relation types
        let store = setup();

        let alice_id = store
            .add_entity(Entity {
                id: "alice".to_string(),
                entity_type: EntityType::Person,
                name: "Alice".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        let bob_id = store
            .add_entity(Entity {
                id: "bob".to_string(),
                entity_type: EntityType::Person,
                name: "Bob".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        let carol_id = store
            .add_entity(Entity {
                id: "carol".to_string(),
                entity_type: EntityType::Person,
                name: "Carol".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        // Alice knows Bob
        store
            .add_relation(Relation {
                source: alice_id.clone(),
                relation: RelationType::Knows,
                target: bob_id.clone(),
                properties: HashMap::new(),
                confidence: 0.9,
                created_at: Utc::now(),
            })
            .unwrap();

        // Bob works with Carol (different relation type)
        store
            .add_relation(Relation {
                source: bob_id.clone(),
                relation: RelationType::WorksAt,
                target: carol_id.clone(),
                properties: HashMap::new(),
                confidence: 0.9,
                created_at: Utc::now(),
            })
            .unwrap();

        // Query with relation filter should only return matching relations
        let matches = store
            .query_graph(GraphPattern {
                source: Some(alice_id),
                relation: Some(RelationType::Knows),
                target: None,
                max_depth: 3,
            })
            .unwrap();

        // Should only return Alice -> Bob (Knows relation)
        // Bob -> Carol is WorksAt, not Knows
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].target.name, "Bob");
    }

    #[test]
    fn test_max_depth_safety_limit() {
        // Test that max_depth is capped at MAX_ALLOWED_DEPTH (10)
        let store = setup();

        let alice_id = store
            .add_entity(Entity {
                id: "alice".to_string(),
                entity_type: EntityType::Person,
                name: "Alice".to_string(),
                properties: HashMap::new(),
                created_at: Utc::now(),
                updated_at: Utc::now(),
            })
            .unwrap();

        // Request depth of 100 (should be capped to 10)
        let _matches = store
            .query_graph(GraphPattern {
                source: Some(alice_id),
                relation: None,
                target: None,
                max_depth: 100, // Should be capped to 10
            })
            .unwrap();

        // If we reach here without hanging, the safety limit worked
        // (A real test would verify the actual depth, but this is a basic sanity check)
    }
}
