use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::error::{Error, Result};

/// A single tool exposed to an MCP client.
///
/// Implementors must declare `kind()` — either `Read` or `Write(capability)`.
/// There is no default, because forgetting the classification is the sort of
/// hazard we can prevent at compile time.
#[async_trait]
pub trait Tool: Send + Sync + 'static {
    fn name(&self) -> &str;
    fn description(&self) -> &str;

    /// Read or Write. Writes carry a capability that must be granted per session.
    fn kind(&self) -> ToolKind;

    /// JSON Schema for the tool's arguments. Should describe a single object.
    fn input_schema(&self) -> Value;

    async fn call(&self, args: Value) -> Result<ToolCallResult>;

    /// If this tool is a hard refusal, return a stable, model-visible reason.
    /// Registered refusals appear in `tools/list` so external clients can see
    /// what the app has chosen not to offer, and the model gets a clear signal
    /// instead of inventing tool names.
    fn refusal_reason(&self) -> Option<&str> {
        None
    }

    /// Opt-in flag for the compact projection used by small models. Defaults
    /// to `false` — conservative: apps opt tools into the compact surface only
    /// after verifying a small model can call them reliably.
    fn small_model_safe(&self) -> bool {
        false
    }
}

/// Classification of a tool's side-effect profile.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ToolKind {
    /// No state change. Freely callable subject to app-level connection auth.
    Read,
    /// Mutates app state. Requires the caller to hold the named capability.
    Write(WriteCapability),
}

impl ToolKind {
    pub fn is_write(&self) -> bool {
        matches!(self, ToolKind::Write(_))
    }

    pub fn capability_id(&self) -> Option<&str> {
        match self {
            ToolKind::Write(cap) => Some(&cap.id),
            ToolKind::Read => None,
        }
    }
}

/// A named, pre-declared write capability.
///
/// `id` is a stable public contract owned by the app — treat it like an API
/// surface identifier. If you rename a capability, you introduce a new one and
/// deprecate the old, you do not silently reassign an existing id.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WriteCapability {
    pub id: String,
    pub description: String,
}

impl WriteCapability {
    pub fn new(id: impl Into<String>, description: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            description: description.into(),
        }
    }
}

/// Result of a single tool invocation.
///
/// Shape mirrors MCP's `CallToolResult`: a list of content parts plus an
/// `is_error` flag. MVP carries only text content; media/resource parts land
/// alongside audiofiles integration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallResult {
    pub content: Vec<ContentPart>,
    #[serde(default, rename = "isError")]
    pub is_error: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentPart {
    Text { text: String },
}

impl ToolCallResult {
    pub fn text(text: impl Into<String>) -> Self {
        Self {
            content: vec![ContentPart::Text { text: text.into() }],
            is_error: false,
        }
    }

    pub fn error(text: impl Into<String>) -> Self {
        Self {
            content: vec![ContentPart::Text { text: text.into() }],
            is_error: true,
        }
    }
}

/// Serializable description of a tool for `tools/list` responses and for
/// translating to provider-native schemas (Ollama, OpenAI, Anthropic).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolSpec {
    pub name: String,
    pub description: String,
    #[serde(rename = "inputSchema")]
    pub input_schema: Value,

    /// `"read"` or `"write"`. Extension over the MCP base spec; clients that
    /// don't understand this can ignore it.
    #[serde(rename = "kburgKind")]
    pub kind: &'static str,

    /// Only set when `kind == "write"`. Extension.
    #[serde(rename = "kburgCapability", skip_serializing_if = "Option::is_none")]
    pub capability: Option<WriteCapability>,

    /// Set when this tool is a registered refusal. Extension.
    #[serde(rename = "kburgRefusalReason", skip_serializing_if = "Option::is_none")]
    pub refusal_reason: Option<String>,
}

/// Which projection of the registry to expose to a given client.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SurfaceProjection {
    /// All registered tools.
    Full,
    /// Only tools annotated `small_model_safe`.
    Compact,
}

impl SurfaceProjection {
    fn includes(self, small_model_safe: bool) -> bool {
        match self {
            SurfaceProjection::Full => true,
            SurfaceProjection::Compact => small_model_safe,
        }
    }
}

/// A registered "hard refusal" — appears in `tools/list` and returns a stable
/// error on invocation, preventing the model from inventing plausible names.
pub struct Refusal {
    name: String,
    description: String,
    reason: String,
}

impl Refusal {
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        reason: impl Into<String>,
    ) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            reason: reason.into(),
        }
    }
}

#[async_trait]
impl Tool for Refusal {
    fn name(&self) -> &str {
        &self.name
    }
    fn description(&self) -> &str {
        &self.description
    }
    fn kind(&self) -> ToolKind {
        ToolKind::Read
    }
    fn input_schema(&self) -> Value {
        serde_json::json!({ "type": "object", "properties": {} })
    }
    fn refusal_reason(&self) -> Option<&str> {
        Some(&self.reason)
    }
    async fn call(&self, _args: Value) -> Result<ToolCallResult> {
        Err(Error::Refused {
            tool: self.name.clone(),
            reason: self.reason.clone(),
        })
    }
}

/// Registry of tools available to clients.
///
/// Cheaply cloneable — the underlying map lives behind `Arc` so the server
/// and any agents can share the same registry without copying.
#[derive(Clone, Default)]
pub struct ToolRegistry {
    tools: Arc<HashMap<String, Arc<dyn Tool>>>,
}

impl ToolRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a tool. Panics on duplicate names or if called after cloning
    /// (register all tools before serving).
    pub fn register<T: Tool>(&mut self, tool: T) {
        let name = tool.name().to_string();
        let map = Arc::get_mut(&mut self.tools).expect(
            "ToolRegistry::register called after cloning; register all tools before serving",
        );
        assert!(
            !map.contains_key(&name),
            "duplicate tool registration: {name}"
        );
        map.insert(name, Arc::new(tool));
    }

    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
        self.tools.get(name).cloned()
    }

    /// All specs, unfiltered. Equivalent to `specs_projected(Full)`.
    pub fn specs(&self) -> Vec<ToolSpec> {
        self.specs_projected(SurfaceProjection::Full)
    }

    pub fn specs_projected(&self, projection: SurfaceProjection) -> Vec<ToolSpec> {
        self.tools
            .values()
            .filter(|t| projection.includes(t.small_model_safe()))
            .map(|t| ToolSpec {
                name: t.name().to_string(),
                description: t.description().to_string(),
                input_schema: t.input_schema(),
                kind: match t.kind() {
                    ToolKind::Read => "read",
                    ToolKind::Write(_) => "write",
                },
                capability: match t.kind() {
                    ToolKind::Write(cap) => Some(cap),
                    ToolKind::Read => None,
                },
                refusal_reason: t.refusal_reason().map(str::to_string),
            })
            .collect()
    }

    /// Distinct write capabilities declared by registered tools.
    pub fn write_capabilities(&self) -> Vec<WriteCapability> {
        let mut seen = HashSet::new();
        let mut out = Vec::new();
        for tool in self.tools.values() {
            if let ToolKind::Write(cap) = tool.kind()
                && seen.insert(cap.id.clone())
            {
                out.push(cap);
            }
        }
        out
    }

    /// Call a tool, enforcing the granted-capabilities set. `grants` may be
    /// `None` to bypass the check (e.g., for a fully-trusted in-process
    /// caller); pass `Some(set)` from any surface exposed to an LLM.
    pub async fn call(
        &self,
        name: &str,
        args: Value,
        grants: Option<&HashSet<String>>,
    ) -> Result<ToolCallResult> {
        let tool = self
            .get(name)
            .ok_or_else(|| Error::ToolNotFound(name.to_string()))?;

        if let (ToolKind::Write(cap), Some(grants)) = (tool.kind(), grants)
            && !grants.contains(&cap.id)
        {
            return Err(Error::CapabilityDenied {
                tool: name.to_string(),
                capability: cap.id,
            });
        }

        tool.call(args).await
    }
}

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

    struct Echo;
    #[async_trait]
    impl Tool for Echo {
        fn name(&self) -> &'static str {
            "echo"
        }
        fn description(&self) -> &'static str {
            "read-only echo"
        }
        fn kind(&self) -> ToolKind {
            ToolKind::Read
        }
        fn small_model_safe(&self) -> bool {
            true
        }
        fn input_schema(&self) -> Value {
            json!({ "type": "object", "properties": {} })
        }
        async fn call(&self, _args: Value) -> Result<ToolCallResult> {
            Ok(ToolCallResult::text("ok"))
        }
    }

    struct Mutate;
    #[async_trait]
    impl Tool for Mutate {
        fn name(&self) -> &'static str {
            "mutate"
        }
        fn description(&self) -> &'static str {
            "write, gated on `test.write`"
        }
        fn kind(&self) -> ToolKind {
            ToolKind::Write(WriteCapability::new("test.write", "mutate test state"))
        }
        fn input_schema(&self) -> Value {
            json!({ "type": "object", "properties": {} })
        }
        async fn call(&self, _args: Value) -> Result<ToolCallResult> {
            Ok(ToolCallResult::text("mutated"))
        }
    }

    fn registry() -> ToolRegistry {
        let mut r = ToolRegistry::new();
        r.register(Echo);
        r.register(Mutate);
        r.register(Refusal::new(
            "delete_everything",
            "not offered",
            "destructive bulk ops are never exposed",
        ));
        r
    }

    fn grants(ids: &[&str]) -> HashSet<String> {
        ids.iter().map(std::string::ToString::to_string).collect()
    }

    #[tokio::test]
    async fn read_is_callable_without_grants() {
        let r = registry();
        let out = r
            .call("echo", json!({}), Some(&HashSet::new()))
            .await
            .unwrap();
        assert!(!out.is_error);
    }

    #[tokio::test]
    async fn write_denied_without_grant() {
        let r = registry();
        let err = r
            .call("mutate", json!({}), Some(&HashSet::new()))
            .await
            .unwrap_err();
        match err {
            Error::CapabilityDenied { capability, .. } => assert_eq!(capability, "test.write"),
            other => panic!("expected CapabilityDenied, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn write_allowed_with_grant() {
        let r = registry();
        let out = r
            .call("mutate", json!({}), Some(&grants(&["test.write"])))
            .await
            .unwrap();
        assert_eq!(out.content.len(), 1);
        assert!(!out.is_error);
    }

    #[tokio::test]
    async fn none_grants_bypasses_the_check() {
        let r = registry();
        let out = r.call("mutate", json!({}), None).await.unwrap();
        assert!(!out.is_error);
    }

    #[tokio::test]
    async fn unknown_tool_is_not_found() {
        let r = registry();
        let err = r.call("nope", json!({}), None).await.unwrap_err();
        assert!(matches!(err, Error::ToolNotFound(_)));
    }

    #[tokio::test]
    async fn refusal_surfaces_in_specs_and_errors_on_call() {
        let r = registry();
        let spec = r
            .specs()
            .into_iter()
            .find(|s| s.name == "delete_everything")
            .expect("refusal appears in tools/list");
        assert!(spec.refusal_reason.is_some());

        let err = r
            .call("delete_everything", json!({}), None)
            .await
            .unwrap_err();
        assert!(matches!(err, Error::Refused { .. }));
    }

    #[test]
    fn compact_projection_filters_to_small_model_safe() {
        let r = registry();
        let compact: Vec<_> = r
            .specs_projected(SurfaceProjection::Compact)
            .into_iter()
            .map(|s| s.name)
            .collect();
        assert!(compact.contains(&"echo".to_string()));
        assert!(!compact.contains(&"mutate".to_string()));
    }

    #[test]
    fn write_capabilities_are_deduped() {
        let r = registry();
        let caps = r.write_capabilities();
        assert_eq!(caps.len(), 1);
        assert_eq!(caps[0].id, "test.write");
    }

    #[test]
    #[should_panic(expected = "duplicate tool registration")]
    fn duplicate_registration_panics() {
        let mut r = ToolRegistry::new();
        r.register(Echo);
        r.register(Echo);
    }
}