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 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", "Execute query_graph_path")
    }

    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())?;
        state.read_graph(|graph| {
            let max_depth = req.max_depth.unwrap_or(5);
            let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new();
            let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
            let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> =
                std::collections::HashMap::new();

            queue.push_back(req.start_node.as_str());
            visited.insert(req.start_node.as_str());

            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 == req.end_node {
                    found = true;
                    break;
                }
                nodes_at_current_depth -= 1;
                if current_depth < max_depth {
                    for rel in &graph.relations {
                        if rel.from == current && !visited.contains(rel.to.as_str()) {
                            visited.insert(rel.to.as_str());
                            parents.insert(
                                rel.to.as_str(),
                                (current, rel.relation_type.as_str(), false),
                            );
                            queue.push_back(rel.to.as_str());
                            nodes_at_next_depth += 1;
                        } else if rel.to == current && !visited.contains(rel.from.as_str()) {
                            visited.insert(rel.from.as_str());
                            parents.insert(
                                rel.from.as_str(),
                                (current, rel.relation_type.as_str(), true),
                            );
                            queue.push_back(rel.from.as_str());
                            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 = req.end_node.as_str();
                while curr != req.start_node {
                    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();
                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
                ))
            }
        })
    }
}

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", "Execute create_entities")
    }

    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 mut entity in req.entities {
                if !entity.name.is_empty() {
                    entity.entity_type = crate::models::normalize_entity_type(&entity.entity_type);
                    inserted.push(entity.clone());
                    g.entities.insert(entity.name.clone(), entity);
                }
            }
        });
        let idx = state.get_search_index();
        for entity in inserted {
            drop(idx.index_entity(&entity));
        }
        Ok("Entities created".to_string())
    }
}

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", "Execute create_relations")
    }

    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.modify_graph(|g| {
            for mut relation in req.relations {
                if !relation.from.is_empty() && !relation.to.is_empty() {
                    relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
                    let from_exists = g.entities.contains_key(&relation.from);
                    let to_exists = g.entities.contains_key(&relation.to);
                    if from_exists && to_exists {
                        g.relations.push(relation);
                    } else {
                        if !from_exists {
                            missing_nodes.insert(relation.from);
                        }
                        if !to_exists {
                            missing_nodes.insert(relation.to);
                        }
                    }
                }
            }
        });
        if !missing_nodes.is_empty() {
            let missing: Vec<_> = missing_nodes.into_iter().collect();
            return Err(crate::error::AppError::Internal(format!(
                "Error: Relations dropped due to missing entities: {}",
                missing.join(", ")
            )));
        }
        Ok("Relations created".to_string())
    }
}

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", "Execute add_observations")
    }

    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.modify_graph(|g| {
            for o in req.observations {
                if let Some(e) = g.entities.get_mut(&o.entity_name) {
                    e.observations.extend(o.contents);
                } else {
                    missing_entities.push(o.entity_name);
                }
            }
        });
        if !missing_entities.is_empty() {
            return Err(crate::error::AppError::Internal(format!(
                "Error: Observations dropped for missing entities: {}",
                missing_entities.join(", ")
            )));
        }
        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", "Execute delete_entities")
    }

    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.modify_graph(|master| {
            for name in &to_delete {
                if master.entities.remove(name).is_none() {
                    missing.push(name.clone());
                }
            }
            master
                .relations
                .retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
        });

        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(", ")
            )));
        }

        let idx = state.get_search_index();
        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",
            "Execute delete_observations",
        )
    }

    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 = Vec::new();
        state.modify_graph(|master| {
            for d in req.deletions {
                if let Some(e) = master.entities.get_mut(&d.entity_name) {
                    let to_rem: HashSet<_> = d.observations.into_iter().collect();
                    e.observations.retain(|o| !to_rem.contains(o));
                } else {
                    missing.push(d.entity_name);
                }
            }
        });
        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(", ")
            )));
        }
        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", "Execute delete_relations")
    }

    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 missing_count = 0;
        state.modify_graph(|master| {
            let to_rem: HashSet<_> = req.relations.into_iter().collect();
            let initial_len = master.relations.len();
            master.relations.retain(|r| !to_rem.contains(r));
            missing_count = to_rem.len() - (initial_len - master.relations.len());
        });
        if missing_count > 0 {
            return Err(crate::error::AppError::Internal(format!(
                "Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.",
                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", "Execute read_graph")
    }

    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 data = state.read_graph(|full| -> crate::error::Result<String> {
            let mut result_json = if let Some(ns) = req.namespace {
                let mut filtered = BorrowedGraph::default();
                for (k, v) in &full.entities {
                    if v.namespace == ns {
                        filtered.entities.insert(k, v);
                    }
                }
                for r in &full.relations {
                    if r.namespace == ns {
                        filtered.relations.push(r);
                    }
                }
                serde_json::to_string(&filtered)?
            } else {
                serde_json::to_string(full)?
            };

            if let Some(max_tok) = max_tokens {
                let max_chars = max_tok * 4;
                if result_json.len() > max_chars {
                    result_json.truncate(max_chars);
                    result_json.push_str("... [TRUNCATED_TO_MAX_TOKENS]");
                }
            }
            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", "Execute search_nodes")
    }

    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 matches = if let Ok(idx) = state.search_index.read() {
            idx.search(&req.query, req.namespace.as_deref())
                .unwrap_or_default()
        } else {
            vec![]
        };

        let data = state.read_graph(|full| -> crate::error::Result<String> {
            let mut result = BorrowedGraph::default();
            for (id, doc_type, _, _, _) in &matches {
                if doc_type == "entity"
                    && let Some(e) = full.entities.get(id)
                {
                    result.entities.insert(id, e);
                }
            }
            Ok::<String, crate::error::AppError>(serde_json::to_string(&result)?)
        })?;
        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", "Execute open_nodes")
    }

    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_json::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", "Execute visualize_graph")
    }

    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", "Execute condense_entity")
    }

    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).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", "Execute merge_entities")
    }

    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();
                }

                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", "Execute find_orphans")
    }

    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_json::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", "Execute get_subgraph")
    }

    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
            .or(req.root_node)
            .ok_or_else(|| crate::error::AppError::Internal("root_entity or root_node is required".to_string()))?;
        let depth = req.depth.unwrap_or(2);
        let format = req.format.unwrap_or(SubgraphFormat::Json);

        if format == SubgraphFormat::MarkdownTree {
            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((root.clone(), 0));
                visited.insert(root.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", root);
            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));
            }
            return Ok(markdown);
        }

        let data = state.read_graph(|full| {
            let mut visited_nodes: std::collections::HashSet<String> = std::collections::HashSet::new();
            let mut matched_entities = std::collections::HashMap::new();
            let mut matched_relations = Vec::new();

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

            while let Some((curr, current_depth)) = queue.pop_front() {
                if let Some(entity) = full.entities.get(&curr) {
                    matched_entities.insert(curr.clone(), entity.clone());
                }

                if current_depth < depth {
                    for rel in &full.relations {
                        if rel.from == curr {
                            matched_relations.push(rel.clone());
                            if !visited_nodes.contains(&rel.to) {
                                visited_nodes.insert(rel.to.clone());
                                queue.push_back((rel.to.clone(), current_depth + 1));
                            }
                        } else if rel.to == curr {
                            matched_relations.push(rel.clone());
                            if !visited_nodes.contains(&rel.from) {
                                visited_nodes.insert(rel.from.clone());
                                queue.push_back((rel.from.clone(), current_depth + 1));
                            }
                        }
                    }
                }
            }

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

            Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?)
        })?;

        Ok(data)
    }
}

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 mut orphans = Vec::new();
        let mut duplicates = 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());
                }
            }

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

            // 2. Compute similarity pairs for duplicate detection
            let names: Vec<_> = g.entities.keys().cloned().collect();
            for i in 0..names.len() {
                for j in (i + 1)..names.len() {
                    let n1 = &names[i];
                    let n2 = &names[j];

                    let l1 = n1.to_lowercase();
                    let l2 = n2.to_lowercase();
                    if l1 == l2 || ((l1.contains(&l2) || l2.contains(&l1)) && 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,
            "potential_duplicates": duplicates,
            "health_score": if orphans.is_empty() && duplicates.is_empty() { "100%" } else { "Needs Maintenance" }
        });

        Ok(serde_json::to_string_pretty(&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, _entity) in g.entities.iter_mut() {
                if let Some(target) = &req.target {
                    if 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_json::to_string_pretty(&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 {
            markdown.truncate(max_chars);
            markdown.push_str("\n... [Truncated to fit token budget]");
        }

        Ok(markdown)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::handlers::meta::{BroadcastAgentSignalHandler, QueryAgentSignalsHandler};
    use serde_json::json;
    use tempfile::tempdir;


    #[tokio::test]
    async fn test_create_and_read_entities() {
        let dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));

        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_eq!(res, "Entities created");

        // 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 dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));

        // 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,
                },
            );
            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,
                },
            );
        });

        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_eq!(res, "Relations created");

        // Test semantic LLM schema feedback (User request)
        let bad_args = json!({
            "relations": [
                {"source": "A", "target": "B", "relationType": "knows"}
            ]
        });
        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 dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));

        // 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,
                },
            );
        });

        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 dir = tempfile::tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));

        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();
        assert!(!res_orphans.contains("Y"));
    }

    #[tokio::test]
    async fn test_more_graph_handlers() {
        let dir = tempfile::tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));

        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 = BroadcastAgentSignalHandler;
        let bcast_res = bcast_handler.execute(json!({
            "sender": "agent1",
            "signal_type": "task_completed",
            "payload": "fix_bug"
        }), state.clone()).await.unwrap();
        assert!(bcast_res.contains("Broadcasted signal"));


        let qsignal_handler = QueryAgentSignalsHandler;
        let qsignal_res = qsignal_handler.execute(json!({"sender": "agent1"}), state.clone()).await.unwrap();
        assert!(qsignal_res.contains("task_completed"));

    }
}