use crate::state::MemoryState;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;

#[async_trait]
pub trait McpTool: Send + Sync {
    /// The unique name of the tool
    fn name(&self) -> &'static str;

    /// The JSON schema for the tool
    fn schema(&self) -> Value;

    /// Execute the tool with the given arguments
    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String>;
}

pub struct MemoryHandler {
    pub state: Arc<MemoryState>,
    pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
}

impl MemoryHandler {
    pub fn new(state: Arc<MemoryState>) -> Self {
        let mut tools: std::collections::HashMap<String, Box<dyn McpTool>> =
            std::collections::HashMap::new();

        macro_rules! register {
            ($module:ident::$handler:ident) => {
                let h = crate::handlers::$module::$handler;
                tools.insert(h.name().to_string(), Box::new(h));
            };
        }

        register!(graph::QueryGraphPathHandler);
        register!(graph::CreateEntitiesHandler);
        register!(graph::CreateRelationsHandler);
        register!(graph::AddObservationsHandler);
        register!(graph::DeleteEntitiesHandler);
        register!(graph::DeleteObservationsHandler);
        register!(graph::DeleteRelationsHandler);
        register!(graph::ReadGraphHandler);
        register!(graph::SearchNodesHandler);
        register!(graph::OpenNodesHandler);
        register!(graph::VisualizeGraphHandler);
        register!(graph::CondenseEntityHandler);
        register!(graph::MergeEntitiesHandler);
        register!(graph::FindOrphansHandler);

        register!(tasks::AddTaskHandler);
        register!(tasks::DeleteTaskHandler);
        register!(tasks::UpdateTaskStatusHandler);
        register!(tasks::ListActiveTasksHandler);
        register!(tasks::SetAcceptanceCriteriaHandler);
        register!(tasks::VerifyAcceptanceCriteriaHandler);
        register!(tasks::AddMilestoneHandler);
        register!(tasks::UpdateMilestoneHandler);
        register!(tasks::ListMilestonesHandler);

        register!(notes::AddStickyNoteHandler);
        register!(notes::ReadStickyNotesHandler);
        register!(notes::DeleteStickyNoteHandler);
        register!(notes::ClearStickyNotesHandler);
        register!(notes::LeaveHandoffMemoHandler);
        register!(notes::ReadHandoffMemosHandler);
        register!(notes::ClearHandoffMemosHandler);
        register!(notes::AddSessionSummaryHandler);
        register!(notes::GenerateStandupReportHandler);

        register!(meta::LogDecisionHandler);
        register!(meta::QueryDecisionsHandler);
        register!(meta::LogErrorFixHandler);
        register!(meta::SearchErrorFixesHandler);
        register!(meta::LogCodeChangeHandler);
        register!(meta::QueryRecentChangesHandler);
        register!(meta::LearnPreferenceHandler);
        register!(meta::ReadPreferencesHandler);
        register!(meta::LogTechDebtHandler);
        register!(meta::ResolveTechDebtHandler);
        register!(meta::ListTechDebtHandler);
        register!(meta::OmniSearchHandler);
        register!(meta::GetProjectHealthHandler);

        register!(env::UpdateEnvFingerprintHandler);
        register!(env::ReadEnvFingerprintHandler);
        register!(env::LogEnvRequirementHandler);
        register!(env::RegisterEnvironmentHandler);
        register!(env::GetEnvironmentDetailsHandler);

        register!(workspaces::PinFileHandler);
        register!(workspaces::UnpinFileHandler);
        register!(workspaces::ListPinnedFilesHandler);
        register!(workspaces::StoreSnippetHandler);
        register!(workspaces::SearchSnippetsHandler);
        register!(workspaces::DeleteSnippetHandler);
        register!(workspaces::SaveContextWorkspaceHandler);
        register!(workspaces::LoadContextWorkspaceHandler);
        register!(workspaces::ListContextWorkspacesHandler);
        register!(workspaces::AddPrChecklistItemHandler);
        register!(workspaces::GetPrChecklistHandler);
        register!(workspaces::ClearPrChecklistHandler);

        Self { state, tools }
    }

    pub async fn handle_request(&self, req: serde_json::Value) -> Option<serde_json::Value> {
        let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
        let id_clone = id.clone();
        let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");

        match method {
            "server/discover" => {
                let payload = serde_json::json!({
                    "resultType": "complete",
                    "ttlMs": 0,
                    "cacheScope": "public",
                    "supportedVersions": ["2026-07-28", "2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"],
                    "capabilities": {
                        "tools": serde_json::json!({}),
                        "resources": serde_json::json!({}),
                        "prompts": serde_json::json!({})
                    },
                    "_meta": {
                        "io.modelcontextprotocol/serverInfo": {
                            "name": "gemini-mcp-memory",
                            "version": "3.0.0"
                        }
                    }
                });
                Some(crate::mcp::success(id, payload))
            }
            "initialize" => {
                let init = rmcp::model::InitializeResult::new(
                    rmcp::model::ServerCapabilities::builder()
                        .enable_tools()
                        .enable_resources()
                        .enable_prompts()
                        .build(),
                )
                .with_server_info(rmcp::model::Implementation::new(
                    "gemini-mcp-memory",
                    "3.0.0",
                ))
                .with_instructions(include_str!("instructions.md"));
                Some(crate::mcp::success(
                    id,
                    serde_json::to_value(&init).unwrap_or_default(),
                ))
            }
            "notifications/initialized" => None,
            "tools/list" => {
                let mut tools: Vec<serde_json::Value> =
                    self.tools.values().map(|t| t.schema()).collect();
                tools.sort_by_key(|t| {
                    t.get("name")
                        .and_then(|n| n.as_str())
                        .unwrap_or("")
                        .to_string()
                });
                Some(crate::mcp::success(
                    id,
                    serde_json::json!({ "tools": tools }),
                ))
            }
            "resources/list" => {
                let payload = serde_json::json!({
                    "resources": [
                        {
                            "uri": "memory://graph/entities",
                            "name": "Graph Entities",
                            "mimeType": "application/json",
                            "description": "All nodes and entities currently stored in the knowledge graph"
                        },
                        {
                            "uri": "memory://graph/relations",
                            "name": "Graph Relations",
                            "mimeType": "application/json",
                            "description": "All edge relationships between entities in the knowledge graph"
                        },
                        {
                            "uri": "memory://tasks/active",
                            "name": "Active Tasks",
                            "mimeType": "application/json",
                            "description": "All currently active or uncompleted tracking tasks"
                        }
                    ]
                });
                Some(crate::mcp::success(id, payload))
            }
            "resources/templates/list" => {
                let payload = serde_json::json!({
                    "resourceTemplates": []
                });
                Some(crate::mcp::success(id, payload))
            }
            "resources/read" => {
                let params = req.get("params").unwrap_or(&serde_json::Value::Null);
                let uri = params
                    .get("uri")
                    .and_then(|u| u.as_str())
                    .unwrap_or("")
                    .to_string();

                let state_clone = Arc::clone(&self.state);
                let uri_clone = uri.clone();
                let text = match tokio::task::spawn_blocking(move || match uri_clone.as_str() {
                    "memory://graph/entities" => {
                        let graph = state_clone.graph.cache.read().unwrap();
                        let data: Vec<_> = graph.entities.values().collect();
                        Some(serde_json::to_string_pretty(&data).unwrap_or_default())
                    }
                    "memory://graph/relations" => {
                        let graph = state_clone.graph.cache.read().unwrap();
                        let data = &graph.relations;
                        Some(serde_json::to_string_pretty(&data).unwrap_or_default())
                    }
                    "memory://tasks/active" => {
                        let tasks = state_clone.tasks.cache.read().unwrap();
                        let data: Vec<_> = tasks
                            .iter()
                            .filter(|t| t.status != "completed" && t.status != "done")
                            .collect();
                        Some(serde_json::to_string_pretty(&data).unwrap_or_default())
                    }
                    _ => None,
                })
                .await
                {
                    Ok(Some(text)) => text,
                    _ => return Some(crate::mcp::error(id, -32602, "Resource not found")),
                };

                let payload = serde_json::json!({
                    "contents": [{
                        "uri": uri,
                        "mimeType": "application/json",
                        "text": text
                    }]
                });
                Some(crate::mcp::success(id, payload))
            }
            "prompts/list" => {
                let payload = serde_json::json!({
                    "prompts": [
                        {
                            "name": "analyze_tech_debt",
                            "description": "Analyze the project's current technical debt",
                            "arguments": []
                        }
                    ]
                });
                Some(crate::mcp::success(id, payload))
            }
            "prompts/get" => {
                let params = req.get("params").unwrap_or(&serde_json::Value::Null);
                let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");

                if name == "analyze_tech_debt" {
                    let payload = serde_json::json!({
                        "messages": [
                            {
                                "role": "user",
                                "content": {
                                    "type": "text",
                                    "text": "Please read the current technical debt from the knowledge graph (or use the list_tech_debt tool) and provide a prioritization plan for addressing it."
                                }
                            }
                        ]
                    });
                    Some(crate::mcp::success(id, payload))
                } else {
                    Some(crate::mcp::error(id, -32602, "Prompt not found"))
                }
            }
            "tools/call" => {
                let params = req.get("params").unwrap_or(&serde_json::Value::Null);
                let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
                let args = params
                    .get("arguments")
                    .cloned()
                    .unwrap_or(serde_json::Value::Object(Default::default()));

                self.state
                    .broadcast_activity(&format!("Agent executed tool: {}", name));

                let result: Result<String, String> = if let Some(tool) = self.tools.get(name) {
                    tool.execute(args, self.state.clone()).await
                } else {
                    Err(format!("Unknown tool: {}", name))
                };

                match result {
                    Ok(text) => {
                        let payload = serde_json::json!({
                            "content": [{"type": "text", "text": text}],
                            "isError": false
                        });
                        Some(crate::mcp::success(id_clone, payload))
                    }
                    Err(e) => {
                        tracing::error!("Tool {} failed: {}", name, e);
                        let payload = serde_json::json!({
                            "content": [{"type": "text", "text": e}],
                            "isError": true
                        });
                        Some(crate::mcp::success(id_clone, payload))
                    }
                }
            }
            m if m.starts_with("notifications/") => None,
            "ping" => Some(crate::mcp::success(id, serde_json::json!({}))),
            _ => {
                if id.is_null() {
                    None
                } else {
                    Some(crate::mcp::error(
                        id,
                        -32601,
                        &format!("Method {} not found", method),
                    ))
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use tempfile::tempdir;

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

        // Assert some known tools are registered
        assert!(handler.tools.contains_key("create_entities"));
        assert!(handler.tools.contains_key("add_task"));

        // Ensure we can fetch list of tools
        let list_tools_req = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tools/list",
            "params": {}
        });

        let res_list = handler.handle_request(list_tools_req).await.unwrap();
        assert_eq!(res_list["jsonrpc"], "2.0");
        assert_eq!(res_list["id"], 1);
        assert!(res_list["result"]["tools"].as_array().unwrap().len() > 10);
    }

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

        // 1. Test successful tool call (e.g. read_graph)
        let req_success = json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/call",
            "params": {
                "name": "read_graph",
                "arguments": {}
            }
        });
        let res_success = handler.handle_request(req_success).await.unwrap();
        assert_eq!(res_success["jsonrpc"], "2.0");
        assert_eq!(res_success["id"], 2);
        // A successful tool call should return a result with isError: false
        assert_eq!(res_success["result"]["isError"], false);
        assert!(res_success["result"]["content"].as_array().is_some());

        // 2. Test tool call semantic failure (e.g. updating a task that does not exist)
        let req_fail = json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "update_task_status",
                "arguments": {
                    "id": "nonexistent_task_123",
                    "status": "in_progress"
                }
            }
        });
        let res_fail = handler.handle_request(req_fail).await.unwrap();
        assert_eq!(res_fail["jsonrpc"], "2.0");
        assert_eq!(res_fail["id"], 3);
        // Semantic failures must explicitly return isError: true inside the result to halt the LLM
        assert_eq!(res_fail["result"]["isError"], true);
        assert!(res_fail["result"]["content"][0]["text"].as_str().unwrap().contains("not found"));

        // 3. Test unknown JSON-RPC method returns JSON-RPC protocol error
        let req_unknown = json!({
            "jsonrpc": "2.0",
            "id": 4,
            "method": "unknown_method_xyz"
        });
        let res_unknown = handler.handle_request(req_unknown).await.unwrap();
        assert!(res_unknown.get("error").is_some());
        assert_eq!(res_unknown["error"]["code"], -32601);
    }
}