use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::*;
use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashSet;
use std::sync::Arc;

#[derive(serde::Serialize, Default)]
struct BorrowedGraph<'a> {
    entities: std::collections::HashMap<&'a String, &'a crate::models::Entity>,
    relations: Vec<&'a crate::models::Relation>,
}

pub struct GraphQueryBuilder<'a> {
    graph: &'a crate::models::KnowledgeGraph,
    max_depth: usize,
    relation_filter: Option<&'a str>,
}

impl<'a> GraphQueryBuilder<'a> {
    pub fn new(graph: &'a crate::models::KnowledgeGraph) -> Self {
        Self {
            graph,
            max_depth: 5,
            relation_filter: None,
        }
    }

    pub fn max_depth(mut self, depth: usize) -> Self {
        self.max_depth = depth;
        self
    }

    pub fn relation_filter(mut self, filter: Option<&'a str>) -> Self {
        self.relation_filter = filter;
        self
    }

    pub fn find_shortest_path(&self, start: &str, end: &str) -> Option<Vec<String>> {
        let adj = self.graph.build_adjacency_map(self.relation_filter);

        let mut queue = std::collections::VecDeque::new();
        let mut visited = std::collections::HashSet::new();
        let mut parents = std::collections::HashMap::new();

        queue.push_back(start);
        visited.insert(start);

        let mut found = false;
        let mut current_depth = 0;
        let mut nodes_at_current_depth = 1;
        let mut nodes_at_next_depth = 0;

        while let Some(current) = queue.pop_front() {
            if current == end {
                found = true;
                break;
            }
            // Visited node upper-bound cap to guarantee deterministic BFS bounds on dense graphs
            if visited.len() > 10_000 {
                break;
            }
            nodes_at_current_depth -= 1;
            if current_depth < self.max_depth
                && let Some(neighbors) = adj.get(current)
            {
                for &(neighbor, rel_type, is_inverse) in neighbors {
                    if !visited.contains(neighbor) {
                        visited.insert(neighbor);
                        parents.insert(neighbor, (current, rel_type, is_inverse));
                        queue.push_back(neighbor);
                        nodes_at_next_depth += 1;
                    }
                }
            }
            if nodes_at_current_depth == 0 {
                current_depth += 1;
                nodes_at_current_depth = nodes_at_next_depth;
                nodes_at_next_depth = 0;
            }
        }

        if found {
            let mut path = Vec::new();
            let mut curr = end;
            while curr != start {
                if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
                    if *is_inverse {
                        path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
                    } else {
                        path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
                    }
                    curr = parent;
                } else {
                    break;
                }
            }
            path.reverse();
            Some(path)
        } else {
            None
        }
    }
}

pub struct QueryGraphPathHandler;

#[async_trait]
impl McpTool for QueryGraphPathHandler {
    fn name(&self) -> &'static str {
        "query_graph_path"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<QueryGraphPathTool>(
            "query_graph_path",
            "Find the shortest relationship path between two entities in the knowledge graph within a maximum depth.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: crate::tools::QueryGraphPathTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;

        tokio::task::spawn_blocking(move || {
            state.read_graph(|graph| {
                let max_depth = req.max_depth.unwrap_or(5);
                let builder = GraphQueryBuilder::new(graph).max_depth(max_depth as usize);
                if let Some(path) = builder.find_shortest_path(&req.start_node, &req.end_node) {
                    Ok(format!("Path found:\n{}", path.join("\n")))
                } else {
                    Ok(format!(
                        "No path found between {} and {} within depth {}",
                        req.start_node, req.end_node, max_depth
                    ))
                }
            })
        })
        .await
        .map_err(|e| {
            crate::error::AppError::Internal(format!("Graph traversal task failed: {}", e))
        })?
    }
}

pub struct CreateEntitiesHandler;

#[async_trait]
impl McpTool for CreateEntitiesHandler {
    fn name(&self) -> &'static str {
        "create_entities"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<CreateEntitiesTool>(
            "create_entities",
            "Create new entities in the knowledge graph with normalized entity types.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut inserted = Vec::new();
        state.modify_graph(|g| {
            for entity_input in req.entities {
                if !entity_input.name.is_empty() {
                    let entity = crate::models::Entity {
                        name: entity_input.name.clone(),
                        entity_type: crate::models::normalize_entity_type(&entity_input.entity_type),
                        observations: entity_input.observations,
                        namespace: entity_input.namespace.unwrap_or_else(crate::models::default_namespace),
                        ..Default::default()
                    };
                    inserted.push(entity.clone());
                    g.entities.insert(entity.name.clone(), entity);
                }
            }
        });
        let names: Vec<String> = inserted
            .iter()
            .map(|e| format!("{} ({})", e.name, e.entity_type))
            .collect();
        if !inserted.is_empty() {
            let idx = state.get_search_index().await;
            let _ = idx.index_entities_batch(&inserted).await;
        }
        Ok(format!(
            "Successfully created {} entity/entities: {}",
            names.len(),
            names.join(", ")
        ))
    }
}

pub struct UpdateEntitiesHandler;

#[async_trait]
impl McpTool for UpdateEntitiesHandler {
    fn name(&self) -> &'static str {
        "update_entities"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<UpdateEntitiesTool>(
            "update_entities",
            "Update existing entities in the knowledge graph (e.g., to rename an entity or change its type).",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: UpdateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        
        let mut updated_entities = Vec::new();
        let mut deleted_names = Vec::new();
        let mut not_found = Vec::new();
        let mut conflict_names = Vec::new();

        state.modify_graph(|g| {
            for update in req.updates {
                if !g.entities.contains_key(&update.name) {
                    not_found.push(update.name.clone());
                    continue;
                }

                if let Some(new_name) = &update.new_name {
                    if update.name != *new_name && g.entities.contains_key(new_name) {
                        conflict_names.push(new_name.clone());
                        continue;
                    }
                }

                if let Some(mut entity) = g.entities.remove(&update.name) {
                    let mut renamed = false;
                    let old_name = entity.name.clone();

                    if let Some(new_name) = update.new_name {
                        if old_name != new_name {
                            entity.name = new_name.clone();
                            renamed = true;
                            deleted_names.push(old_name.clone());
                        }
                    }

                    if let Some(new_type) = update.new_entity_type {
                        entity.entity_type = crate::models::normalize_entity_type(&new_type);
                    }

                    let current_name = entity.name.clone();
                    
                    if renamed {
                        // update relations
                        let mut seen = std::collections::HashSet::new();
                        g.relations.retain_mut(|r| {
                            if r.from == old_name {
                                r.from = current_name.clone();
                            }
                            if r.to == old_name {
                                r.to = current_name.clone();
                            }
                            // Keep unique relations
                            seen.insert(r.clone())
                        });
                    }

                    updated_entities.push(entity.clone());
                    g.entities.insert(current_name, entity);
                }
            }
        });

        if !not_found.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Entities not found: {}", not_found.join(", ")
            )));
        }

        if !conflict_names.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Cannot rename to existing entity names: {}", conflict_names.join(", ")
            )));
        }

        let idx = state.get_search_index().await;
        for old_name in deleted_names {
            drop(idx.delete_document(&old_name));
        }

        if !updated_entities.is_empty() {
            let _ = idx.index_entities_batch(&updated_entities).await;
        }

        let names: Vec<String> = updated_entities.iter().map(|e| e.name.clone()).collect();
        Ok(format!("Successfully updated {} entity/entities: {}", names.len(), names.join(", ")))
    }
}

pub struct CreateRelationsHandler;

#[async_trait]
impl McpTool for CreateRelationsHandler {
    fn name(&self) -> &'static str {
        "create_relations"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<CreateRelationsTool>(
            "create_relations",
            "Create directed relationships between entities in the knowledge graph. Requires 'from', 'to', and 'relation_type'.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: CreateRelationsTool = match serde_json::from_value(args.clone()) {
            Ok(r) => r,
            Err(e) => {
                let err_msg = e.to_string();
                if err_msg.contains("missing field `from`")
                    || err_msg.contains("missing field `to`")
                    || err_msg.contains("missing field `relation_type`")
                {
                    return Err(crate::error::AppError::Internal(format!(
                        "Schema error: {}. Note that the relation schema strictly uses 'from', 'to', and 'relation_type' (not 'source', 'target', or 'relationType'). Please correct your tool call arguments.",
                        err_msg
                    )));
                }
                return Err(err_msg.into());
            }
        };
        let mut missing_nodes = std::collections::HashSet::new();
        state.read_graph(|g| {
            for relation in &req.relations {
                if !relation.from.is_empty() && !relation.to.is_empty() {
                    let from_exists = g.entities.contains_key(&relation.from);
                    let to_exists = g.entities.contains_key(&relation.to);
                    if !from_exists {
                        missing_nodes.insert(relation.from.clone());
                    }
                    if !to_exists {
                        missing_nodes.insert(relation.to.clone());
                    }
                }
            }
        });

        let mut auto_created = Vec::new();
        let mut added_relations = Vec::new();
        state.modify_graph(|g| {
            for node_name in missing_nodes {
                if !g.entities.contains_key(&node_name) {
                    g.entities.insert(
                        node_name.clone(),
                        crate::models::Entity {
                            name: node_name.clone(),
                            entity_type: "Entity".to_string(),
                            observations: vec![
                                "Auto-created stub entity for relation endpoint".to_string(),
                            ],
                            namespace: crate::models::default_namespace(),
                            git_branch: None,
                            ..Default::default()
                        },
                    );
                    auto_created.push(node_name);
                }
            }
            for relation_input in req.relations {
                if !relation_input.from.is_empty() && !relation_input.to.is_empty() {
                    let relation = crate::models::Relation {
                        from: relation_input.from,
                        to: relation_input.to,
                        relation_type: crate::models::normalize_relation_type(&relation_input.relation_type),
                        namespace: relation_input.namespace.unwrap_or_else(crate::models::default_namespace),
                        ..Default::default()
                    };
                    added_relations.push(format!(
                        "{} -[{}]-> {}",
                        relation.from, relation.relation_type, relation.to
                    ));
                    g.relations.push(relation);
                }
            }
        });

        let mut msg = format!(
            "Successfully created {} relation(s):\n{}",
            added_relations.len(),
            added_relations.join("\n")
        );
        if !auto_created.is_empty() {
            msg.push_str(&format!(
                "\nNote: Auto-created {} missing stub entity/entities: {}",
                auto_created.len(),
                auto_created.join(", ")
            ));
        }
        Ok(msg)
    }
}

pub struct AddObservationsHandler;

#[async_trait]
impl McpTool for AddObservationsHandler {
    fn name(&self) -> &'static str {
        "add_observations"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<AddObservationsTool>(
            "add_observations",
            "Add new observations and factual statements to existing entities in the knowledge graph.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut missing_entities = Vec::new();
        state.read_graph(|g| {
            for o in &req.observations {
                if !g.entities.contains_key(&o.entity_name) {
                    missing_entities.push(o.entity_name.clone());
                }
            }
        });

        if !missing_entities.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Observations dropped for missing entities: {}",
                missing_entities.join(", ")
            )));
        }

        state.modify_graph(|g| {
            for o in req.observations {
                if let Some(e) = g.entities.get_mut(&o.entity_name) {
                    e.observations.extend(o.contents);
                }
            }
        });
        Ok("Observations added".to_string())
    }
}

pub struct DeleteEntitiesHandler;

#[async_trait]
impl McpTool for DeleteEntitiesHandler {
    fn name(&self) -> &'static str {
        "delete_entities"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<DeleteEntitiesTool>(
            "delete_entities",
            "Delete specified entities and all their associated relations and observations from the knowledge graph.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect();
        let mut missing = Vec::new();
        state.read_graph(|g| {
            for name in &to_delete {
                if !g.entities.contains_key(name) {
                    missing.push(name.clone());
                }
            }
        });

        if !missing.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
                missing.join(", ")
            )));
        }

        state.modify_graph(|master| {
            for name in &to_delete {
                master.entities.remove(name);
            }
            master
                .relations
                .retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
        });

        let idx = state.get_search_index().await;
        for name in to_delete {
            drop(idx.delete_document(&name));
        }
        Ok("Entities deleted".to_string())
    }
}

pub struct DeleteObservationsHandler;

#[async_trait]
impl McpTool for DeleteObservationsHandler {
    fn name(&self) -> &'static str {
        "delete_observations"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<DeleteObservationsTool>(
            "delete_observations",
            "Delete specific observations from existing entities in the knowledge graph without removing the entity itself.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: DeleteObservationsTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        
        let mut missing_entities = Vec::new();
        let mut missing_observations = Vec::new();
        
        state.modify_graph(|master| {
            for d in req.deletions {
                if let Some(e) = master.entities.get_mut(&d.entity_name) {
                    for target in d.observations {
                        let target_clean = target.trim();
                        if target_clean.is_empty() { continue; }
                        
                        let initial_len = e.observations.len();
                        e.observations.retain(|o| {
                            let o_clean = o.trim();
                            // Exact match OR substantial substring match (to forgive LLM whitespace/newline errors)
                            !(o_clean == target_clean || (target_clean.len() >= 10 && o_clean.contains(target_clean)))
                        });
                        
                        if e.observations.len() == initial_len {
                            missing_observations.push(target);
                        }
                    }
                } else {
                    missing_entities.push(d.entity_name);
                }
            }
        });

        if !missing_entities.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Entities not found: {}. Please use search_nodes to verify the exact entity names.",
                missing_entities.join(", ")
            )));
        }
        
        if !missing_observations.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Failed to find {} observation(s) to delete. Make sure you provide a sufficiently long exact substring match.",
                missing_observations.len()
            )));
        }

        Ok("Observations deleted".to_string())
    }
}

pub struct DeleteRelationsHandler;

#[async_trait]
impl McpTool for DeleteRelationsHandler {
    fn name(&self) -> &'static str {
        "delete_relations"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<DeleteRelationsTool>(
            "delete_relations",
            "Delete specific relationships between entities in the knowledge graph by specifying source 'from', target 'to', and 'relation_type'.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut deleted_count = 0;
        let requested_count = req.relations.len();
        
        state.modify_graph(|master| {
            let initial_len = master.relations.len();
            master.relations.retain(|r| {
                let should_delete = req.relations.iter().any(|target| {
                    target.from == r.from && target.to == r.to && target.relation_type == r.relation_type
                });
                !should_delete
            });
            deleted_count = initial_len - master.relations.len();
        });
        
        let missing_count = requested_count.saturating_sub(deleted_count);
        if missing_count > 0 {
            return Err(crate::error::AppError::Internal(format!(
                "Error: {} relation(s) not found in graph. Please verify exact relation properties (from, to, relation_type) using read_graph or get_subgraph.",
                missing_count
            )));
        }
        Ok("Relations deleted".to_string())
    }
}

pub struct ReadGraphHandler;

#[async_trait]
impl McpTool for ReadGraphHandler {
    fn name(&self) -> &'static str {
        "read_graph"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ReadGraphTool>(
            "read_graph",
            "Read entities and relations from the knowledge graph with optional namespace filtering and token truncation. For large graphs, specify 'namespace' or use 'search_nodes' or 'get_subgraph' for targeted discovery.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let max_tokens = req.max_tokens;
        let limit = req.limit;
        let offset = req.offset.unwrap_or(0);

        let data = state.read_graph(|full| -> crate::error::Result<String> {
            let filtered_entities: Vec<(&String, &crate::models::Entity)> = full
                .entities
                .iter()
                .filter(|(_, v)| req.namespace.as_ref().is_none_or(|ns| &v.namespace == ns))
                .filter(|(_, v)| req.entity_type.as_ref().is_none_or(|et| &v.entity_type == et))
                .collect();

            let total_count = filtered_entities.len();

            let (entities_map, relations_vec, meta_json) = if limit.is_some() || req.offset.is_some() {
                let effective_limit = limit.unwrap_or(50);
                let paged_entities: std::collections::HashMap<&String, &crate::models::Entity> = filtered_entities
                    .into_iter()
                    .skip(offset)
                    .take(effective_limit)
                    .collect();

                let has_more = offset + paged_entities.len() < total_count;

                let relevant_relations: Vec<&crate::models::Relation> = full
                    .relations
                    .iter()
                    .filter(|r| {
                        let ns_match = req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns);
                        ns_match && (paged_entities.contains_key(&r.from) || paged_entities.contains_key(&r.to))
                    })
                    .collect();

                let meta = serde_json::json!({
                    "total_count": total_count,
                    "offset": offset,
                    "limit": effective_limit,
                    "has_more": has_more,
                });
                (paged_entities, relevant_relations, Some(meta))
            } else {
                let ent_map: std::collections::HashMap<&String, &crate::models::Entity> =
                    filtered_entities.into_iter().collect();
                let rel_vec: Vec<&crate::models::Relation> = full
                    .relations
                    .iter()
                    .filter(|r| req.namespace.as_ref().is_none_or(|ns| &r.namespace == ns))
                    .collect();
                (ent_map, rel_vec, None)
            };

            let serialized = if let Some(meta) = &meta_json {
                serde_json::json!({
                    "entities": entities_map,
                    "relations": relations_vec,
                    "total_count": meta["total_count"],
                    "offset": meta["offset"],
                    "limit": meta["limit"],
                    "has_more": meta["has_more"],
                })
            } else {
                serde_json::json!({
                    "entities": entities_map,
                    "relations": relations_vec,
                })
            };

            let mut result_json = serde_yaml::to_string(&serialized)?;

            if let Some(max_tok) = max_tokens {
                let max_chars = max_tok * 4;
                if result_json.len() > max_chars {
                    let total_ent = entities_map.len();
                    let total_rel = relations_vec.len();
                    let frac = (max_chars as f64) / (result_json.len() as f64);
                    let mut keep_count = (total_ent as f64 * frac * 0.85).max(1.0) as usize;

                    let ent_vec: Vec<_> = entities_map.into_iter().collect();

                    while keep_count > 0 {
                        let sliced_ent: std::collections::HashMap<_, _> = ent_vec.iter().take(keep_count).cloned().collect();
                        let sliced_rel: Vec<_> = relations_vec.iter().filter(|r| sliced_ent.contains_key(&r.from) && sliced_ent.contains_key(&r.to)).copied().collect();

                        let mut truncated_val = serde_json::json!({
                            "entities": sliced_ent,
                            "relations": sliced_rel,
                            "_meta": {
                                "truncated": true,
                                "total_entities": total_ent,
                                "returned_entities": sliced_ent.len(),
                                "total_relations": total_rel,
                                "returned_relations": sliced_rel.len(),
                                "note": "Output structurally truncated to fit within max_tokens. Use pagination (limit, offset), search_nodes, or get_subgraph for targeted discovery."
                            }
                        });

                        if let Some(meta) = &meta_json
                            && let Some(obj) = truncated_val.as_object_mut() {
                                obj.insert("total_count".to_string(), meta["total_count"].clone());
                                obj.insert("offset".to_string(), meta["offset"].clone());
                                obj.insert("limit".to_string(), meta["limit"].clone());
                                obj.insert("has_more".to_string(), serde_json::Value::Bool(true));
                            }

                        if let Ok(trunc_str) = serde_yaml::to_string(&truncated_val)
                            && (trunc_str.len() <= max_chars || keep_count == 1) {
                                result_json = trunc_str;
                                break;
                            }
                        keep_count = keep_count.saturating_sub(keep_count / 3 + 1);
                    }
                }
            }

            Ok(result_json)
        })?;
        Ok(data)
    }
}
pub struct SearchNodesHandler;

#[async_trait]
impl McpTool for SearchNodesHandler {
    fn name(&self) -> &'static str {
        "search_nodes"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<SearchNodesTool>(
            "search_nodes",
            "Search knowledge graph entities and nodes by name or observation query. Returns matching entity names, types, and observations.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let limit = req.limit.unwrap_or(10);
        let include_body = req.include_body.unwrap_or(false);

        let idx = state.get_search_index().await;
        let matches = idx
            .search(&req.query, req.namespace.as_deref())
            .unwrap_or_default();

        let data = state.read_graph(|full| -> crate::error::Result<String> {
            let mut matched_entities = Vec::new();
            for (id, doc_type, _, _, _) in matches.iter().take(limit) {
                if doc_type == "entity"
                    && let Some(e) = full.entities.get(id)
                {
                    if include_body {
                        matched_entities.push(serde_json::to_value(e)?);
                    } else {
                        matched_entities.push(serde_json::json!({
                            "name": e.name,
                            "entity_type": e.entity_type,
                            "namespace": e.namespace,
                            "git_branch": e.git_branch,
                            "observations_count": e.observations.len()
                        }));
                    }
                }
            }
            Ok::<String, crate::error::AppError>(serde_yaml::to_string(&matched_entities)?)
        })?;
        Ok(data)
    }
}

pub struct OpenNodesHandler;

#[async_trait]
impl McpTool for OpenNodesHandler {
    fn name(&self) -> &'static str {
        "open_nodes"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<OpenNodesTool>(
            "open_nodes",
            "Retrieve complete node details, observations, and connected relations for specific entities in the knowledge graph.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.read_graph(|full| -> crate::error::Result<String> {
            let targets: HashSet<&str> = req.names.iter().map(|s| s.as_str()).collect();
            let mut result = BorrowedGraph::default();
            let mut connected: HashSet<&str> = HashSet::new();
            for r in &full.relations {
                if targets.contains(r.from.as_str()) {
                    connected.insert(r.to.as_str());
                    result.relations.push(r);
                } else if targets.contains(r.to.as_str()) {
                    connected.insert(r.from.as_str());
                    result.relations.push(r);
                }
            }
            for (name, e) in &full.entities {
                if targets.contains(name.as_str()) || connected.contains(name.as_str()) {
                    result.entities.insert(name, e);
                }
            }
            Ok::<String, crate::error::AppError>(serde_yaml::to_string(&result)?)
        })?;
        Ok(data)
    }
}

pub struct VisualizeGraphHandler;

#[async_trait]
impl McpTool for VisualizeGraphHandler {
    fn name(&self) -> &'static str {
        "visualize_graph"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<VisualizeGraphTool>(
            "visualize_graph",
            "Generate a visual Mermaid representation of the knowledge graph topology or a filtered query subset.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: VisualizeGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let query = req.query.unwrap_or_default();
        let mut output = state.read_graph(|full| {
            let mut included: HashSet<&str> = HashSet::new();
            let mut to_draw = Vec::new();

            for (name, e) in &full.entities {
                if let Some(ns) = &req.namespace
                    && e.namespace != *ns
                {
                    continue;
                }
                if query.is_empty()
                    || contains_ignore_ascii_case(name, &query)
                    || contains_ignore_ascii_case(&e.entity_type, &query)
                {
                    included.insert(name.as_str());
                }
            }

            for r in &full.relations {
                if let Some(ns) = &req.namespace
                    && r.namespace != *ns
                {
                    continue;
                }
                if query.is_empty()
                    || included.contains(r.from.as_str())
                    || included.contains(r.to.as_str())
                {
                    included.insert(r.from.as_str());
                    included.insert(r.to.as_str());
                    to_draw.push(r);
                }
            }

            let mut out = String::with_capacity(included.len() * 40 + to_draw.len() * 60);
            out.push_str("graph TD;\n");

            let sanitize_to = |out_str: &mut String, s: &str, id_mode: bool| {
                for c in s.chars() {
                    if c != '"' && c != '(' && c != ')' {
                        if id_mode && (c == ' ' || c == '-' || c == '.') {
                            out_str.push('_');
                        } else {
                            out_str.push(c);
                        }
                    }
                }
            };

            for name in &included {
                out.push_str("    id_");
                sanitize_to(&mut out, name, true);
                out.push_str("[\"");
                sanitize_to(&mut out, name, false);
                out.push_str("\"];\n");
            }
            for r in to_draw {
                out.push_str("    id_");
                sanitize_to(&mut out, &r.from, true);
                out.push_str("-->|\"");
                out.push_str(&r.relation_type.replace("\"", ""));
                out.push_str("\"|id_");
                sanitize_to(&mut out, &r.to, true);
                out.push_str(";\n");
            }
            out
        });
        if output == "graph TD;\n" {
            output = "No nodes found to visualize.".to_string();
        }
        Ok(output.to_string())
    }
}

pub struct CondenseEntityHandler;

#[async_trait]
impl McpTool for CondenseEntityHandler {
    fn name(&self) -> &'static str {
        "condense_entity"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<CondenseEntityTool>(
            "condense_entity",
            "Condense, summarize, and deduplicate redundant observations on an entity to optimize context size while preserving key facts.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut missing = false;

        if req.summarized_observations.is_empty() {
            let mut obs_to_condense = Vec::new();
            state.read_graph(|g| {
                if let Some(e) = g.entities.get(&req.entity_name) {
                    obs_to_condense = e.observations.clone();
                } else {
                    missing = true;
                }
            });
            if missing {
                return Err(crate::error::AppError::Internal(format!(
                    "Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
                    req.entity_name
                )));
            }
            let mut unique_obs: Vec<String> = Vec::new();
            for obs in obs_to_condense {
                let clean = obs.trim().to_string();
                if !clean.is_empty() && !unique_obs.contains(&clean) {
                    unique_obs.push(clean);
                }
            }
            if unique_obs.len() > 3 && state.ollama.is_available().await {
                let prompt = format!(
                    "Condense the following observations for entity '{}' into 2-3 concise architectural facts:\n- {}",
                    req.entity_name,
                    unique_obs.join("\n- ")
                );
                if let Ok(summary) = state.ollama.generate(&prompt, None, None, None).await {
                    let lines: Vec<String> = summary
                        .lines()
                        .map(|l| l.trim().trim_start_matches('-').trim().to_string())
                        .filter(|l| !l.is_empty())
                        .collect();
                    if !lines.is_empty() {
                        unique_obs = lines;
                    }
                }
            }
            state.modify_graph(|master| {
                if let Some(e) = master.entities.get_mut(&req.entity_name) {
                    e.observations = unique_obs;
                }
            });
            return Ok(format!(
                "Entity '{}' observations condensed with deduplication/AI summary.",
                req.entity_name
            ));
        }

        state.modify_graph(|master| {
            if let Some(e) = master.entities.get_mut(&req.entity_name) {
                e.observations = req.summarized_observations;
            } else {
                missing = true;
            }
        });
        if missing {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
                req.entity_name
            )));
        }
        Ok("Entity condensed".to_string())
    }
}

pub struct MergeEntitiesHandler;

#[async_trait]
impl McpTool for MergeEntitiesHandler {
    fn name(&self) -> &'static str {
        "merge_entities"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<MergeEntitiesTool>(
            "merge_entities",
            "Merge two duplicate entities into a primary target entity, consolidating all observations and re-pointing relations.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let mut missing = false;
        state.modify_graph(|master| {
            if let Some(src) = master.entities.remove(&req.source_entity) {
                if let Some(tgt) = master.entities.get_mut(&req.target_entity) {
                    tgt.observations.extend(src.observations);
                    tgt.observations.sort();
                    tgt.observations.dedup();
                } else {
                    let mut new_tgt = src.clone();
                    new_tgt.name = req.target_entity.clone();
                    master.entities.insert(req.target_entity.clone(), new_tgt);
                }
            } else {
                missing = true;
            }
            let mut seen = std::collections::HashSet::new();
            master.relations.retain_mut(|r| {
                if r.from == req.source_entity {
                    r.from = req.target_entity.clone();
                }
                if r.to == req.source_entity {
                    r.to = req.target_entity.clone();
                }

                // Filter out self-loops
                if r.from == r.to {
                    return false;
                }

                if r.from == req.target_entity || r.to == req.target_entity {
                    seen.insert(r.clone())
                } else {
                    true
                }
            });
        });
        if missing {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Source entity '{}' not found. Please verify the exact entity name using search_nodes.",
                req.source_entity
            )));
        }
        Ok("Entities merged".to_string())
    }
}

pub struct FindOrphansHandler;

#[async_trait]
impl McpTool for FindOrphansHandler {
    fn name(&self) -> &'static str {
        "find_orphans"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<FindOrphansTool>(
            "find_orphans",
            "Find orphan nodes in the knowledge graph that have no incoming or outgoing relationships.",
        )
    }

    async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let orphans = state.read_graph(|full| {
            let mut connected = std::collections::HashSet::new();
            for r in &full.relations {
                connected.insert(r.from.as_str());
                connected.insert(r.to.as_str());
            }
            full.entities
                .keys()
                .filter(|k| !connected.contains(k.as_str()))
                .cloned()
                .collect::<Vec<String>>()
        });
        let data = Ok::<String, crate::error::AppError>(serde_yaml::to_string(&orphans)?)?;
        Ok(data)
    }
}

pub struct GetSubgraphHandler;

#[async_trait]
impl McpTool for GetSubgraphHandler {
    fn name(&self) -> &'static str {
        "get_subgraph"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<GetSubgraphTool>(
            "get_subgraph",
            "Extract a connected neighborhood subgraph centered around specified seed nodes up to a given hop depth.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: GetSubgraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let root = req.root_entity.ok_or_else(|| {
            crate::error::AppError::Internal("root_entity is required".to_string())
        })?;
        let depth = req.depth.unwrap_or(2);
        let format = req.format.unwrap_or(SubgraphFormat::MarkdownTree);

        let (matched_entities, matched_relations) = state.read_graph(|full| {
            let mut adj: std::collections::HashMap<&str, Vec<&str>> =
                std::collections::HashMap::with_capacity(full.relations.len() * 2);
            for r in &full.relations {
                adj.entry(r.from.as_str()).or_default().push(r.to.as_str());
                adj.entry(r.to.as_str()).or_default().push(r.from.as_str());
            }

            let mut visited = std::collections::HashSet::new();
            let mut queue = std::collections::VecDeque::new();
            queue.push_back((root.as_str(), 0));
            visited.insert(root.clone());

            while let Some((curr, d)) = queue.pop_front() {
                if d < depth
                    && let Some(neighbors) = adj.get(curr)
                {
                    for &next in neighbors {
                        if visited.insert(next.to_string()) {
                            queue.push_back((next, d + 1));
                        }
                    }
                }
            }

            let ent: std::collections::HashMap<String, crate::models::Entity> = visited
                .iter()
                .filter_map(|name| full.entities.get(name).map(|e| (name.clone(), e.clone())))
                .collect();
            let rel: Vec<crate::models::Relation> = full
                .relations
                .iter()
                .filter(|r| visited.contains(&r.from) && visited.contains(&r.to))
                .cloned()
                .collect();

            (ent, rel)
        });

        if format == SubgraphFormat::MarkdownTree {
            let mut markdown = format!("# Subgraph Topology for `{}`\n\n", root);
            markdown.push_str("## Entities\n");
            for e in matched_entities.values() {
                markdown.push_str(&format!("- **{}** (`{}`)\n", e.name, e.entity_type));
                for obs in &e.observations {
                    markdown.push_str(&format!("  - {}\n", obs));
                }
            }

            markdown.push_str("\n## Relations\n");
            for r in &matched_relations {
                markdown.push_str(&format!(
                    "- `{}` --[{}]--> `{}`\n",
                    r.from, r.relation_type, r.to
                ));
            }
            return Ok(markdown);
        }

        let result = serde_json::json!({
            "root_entity": root,
            "depth": depth,
            "entities": matched_entities,
            "relations": matched_relations,
        });

        Ok(serde_yaml::to_string(&result)?)
    }
}

use crate::handlers::utils::*;

pub struct SweepGraphHealthHandler;

#[async_trait]
impl McpTool for SweepGraphHealthHandler {
    fn name(&self) -> &'static str {
        "sweep_graph_health"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<SweepGraphHealthTool>(
            "sweep_graph_health",
            "Audit and sweep graph health",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: SweepGraphHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let auto_prune = req.auto_prune_orphans.unwrap_or(false);
        let auto_prune_stale = req.auto_prune_stale_files.unwrap_or(false);

        let mut orphans = Vec::new();
        let mut duplicates = Vec::new();
        let mut stale_entities = Vec::new();

        state.modify_graph(|g| {
            // 1. Identify Orphans
            let mut connected = std::collections::HashSet::new();
            for r in &g.relations {
                connected.insert(r.from.clone());
                connected.insert(r.to.clone());
            }

            for name in g.entities.keys() {
                if !connected.contains(name) {
                    orphans.push(name.clone());
                }
            }

            // ADR-0111: Automated Stale Symbol Pruning & Graph Tombstoning
            for (name, entity) in g.entities.iter() {
                if entity.entity_type == "File" || name.contains(".rs") || name.contains(".ts") || name.contains(".js") || name.contains(".py") {
                    let check_path = entity.file_path.as_deref().unwrap_or(name.as_str());
                    // Strip any symbol qualifiers like File::symbol
                    let base_file = check_path.split("::").next().unwrap_or(check_path);
                    if (base_file.contains('/') || base_file.contains('\\') || base_file.ends_with(".rs") || base_file.ends_with(".ts")) && !std::path::Path::new(base_file).exists() {
                        stale_entities.push(name.clone());
                    }
                }
            }

            if auto_prune {
                for orphan in &orphans {
                    g.entities.remove(orphan);
                }
            }

            if auto_prune_stale {
                for stale in &stale_entities {
                    g.entities.remove(stale);
                    g.relations.retain(|r| &r.from != stale && &r.to != stale);
                }
            }

            // 2. Compute similarity pairs for duplicate detection using pre-computed lowercase names
            let names: Vec<_> = g.entities.keys().cloned().collect();
            let lower_names: Vec<String> = names.iter().map(|n| n.to_lowercase()).collect();
            for i in 0..names.len() {
                for j in (i + 1)..names.len() {
                    let n1 = &names[i];
                    let n2 = &names[j];
                    let l1 = &lower_names[i];
                    let l2 = &lower_names[j];

                    if l1 == l2 || ((l1.contains(l2.as_str()) || l2.contains(l1.as_str())) && l1.len().min(l2.len()) > 3) {
                        duplicates.push(serde_json::json!({
                            "entity_a": n1,
                            "entity_b": n2,
                            "suggested_action": format!("merge_entities(source: '{}', target: '{}')", n1, n2)
                        }));
                    }
                }
            }
        });

        let report = serde_json::json!({
            "orphaned_entities": orphans,
            "orphans_pruned": auto_prune,
            "stale_entities": stale_entities,
            "stale_pruned": auto_prune_stale,
            "potential_duplicates": duplicates,
            "health_score": if orphans.is_empty() && duplicates.is_empty() && stale_entities.is_empty() { "100%" } else { "Needs Maintenance" }
        });

        Ok(serde_yaml::to_string(&report)?)
    }
}

pub struct ResolveStaleSymbolsHandler;

#[async_trait]
impl McpTool for ResolveStaleSymbolsHandler {
    fn name(&self) -> &'static str {
        "resolve_stale_symbols"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<ResolveStaleSymbolsTool>(
            "resolve_stale_symbols",
            "Inspect Knowledge Graph entities and tech debt symbol/line references against files on disk, healing stale pointers.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: ResolveStaleSymbolsTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let auto_heal = req.auto_heal.unwrap_or(true);

        let mut checked = 0;
        let mut stale_entities = Vec::new();

        state.modify_graph(|g| {
            for name in g.entities.keys() {
                if let Some(target) = &req.target
                    && name != target
                    && !name.contains(target)
                {
                    continue;
                }
                checked += 1;
                let path = std::path::Path::new(name);
                if (name.contains('/') || name.contains('\\') || name.contains('.'))
                    && !path.exists()
                {
                    stale_entities.push(name.clone());
                }
            }

            if auto_heal {
                for stale in &stale_entities {
                    g.entities.remove(stale);
                    g.relations.retain(|r| &r.from != stale && &r.to != stale);
                }
            }
        });

        let report = serde_json::json!({
            "checked_count": checked,
            "stale_entities_found": stale_entities,
            "auto_healed": auto_heal,
            "status": "success"
        });

        Ok(serde_yaml::to_string(&report)?)
    }
}

pub struct SummarizeSubgraphHandler;

#[async_trait]
impl McpTool for SummarizeSubgraphHandler {
    fn name(&self) -> &'static str {
        "summarize_subgraph"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<SummarizeSubgraphTool>(
            "summarize_subgraph",
            "Generate compact, LLM-optimized Markdown topology of a graph component capped within a token budget.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: SummarizeSubgraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let depth = req.depth.unwrap_or(2);

        let (sub_entities, sub_relations) = state.read_graph(|g| {
            let mut visited = std::collections::HashSet::new();
            let mut queue = std::collections::VecDeque::new();
            queue.push_back((req.root_entity.clone(), 0));
            visited.insert(req.root_entity.clone());

            while let Some((curr, d)) = queue.pop_front() {
                if d >= depth {
                    continue;
                }
                for r in &g.relations {
                    if r.from == curr && !visited.contains(&r.to) {
                        visited.insert(r.to.clone());
                        queue.push_back((r.to.clone(), d + 1));
                    } else if r.to == curr && !visited.contains(&r.from) {
                        visited.insert(r.from.clone());
                        queue.push_back((r.from.clone(), d + 1));
                    }
                }
            }

            let ent: Vec<_> = g
                .entities
                .values()
                .filter(|e| visited.contains(&e.name))
                .cloned()
                .collect();
            let rel: Vec<_> = g
                .relations
                .iter()
                .filter(|r| visited.contains(&r.from) && visited.contains(&r.to))
                .cloned()
                .collect();
            (ent, rel)
        });

        let mut markdown = format!("# Subgraph Topology for `{}`\n\n", req.root_entity);
        markdown.push_str("## Entities\n");
        for e in &sub_entities {
            markdown.push_str(&format!("- **{}** (`{}`)\n", e.name, e.entity_type));
            for obs in &e.observations {
                markdown.push_str(&format!("  - {}\n", obs));
            }
        }

        markdown.push_str("\n## Relations\n");
        for r in &sub_relations {
            markdown.push_str(&format!(
                "- `{}` --[{}]--> `{}`\n",
                r.from, r.relation_type, r.to
            ));
        }

        let max_tokens = req.max_tokens.unwrap_or(1000);
        let max_chars = max_tokens * 4;
        if markdown.len() > max_chars {
            let valid_boundary = markdown.floor_char_boundary(max_chars);
            markdown.truncate(valid_boundary);
            markdown.push_str("\n... [Truncated to fit token budget]");
        }

        Ok(markdown)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::handlers::meta::AgentSignalsHandler;
    use serde_json::json;

    #[tokio::test]
    async fn test_create_and_read_entities() {
        let state = Arc::new(MemoryState::new_in_memory());

        let create_handler = CreateEntitiesHandler;
        let args = json!({
            "entities": [
                {"name": "Alice", "entity_type": "Person", "observations": ["Likes Bob"]}
            ]
        });

        let res = create_handler
            .execute(args, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res.contains("Successfully created 1 entity/entities"));

        // Ensure graph contains the entity
        state.graph.read_with(|g| {
            assert!(g.entities.contains_key("Alice"));
            assert_eq!(g.entities["Alice"].entity_type, "Person");
        });
    }

    #[tokio::test]
    async fn test_create_relations() {
        let state = Arc::new(MemoryState::new_in_memory());

        // Needs entities first
        state.graph.modify(|g| {
            g.entities.insert(
                "A".to_string(),
                crate::models::Entity {
                    name: "A".to_string(),
                    entity_type: "Node".to_string(),
                    observations: vec![],
                    namespace: "global".to_string(),
                    git_branch: None,
                    ..Default::default()
                },
            );
            g.entities.insert(
                "B".to_string(),
                crate::models::Entity {
                    name: "B".to_string(),
                    entity_type: "Node".to_string(),
                    observations: vec![],
                    namespace: "global".to_string(),
                    git_branch: None,
                    ..Default::default()
                },
            );
        });

        let handler = CreateRelationsHandler;
        let args = json!({
            "relations": [
                {"from": "A", "to": "B", "relation_type": "knows"}
            ]
        });
        let res = handler
            .execute(args, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res.contains("Successfully created 1 relation(s)"));

        let alias_args = json!({
            "relations": [
                {"from": "A", "to": "B", "relation_type": "knows"}
            ]
        });
        let alias_res = handler.execute(alias_args, state.clone()).await.unwrap();
        assert!(alias_res.contains("Successfully created 1 relation(s)"));

        // Test semantic LLM schema feedback on missing fields
        let bad_args = json!({
            "relations": [
                {"invalid_field": "X"}
            ]
        });
        let err_res = handler
            .execute(bad_args, state.clone())
            .await
            .unwrap_err()
            .to_string();
        assert!(err_res.contains("Schema error:"));
        assert!(err_res.contains("strictly uses 'from', 'to', and 'relation_type'"));
    }

    #[tokio::test]
    async fn test_observations_and_reads() {
        let state = Arc::new(MemoryState::new_in_memory());

        // Need entity first
        state.graph.modify(|g| {
            g.entities.insert(
                "A".to_string(),
                crate::models::Entity {
                    name: "A".to_string(),
                    entity_type: "Node".to_string(),
                    observations: vec![],
                    namespace: "global".to_string(),
                    git_branch: None,
                    ..Default::default()
                },
            );
        });

        let add_obs = AddObservationsHandler;
        let args_obs = json!({
            "observations": [
                {"entity_name": "A", "contents": ["Obs 1", "Obs 2"]}
            ]
        });
        let res1 = add_obs
            .execute(args_obs, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res1, "Observations added");

        let read_graph = ReadGraphHandler;
        let res2 = read_graph
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res2.contains("Obs 1"));
        assert!(res2.contains("Obs 2"));

        let del_entity = DeleteEntitiesHandler;
        let res4 = del_entity
            .execute(json!({"entity_names": ["A"]}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res4, "Entities deleted");

        let res5 = read_graph
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(!res5.contains("A"));
    }

    #[tokio::test]
    async fn test_advanced_graph_operations() {
        let state = Arc::new(MemoryState::new_in_memory());

        let create_handler = CreateEntitiesHandler;
        let args_ent = json!({
            "entities": [
                {"name": "X", "entity_type": "File", "observations": ["Obs X"], "namespace": "global"},
                {"name": "Y", "entity_type": "File", "observations": ["Obs Y"], "namespace": "global"}
            ]
        });
        create_handler
            .execute(args_ent, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();

        let rel_handler = CreateRelationsHandler;
        let args_rel = json!({
            "relations": [
                {"from": "X", "to": "Y", "relation_type": "depends_on", "namespace": "global"}
            ]
        });
        rel_handler
            .execute(args_rel, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();

        let read_handler = ReadGraphHandler;
        let res_read = read_handler
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res_read.contains("X"));
        assert!(res_read.contains("depends_on"));

        let open_handler = OpenNodesHandler;
        let res_open = open_handler
            .execute(json!({"names": ["X"]}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res_open.contains("Y"));

        let viz_handler = VisualizeGraphHandler;
        let res_viz = viz_handler
            .execute(json!({"query": "X"}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(!res_viz.is_empty());

        let condense = CondenseEntityHandler;
        let res_cond = condense
            .execute(
                json!({"entity_name": "X", "summarized_observations": ["X condensed"]}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res_cond, "Entity condensed");

        let merge = MergeEntitiesHandler;
        let res_merge = merge
            .execute(
                json!({"source_entity": "X", "target_entity": "Y"}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res_merge, "Entities merged");

        let orphans = FindOrphansHandler;
        let res_orphans = orphans
            .execute(json!({}), state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        // After merging X into Y and purging self-loops, Y is the sole node and becomes an orphan
        assert!(res_orphans.contains("Y"));
    }

    #[tokio::test]
    async fn test_more_graph_handlers() {
        let state = Arc::new(MemoryState::new_in_memory());

        let create_handler = CreateEntitiesHandler;
        let args_ent = json!({
            "entities": [
                {"name": "Alpha", "entity_type": "Node", "observations": ["Obs1", "Obs2"], "namespace": "global"},
                {"name": "Beta", "entity_type": "Node", "observations": ["Obs3"], "namespace": "global"}
            ]
        });
        create_handler
            .execute(args_ent, state.clone())
            .await
            .unwrap();

        let rel_handler = CreateRelationsHandler;
        let args_rel = json!({
            "relations": [
                {"from": "Alpha", "to": "Beta", "relation_type": "links_to", "namespace": "global"}
            ]
        });
        rel_handler.execute(args_rel, state.clone()).await.unwrap();

        let path_handler = QueryGraphPathHandler;
        let path_res = path_handler
            .execute(
                json!({"start_node": "Alpha", "end_node": "Beta"}),
                state.clone(),
            )
            .await
            .unwrap();
        assert!(path_res.contains("Alpha"));

        let sub_handler = GetSubgraphHandler;
        let sub_res = sub_handler
            .execute(json!({"root_entity": "Alpha"}), state.clone())
            .await
            .unwrap();
        assert!(sub_res.contains("Alpha"));

        let sum_handler = SummarizeSubgraphHandler;
        let sum_res = sum_handler
            .execute(json!({"root_entity": "Alpha"}), state.clone())
            .await
            .unwrap();
        assert!(!sum_res.is_empty());

        let sweep_handler = SweepGraphHealthHandler;
        let sweep_res = sweep_handler
            .execute(json!({}), state.clone())
            .await
            .unwrap();
        assert!(!sweep_res.is_empty());

        let del_obs_handler = DeleteObservationsHandler;
        let del_obs_res = del_obs_handler
            .execute(
                json!({
                    "deletions": [{"entity_name": "Alpha", "observations": ["Obs1"]}]
                }),
                state.clone(),
            )
            .await
            .unwrap();
        assert_eq!(del_obs_res, "Observations deleted");

        let del_rel_handler = DeleteRelationsHandler;
        let del_rel_res = del_rel_handler
            .execute(
                json!({
                    "relations": [{"from": "Alpha", "to": "Beta", "relation_type": "links_to"}]
                }),
                state.clone(),
            )
            .await
            .unwrap();
        assert_eq!(del_rel_res, "Relations deleted");

        let bcast_handler = AgentSignalsHandler;
        let bcast_res = bcast_handler
            .execute(
                json!({
                    "action": "broadcast",
                    "sender": "agent1",
                    "signal_type": "task_completed",
                    "payload": "fix_bug"
                }),
                state.clone(),
            )
            .await
            .unwrap();
        assert!(bcast_res.contains("Broadcasted signal"));

        let qsignal_handler = AgentSignalsHandler;
        let qsignal_res = qsignal_handler
            .execute(
                json!({
                    "action": "query",
                    "sender": "agent1"
                }),
                state.clone(),
            )
            .await
            .unwrap();
        assert!(qsignal_res.contains("task_completed"));
        let read_paged_handler = ReadGraphHandler;
        let paged_res = read_paged_handler
            .execute(json!({"limit": 1, "offset": 0}), state.clone())
            .await
            .unwrap();
        let parsed_paged: serde_json::Value = serde_yaml::from_str(&paged_res).unwrap();
        assert_eq!(parsed_paged["total_count"], 2);
        assert_eq!(parsed_paged["has_more"], true);
    }
}