use crate::axon::{AxonMessage, AxonMethod, AxonResponse};
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;

type ToolFn = Arc<dyn Fn(Value) -> Value + Send + Sync>;

/// Routes Axon method names to handler functions.
///
/// Each handler takes the `params` field of an AxonMessage and returns
/// a JSON Value to be wrapped in an AxonResponse.
pub struct AxonHandler {
    tools: HashMap<String, ToolFn>,
}

impl AxonHandler {
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// Register a tool handler by name.
    pub fn register<F>(&mut self, name: impl Into<String>, f: F)
    where
        F: Fn(Value) -> Value + Send + Sync + 'static,
    {
        self.tools.insert(name.into(), Arc::new(f));
    }

    /// Dispatch an AxonMessage to the appropriate handler.
    pub fn dispatch(&self, msg: AxonMessage) -> AxonResponse {
        match msg.method {
            AxonMethod::ToolCall => {
                let tool_name = msg
                    .params
                    .get("tool")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                let tool_params = msg
                    .params
                    .get("params")
                    .cloned()
                    .unwrap_or(Value::Object(Default::default()));

                if let Some(f) = self.tools.get(tool_name) {
                    let result = f(tool_params);
                    AxonResponse::ok(msg.id, result)
                } else {
                    AxonResponse::err(msg.id, format!("unknown tool: {}", tool_name))
                }
            }
            AxonMethod::Ping => {
                AxonResponse::ok(msg.id, serde_json::json!({ "pong": true }))
            }
            other => {
                AxonResponse::err(msg.id, format!("unknown method: {}", other.as_str()))
            }
        }
    }

    /// List all registered tool names.
    pub fn tool_names(&self) -> Vec<String> {
        let mut names: Vec<_> = self.tools.keys().cloned().collect();
        names.sort();
        names
    }
}

impl Default for AxonHandler {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn dispatch_ping() {
        let handler = AxonHandler::new();
        let msg = AxonMessage {
            id: "1".into(),
            method: AxonMethod::Ping,
            params: serde_json::json!({}),
        };
        let resp = handler.dispatch(msg);
        assert!(resp.success);
        assert_eq!(resp.result["pong"], true);
    }

    #[test]
    fn dispatch_tool_call() {
        let mut handler = AxonHandler::new();
        handler.register("remember", |params| {
            serde_json::json!({ "stored": params.get("content").cloned().unwrap_or_default() })
        });

        let msg = AxonMessage {
            id: "2".into(),
            method: AxonMethod::ToolCall,
            params: serde_json::json!({ "tool": "remember", "params": { "content": "hello" } }),
        };
        let resp = handler.dispatch(msg);
        assert!(resp.success);
        assert_eq!(resp.result["stored"], "hello");
    }

    #[test]
    fn dispatch_unknown_tool() {
        let handler = AxonHandler::new();
        let msg = AxonMessage {
            id: "3".into(),
            method: AxonMethod::ToolCall,
            params: serde_json::json!({ "tool": "nonexistent", "params": {} }),
        };
        let resp = handler.dispatch(msg);
        assert!(!resp.success);
    }

    #[test]
    fn dispatch_unhandled_method() {
        let handler = AxonHandler::new();
        let msg = AxonMessage {
            id: "4".into(),
            method: AxonMethod::Subscribe,
            params: serde_json::json!({}),
        };
        let resp = handler.dispatch(msg);
        assert!(!resp.success);
        assert!(resp.result["error"].as_str().unwrap().contains("unknown method"));
    }

    #[test]
    fn tool_names_sorted() {
        let mut handler = AxonHandler::new();
        handler.register("remember", |_| Value::Null);
        handler.register("recall", |_| Value::Null);
        handler.register("plan_work", |_| Value::Null);
        let names = handler.tool_names();
        let mut expected = vec!["plan_work", "recall", "remember"];
        expected.sort();
        assert_eq!(names, expected);
    }
}