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>) -> Result<String, 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>) -> Result<String, 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 in req.entities {
                if !entity.name.is_empty() {
                    inserted.push(entity.clone());
                    g.entities.insert(entity.name.clone(), entity);
                }
            }
        });
        let idx = state
            .search_index
            .read()
            .unwrap_or_else(|e| e.into_inner())
            .clone();
        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>) -> Result<String, String> {
        let req: CreateRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.modify_graph(|g| {
            for relation in req.relations {
                if !relation.from.is_empty() && !relation.to.is_empty() {
                    g.relations.push(relation);
                }
            }
        });
        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>) -> Result<String, String> {
        let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        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", "Execute delete_entities")
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, 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();
        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
            .search_index
            .read()
            .unwrap_or_else(|e| e.into_inner())
            .clone();
        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>) -> Result<String, String> {
        let req: DeleteObservationsTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        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));
                }
            }
        });
        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>) -> Result<String, String> {
        let req: DeleteRelationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.modify_graph(|master| {
            let to_rem: HashSet<_> = req.relations.into_iter().collect();
            master.relations.retain(|r| !to_rem.contains(r));
        });
        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>) -> Result<String, String> {
        let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.read_graph(|full| {
            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).unwrap_or_default()
            } else {
                serde_json::to_string(full).unwrap_or_default()
            }
        });
        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>) -> Result<String, 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| {
            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);
                }
            }
            serde_json::to_string(&result).unwrap_or_default()
        });
        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>) -> Result<String, String> {
        let req: OpenNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.read_graph(|full| {
            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);
                }
            }
            serde_json::to_string(&result).unwrap_or_default()
        });
        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>) -> Result<String, 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>) -> Result<String, String> {
        let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.modify_graph(|master| {
            if let Some(e) = master.entities.get_mut(&req.entity_name) {
                e.observations = req.summarized_observations;
            }
        });
        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>) -> Result<String, String> {
        let req: MergeEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        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);
                    MemoryState::deduplicate(&mut tgt.observations);
                } else {
                    let mut new_tgt = src.clone();
                    new_tgt.name = req.target_entity.clone();
                    master.entities.insert(req.target_entity.clone(), new_tgt);
                }
            }
            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
                }
            });
        });
        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>) -> Result<String, 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 = serde_json::to_string(&orphans).unwrap_or_default();
        Ok(data.to_string())
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    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", "entityType": "Person", "observations": ["Likes Bob"]}
            ]
        });

        let res = create_handler.execute(args, state.clone()).await.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", "relationType": "knows"}
            ]
        });
        let res = handler.execute(args, state.clone()).await.unwrap();
        assert_eq!(res, "Relations created");
    }

    #[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": [
                {"entityName": "A", "contents": ["Obs 1", "Obs 2"]}
            ]
        });
        let res1 = add_obs.execute(args_obs, state.clone()).await.unwrap();
        assert_eq!(res1, "Observations added");

        let read_graph = ReadGraphHandler;
        let res2 = read_graph
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .unwrap();
        assert!(res2.contains("Obs 1"));
        assert!(res2.contains("Obs 2"));

        let del_entity = DeleteEntitiesHandler;
        let res4 = del_entity
            .execute(json!({"entityNames": ["A"]}), state.clone())
            .await
            .unwrap();
        assert_eq!(res4, "Entities deleted");

        let res5 = read_graph
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .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", "entityType": "File", "observations": ["Obs X"], "namespace": "global"},
                {"name": "Y", "entityType": "File", "observations": ["Obs Y"], "namespace": "global"}
            ]
        });
        create_handler
            .execute(args_ent, state.clone())
            .await
            .unwrap();

        let rel_handler = CreateRelationsHandler;
        let args_rel = json!({
            "relations": [
                {"from": "X", "to": "Y", "relationType": "depends_on", "namespace": "global"}
            ]
        });
        rel_handler.execute(args_rel, state.clone()).await.unwrap();

        let read_handler = ReadGraphHandler;
        let res_read = read_handler
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .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
            .unwrap();
        assert!(res_open.contains("Y"));

        let viz_handler = VisualizeGraphHandler;
        let res_viz = viz_handler
            .execute(json!({"query": "X"}), state.clone())
            .await
            .unwrap();
        assert!(!res_viz.is_empty());

        let condense = CondenseEntityHandler;
        let res_cond = condense
            .execute(
                json!({"entityName": "X", "summarized_observations": ["X condensed"]}),
                state.clone(),
            )
            .await
            .unwrap();
        assert_eq!(res_cond, "Entity condensed");

        let merge = MergeEntitiesHandler;
        let res_merge = merge
            .execute(
                json!({"sourceEntity": "X", "targetEntity": "Y"}),
                state.clone(),
            )
            .await
            .unwrap();
        assert_eq!(res_merge, "Entities merged");

        let orphans = FindOrphansHandler;
        let res_orphans = orphans.execute(json!({}), state.clone()).await.unwrap();
        assert!(!res_orphans.contains("Y"));
    }
}