//! Tool registry: built-ins (compile-time) + per-session dynamic entries.

use crate::tool::{AnyTool, ClientKind, ProtocolKind, ToolContext};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};

/// Where a dynamic tool came from.
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub enum ToolSource {
    Mcp(Arc<str>),
    Custom(Arc<str>),
}

/// A registered dynamic tool with optional client/protocol filters.
#[derive(Clone)]
pub struct DynamicEntry {
    pub tool: Arc<dyn AnyTool>,
    pub source: ToolSource,
    pub only_for_client: Option<ClientKind>,
    pub only_for_protocol: Option<ProtocolKind>,
}

/// Resolution policy when a name exists in both built-ins and dynamic.
#[derive(Clone, Debug, Default)]
pub enum CollisionPolicy {
    #[default]
    BuiltInWins,
    DynamicWins,
    PerName(HashMap<Arc<str>, Winner>),
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum Winner { BuiltIn, Dynamic }

pub struct ToolRegistry {
    built_ins: HashMap<&'static str, Arc<dyn AnyTool>>,
    dynamic: RwLock<HashMap<Arc<str>, HashMap<String, DynamicEntry>>>,
    collision_policy: CollisionPolicy,
}

impl ToolRegistry {
    /// Canonical scope key for dynamic-tool storage.
    ///
    /// Session takes precedence so tools registered for a single session
    /// don't leak across the connector; absent that, the connector itself
    /// scopes the entry.
    fn scope_key(ctx: &ToolContext) -> Arc<str> {
        ctx.session_id
            .clone()
            .unwrap_or_else(|| ctx.connector_id.clone())
    }

    pub fn new(
        built_ins: impl IntoIterator<Item = Arc<dyn AnyTool>>,
        collision_policy: CollisionPolicy,
    ) -> Self {
        let built_ins = built_ins
            .into_iter()
            .map(|t| (t.name(), t))
            .collect();
        Self {
            built_ins,
            dynamic: RwLock::new(HashMap::new()),
            collision_policy,
        }
    }

    pub fn resolve(&self, name: &str, ctx: &ToolContext) -> Option<Arc<dyn AnyTool>> {
        let dyn_match = self.find_dynamic(name, ctx);
        let builtin = self.built_ins.get(name).cloned();
        match (builtin, dyn_match) {
            (Some(b), None)    => Some(b),
            (None,    Some(d)) => Some(d.tool),
            (None,    None)    => None,
            (Some(b), Some(d)) => match &self.collision_policy {
                CollisionPolicy::BuiltInWins => Some(b),
                CollisionPolicy::DynamicWins => Some(d.tool),
                CollisionPolicy::PerName(map) => match map.get(name) {
                    Some(Winner::BuiltIn) | None => Some(b),
                    Some(Winner::Dynamic) => Some(d.tool),
                },
            },
        }
    }

    /// Enumerate all tools visible in `ctx`'s scope.
    ///
    /// Returns built-in tool names plus any dynamic tools registered under
    /// the canonical scope key for `ctx` (see [`Self::scope_key`]) whose
    /// optional client/protocol filters match. The result is sorted and
    /// deduplicated.
    pub fn list(&self, ctx: &ToolContext) -> Vec<String> {
        let mut names: Vec<String> = self.built_ins.keys().map(|n| n.to_string()).collect();
        let key = Self::scope_key(ctx);
        if let Some(per_scope) = self.dynamic.read().unwrap().get(&key) {
            for (n, e) in per_scope.iter() {
                if entry_matches_ctx(e, ctx) { names.push(n.clone()); }
            }
        }
        names.sort();
        names.dedup();
        names
    }

    /// Insert a dynamic tool under `scope_key`.
    ///
    /// Callers must pass the value [`Self::scope_key`] would produce for the
    /// registering [`ToolContext`]: the `session_id` if present, otherwise
    /// the `connector_id`. Anything else will be invisible to `resolve` and
    /// `list` for that context.
    pub fn register_dynamic(&self, scope_key: impl Into<Arc<str>>, name: String, entry: DynamicEntry) {
        let mut g = self.dynamic.write().unwrap();
        g.entry(scope_key.into()).or_default().insert(name, entry);
    }

    pub fn unregister_scope(&self, scope_key: &str) {
        self.dynamic.write().unwrap().remove(scope_key);
    }

    fn find_dynamic(&self, name: &str, ctx: &ToolContext) -> Option<DynamicEntry> {
        let key = Self::scope_key(ctx);
        let g = self.dynamic.read().unwrap();
        let per_scope = g.get(&key)?;
        let entry = per_scope.get(name)?;
        if entry_matches_ctx(entry, ctx) { Some(entry.clone()) } else { None }
    }
}

fn entry_matches_ctx(entry: &DynamicEntry, ctx: &ToolContext) -> bool {
    if let Some(c) = &entry.only_for_client {
        if c != &ctx.client_kind { return false; }
    }
    if let Some(p) = &entry.only_for_protocol {
        if p != &ctx.protocol { return false; }
    }
    true
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{PermissionConfig, SandboxConfig, WhitelistConfig};
    use crate::permission::check::PermissionContext;
    use crate::permission::whitelist::CompiledWhitelist;
    use crate::tool::{Tool, ToolEventSink, ToolInput, ToolKind};
    use async_trait::async_trait;
    use schemars::JsonSchema;
    use serde::{Deserialize, Serialize};
    use std::path::PathBuf;

    #[derive(Default)] struct A;
    #[derive(Default)] struct B;

    #[derive(Serialize, Deserialize, JsonSchema)]
    struct Empty {}

    macro_rules! impl_t {
        ($t:ty, $n:literal) => {
            #[async_trait]
            impl Tool for $t {
                type Input = Empty;
                type Output = Empty;
                const NAME: &'static str = $n;
                fn kind() -> ToolKind { ToolKind::Other }
                async fn run(
                    self: Arc<Self>, _i: ToolInput<Empty>,
                    _e: ToolEventSink, _c: &ToolContext,
                ) -> Result<Empty, Empty> { Ok(Empty {}) }
            }
        };
    }
    impl_t!(A, "a");
    impl_t!(B, "b");

    fn ctx() -> ToolContext {
        let wl = CompiledWhitelist::compile(&WhitelistConfig::default()).unwrap();
        let pc = PermissionContext::new("conn-1".to_string(), None, wl);
        ToolContext::for_test(
            "conn-1", ClientKind::claude(), ProtocolKind::acp(),
            PathBuf::from("/tmp"),
            SandboxConfig::default(), PermissionConfig::default(), pc,
        )
    }

    fn arcs() -> Vec<Arc<dyn AnyTool>> {
        vec![
            <A as Tool>::erase(Arc::new(A)),
            <B as Tool>::erase(Arc::new(B)),
        ]
    }

    #[test]
    fn resolves_built_in_by_name() {
        let r = ToolRegistry::new(arcs(), CollisionPolicy::BuiltInWins);
        assert!(r.resolve("a", &ctx()).is_some());
        assert!(r.resolve("nope", &ctx()).is_none());
    }

    #[test]
    fn list_includes_built_ins() {
        let r = ToolRegistry::new(arcs(), CollisionPolicy::BuiltInWins);
        let mut names: Vec<String> = r.list(&ctx()).into_iter().collect();
        names.sort();
        assert_eq!(names, vec!["a".to_string(), "b".to_string()]);
    }

    fn dynamic_entry_named(name: &str, only_client: Option<ClientKind>) -> DynamicEntry {
        let _ = name;
        DynamicEntry {
            tool: <A as Tool>::erase(Arc::new(A)),
            source: ToolSource::Mcp(Arc::from("server-1")),
            only_for_client: only_client,
            only_for_protocol: None,
        }
    }

    #[test]
    fn dynamic_resolves_when_no_built_in() {
        let r = ToolRegistry::new(vec![], CollisionPolicy::BuiltInWins);
        r.register_dynamic("conn-1", "extra".to_string(), dynamic_entry_named("extra", None));
        assert!(r.resolve("extra", &ctx()).is_some());
    }

    #[test]
    fn collision_default_built_in_wins() {
        let r = ToolRegistry::new(arcs(), CollisionPolicy::BuiltInWins);
        r.register_dynamic("conn-1", "a".to_string(), dynamic_entry_named("a", None));
        // Both define "a"; built-in wins.
        let resolved = r.resolve("a", &ctx()).unwrap();
        assert_eq!(resolved.name(), "a");
    }

    #[test]
    fn collision_dynamic_wins_when_configured() {
        let r = ToolRegistry::new(arcs(), CollisionPolicy::DynamicWins);
        r.register_dynamic("conn-1", "a".to_string(), dynamic_entry_named("a", None));
        let resolved = r.resolve("a", &ctx()).unwrap();
        assert_eq!(resolved.name(), "a");
    }

    #[test]
    fn dynamic_filters_by_client_kind() {
        let r = ToolRegistry::new(vec![], CollisionPolicy::BuiltInWins);
        r.register_dynamic(
            "conn-1",
            "extra".to_string(),
            dynamic_entry_named("extra", Some(ClientKind::codex())),
        );
        // ctx() has client_kind = Claude — should NOT see this tool.
        assert!(r.resolve("extra", &ctx()).is_none());
    }

    fn ctx_with_session(session: &str) -> ToolContext {
        let wl = CompiledWhitelist::compile(&WhitelistConfig::default()).unwrap();
        let pc = PermissionContext::new("conn-1".to_string(), Some(session.to_string()), wl);
        let mut c = ToolContext::for_test(
            "conn-1", ClientKind::claude(), ProtocolKind::acp(),
            PathBuf::from("/tmp"),
            SandboxConfig::default(), PermissionConfig::default(), pc,
        );
        c.session_id = Some(Arc::from(session));
        c
    }

    #[test]
    fn register_under_session_visible_to_resolve_and_list() {
        let r = ToolRegistry::new(vec![], CollisionPolicy::BuiltInWins);
        let c = ctx_with_session("sess-A");
        r.register_dynamic(
            ToolRegistry::scope_key(&c),
            "extra".to_string(),
            dynamic_entry_named("extra", None),
        );
        assert!(r.resolve("extra", &c).is_some());
        assert!(r.list(&c).iter().any(|n| n == "extra"));
        // A different session under the same connector must NOT see it.
        let other = ctx_with_session("sess-B");
        assert!(r.resolve("extra", &other).is_none());
        assert!(!r.list(&other).iter().any(|n| n == "extra"));
    }

    #[test]
    fn register_under_connector_only_visible_to_resolve_and_list() {
        let r = ToolRegistry::new(vec![], CollisionPolicy::BuiltInWins);
        let c = ctx(); // session_id = None → scope_key falls back to connector_id
        r.register_dynamic(
            ToolRegistry::scope_key(&c),
            "extra".to_string(),
            dynamic_entry_named("extra", None),
        );
        assert!(r.resolve("extra", &c).is_some());
        assert!(r.list(&c).iter().any(|n| n == "extra"));
    }

    #[test]
    fn unregister_scope_drops_dynamic() {
        let r = ToolRegistry::new(vec![], CollisionPolicy::BuiltInWins);
        r.register_dynamic("conn-1", "extra".to_string(), dynamic_entry_named("extra", None));
        assert!(r.resolve("extra", &ctx()).is_some());
        r.unregister_scope("conn-1");
        assert!(r.resolve("extra", &ctx()).is_none());
    }
}
