agent-sdk 0.8.0

//! Tool definition and registry.
//!
//! Tools allow the LLM to perform actions in the real world. This module provides:
//!
//! - [`Tool`] trait - Define custom tools the LLM can call
//! - [`ToolName`] trait - Marker trait for strongly-typed tool names
//! - [`PrimitiveToolName`] - Tool names for SDK's built-in tools
//! - [`DynamicToolName`] - Tool names created at runtime (MCP bridges)
//! - [`ToolRegistry`] - Collection of available tools
//! - [`ToolContext`] - Context passed to tool execution
//! - [`ListenExecuteTool`] - Tools that listen for updates, then execute later
//!
//! # Implementing a Tool
//!
//! ```ignore
//! use agent_sdk::{Tool, ToolContext, ToolResult, ToolTier, PrimitiveToolName};
//!
//! struct MyTool;
//!
//! // No #[async_trait] needed - Rust 1.75+ supports native async traits
//! impl Tool<MyContext> for MyTool {
//!     type Name = PrimitiveToolName;
//!
//!     fn name(&self) -> PrimitiveToolName { PrimitiveToolName::Read }
//!     fn display_name(&self) -> &'static str { "My Tool" }
//!     fn description(&self) -> &'static str { "Does something useful" }
//!     fn input_schema(&self) -> Value { json!({ "type": "object" }) }
//!     fn tier(&self) -> ToolTier { ToolTier::Observe }
//!
//!     async fn execute(&self, ctx: &ToolContext<MyContext>, input: Value) -> Result<ToolResult> {
//!         Ok(ToolResult::success("Done!"))
//!     }
//! }
//! ```

use crate::events::{AgentEvent, AgentEventEnvelope, SequenceCounter};
use crate::llm;
use crate::types::{ToolOutcome, ToolResult, ToolTier};
use anyhow::Result;
use async_trait::async_trait;
use futures::Stream;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;
use std::collections::HashMap;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;
use time::OffsetDateTime;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

// ============================================================================
// Tool Name Types
// ============================================================================

/// Marker trait for tool names.
///
/// Tool names must be serializable (for storage/logging) and deserializable
/// (for parsing from LLM responses). The string representation is derived
/// from serde serialization.
///
/// # Example
///
/// ```ignore
/// #[derive(Serialize, Deserialize)]
/// #[serde(rename_all = "snake_case")]
/// pub enum MyToolName {
///     Read,
///     Write,
/// }
///
/// impl ToolName for MyToolName {}
/// ```
pub trait ToolName: Send + Sync + Serialize + DeserializeOwned + 'static {}

/// Helper to get string representation of a tool name via serde.
///
/// Returns `"<unknown_tool>"` if serialization fails (should never happen
/// with properly implemented `ToolName` types that use `#[derive(Serialize)]`).
#[must_use]
pub fn tool_name_to_string<N: ToolName>(name: &N) -> String {
    serde_json::to_string(name)
        .unwrap_or_else(|_| "\"<unknown_tool>\"".to_string())
        .trim_matches('"')
        .to_string()
}

/// Parse a tool name from string via serde.
///
/// # Errors
/// Returns error if the string doesn't match a valid tool name.
pub fn tool_name_from_str<N: ToolName>(s: &str) -> Result<N, serde_json::Error> {
    serde_json::from_str(&format!("\"{s}\""))
}

/// Tool names for SDK's built-in primitive tools.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PrimitiveToolName {
    Read,
    Write,
    Edit,
    MultiEdit,
    Bash,
    Glob,
    Grep,
    NotebookRead,
    NotebookEdit,
    TodoRead,
    TodoWrite,
    AskUser,
    LinkFetch,
    WebSearch,
}

impl ToolName for PrimitiveToolName {}

/// Dynamic tool name for runtime-created tools (MCP bridges, subagents).
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct DynamicToolName(String);

impl DynamicToolName {
    #[must_use]
    pub fn new(name: impl Into<String>) -> Self {
        Self(name.into())
    }

    #[must_use]
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl ToolName for DynamicToolName {}

// ============================================================================
// Progress Stage Types (for AsyncTool)
// ============================================================================

/// Marker trait for tool progress stages (type-safe, like [`ToolName`]).
///
/// Progress stages are used by async tools to indicate the current phase
/// of a long-running operation. They must be serializable for event streaming.
///
/// # Example
///
/// ```ignore
/// #[derive(Clone, Debug, Serialize, Deserialize)]
/// #[serde(rename_all = "snake_case")]
/// pub enum PixTransferStage {
///     Initiated,
///     Processing,
///     SentToBank,
/// }
///
/// impl ProgressStage for PixTransferStage {}
/// ```
pub trait ProgressStage: Clone + Send + Sync + Serialize + DeserializeOwned + 'static {}

/// Helper to get string representation of a progress stage via serde.
///
/// # Panics
///
/// Panics if the stage cannot be serialized to a string. This should
/// never happen with properly implemented `ProgressStage` types.
#[must_use]
pub fn stage_to_string<S: ProgressStage>(stage: &S) -> String {
    serde_json::to_string(stage)
        .expect("ProgressStage must serialize to string")
        .trim_matches('"')
        .to_string()
}

/// Status update from an async tool operation.
#[derive(Clone, Debug, Serialize)]
pub enum ToolStatus<S: ProgressStage> {
    /// Operation is making progress
    Progress {
        stage: S,
        message: String,
        data: Option<serde_json::Value>,
    },

    /// Operation completed successfully
    Completed(ToolResult),

    /// Operation failed
    Failed(ToolResult),
}

/// Type-erased status for the agent loop.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ErasedToolStatus {
    /// Operation is making progress
    Progress {
        stage: String,
        message: String,
        data: Option<serde_json::Value>,
    },
    /// Operation completed successfully
    Completed(ToolResult),
    /// Operation failed
    Failed(ToolResult),
}

/// Update emitted from a `listen()` stream.
///
/// This models workflows where a runtime prepares an operation over time, and
/// execution happens later using an operation identifier and revision.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ListenToolUpdate {
    /// Preparation is still running and should keep listening.
    Listening {
        /// Opaque operation identifier used for later execute/cancel calls.
        operation_id: String,
        /// Monotonic revision number for optimistic concurrency.
        revision: u64,
        /// Human-readable status message.
        message: String,
        /// Optional current snapshot for UI rendering.
        snapshot: Option<serde_json::Value>,
        /// Optional expiration timestamp (RFC3339).
        #[serde(with = "time::serde::rfc3339::option")]
        expires_at: Option<OffsetDateTime>,
    },

    /// Preparation is complete and execution can be confirmed.
    Ready {
        /// Opaque operation identifier used for later execute/cancel calls.
        operation_id: String,
        /// Monotonic revision number for optimistic concurrency.
        revision: u64,
        /// Human-readable status message.
        message: String,
        /// Snapshot shown in confirmation UI.
        snapshot: serde_json::Value,
        /// Optional expiration timestamp (RFC3339).
        #[serde(with = "time::serde::rfc3339::option")]
        expires_at: Option<OffsetDateTime>,
    },

    /// Operation is no longer valid.
    Invalidated {
        /// Opaque operation identifier.
        operation_id: String,
        /// Human-readable reason.
        message: String,
        /// Whether caller may recover by starting a new listen operation.
        recoverable: bool,
    },
}

/// Reason for stopping a listen session.
#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
pub enum ListenStopReason {
    /// User explicitly rejected confirmation.
    UserRejected,
    /// Agent policy/hook blocked execution before confirmation.
    Blocked,
    /// Consumer disconnected while listen stream was active.
    StreamDisconnected,
    /// Listen stream ended unexpectedly before terminal state.
    StreamEnded,
}

impl<S: ProgressStage> From<ToolStatus<S>> for ErasedToolStatus {
    fn from(status: ToolStatus<S>) -> Self {
        match status {
            ToolStatus::Progress {
                stage,
                message,
                data,
            } => Self::Progress {
                stage: stage_to_string(&stage),
                message,
                data,
            },
            ToolStatus::Completed(r) => Self::Completed(r),
            ToolStatus::Failed(r) => Self::Failed(r),
        }
    }
}

/// Context passed to tool execution
pub struct ToolContext<Ctx> {
    /// Application-specific context (e.g., `user_id`, db connection)
    pub app: Ctx,
    /// Tool-specific metadata
    pub metadata: HashMap<String, Value>,
    /// Optional channel for tools to emit events (e.g., subagent progress)
    event_tx: Option<mpsc::Sender<AgentEventEnvelope>>,
    /// Optional sequence counter for wrapping events in envelopes
    event_seq: Option<SequenceCounter>,
    /// Optional cancellation token for propagating cancellation to subtasks
    cancel_token: Option<CancellationToken>,
}

impl<Ctx> ToolContext<Ctx> {
    #[must_use]
    pub fn new(app: Ctx) -> Self {
        Self {
            app,
            metadata: HashMap::new(),
            event_tx: None,
            event_seq: None,
            cancel_token: None,
        }
    }

    #[must_use]
    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
        self.metadata.insert(key.into(), value);
        self
    }

    /// Set the event channel and sequence counter for tools that need to emit
    /// events during execution.
    #[must_use]
    pub fn with_event_tx(
        mut self,
        tx: mpsc::Sender<AgentEventEnvelope>,
        seq: SequenceCounter,
    ) -> Self {
        self.event_tx = Some(tx);
        self.event_seq = Some(seq);
        self
    }

    /// Emit an event through the event channel (if set).
    ///
    /// The event is wrapped in an [`AgentEventEnvelope`] with a unique ID,
    /// sequence number, and timestamp before sending.
    ///
    /// This uses `try_send` to avoid blocking and to ensure the future is `Send`.
    /// The event is silently dropped if the channel is full.
    pub fn emit_event(&self, event: AgentEvent) {
        if let Some((tx, seq)) = self.event_tx.as_ref().zip(self.event_seq.as_ref()) {
            let envelope = AgentEventEnvelope::wrap(event, seq);
            let _ = tx.try_send(envelope);
        }
    }

    /// Get a clone of the event channel sender (if set).
    ///
    /// This is useful for tools that spawn subprocesses (like subagents)
    /// and need to forward events to the parent's event stream.
    #[must_use]
    pub fn event_tx(&self) -> Option<mpsc::Sender<AgentEventEnvelope>> {
        self.event_tx.clone()
    }

    /// Get a clone of the sequence counter (if set).
    ///
    /// This is useful for tools that spawn subprocesses (like subagents)
    /// and need to assign sequence numbers to events sent to the parent's stream.
    #[must_use]
    pub fn event_seq(&self) -> Option<SequenceCounter> {
        self.event_seq.clone()
    }

    /// Set the cancellation token for propagating cancellation to subtasks.
    #[must_use]
    pub fn with_cancel_token(mut self, token: CancellationToken) -> Self {
        self.cancel_token = Some(token);
        self
    }

    /// Get the cancellation token (if set).
    ///
    /// Used by tools that spawn long-running subtasks (like subagents)
    /// to propagate cancellation from the parent.
    #[must_use]
    pub fn cancel_token(&self) -> Option<CancellationToken> {
        self.cancel_token.clone()
    }
}

// ============================================================================
// Tool Trait
// ============================================================================

/// Definition of a tool that can be called by the agent.
///
/// Tools have a strongly-typed `Name` associated type that determines
/// how the tool name is serialized for LLM communication.
///
/// # Native Async Support
///
/// This trait uses Rust's native async functions in traits (stabilized in Rust 1.75).
/// You do NOT need the `async_trait` crate to implement this trait.
pub trait Tool<Ctx>: Send + Sync {
    /// The type of name for this tool.
    type Name: ToolName;

    /// Returns the tool's strongly-typed name.
    fn name(&self) -> Self::Name;

    /// Human-readable display name for UI (e.g., "Read File" vs "read").
    ///
    /// Defaults to empty string. Override for better UX.
    fn display_name(&self) -> &'static str;

    /// Human-readable description of what the tool does.
    fn description(&self) -> &'static str;

    /// JSON schema for the tool's input parameters.
    fn input_schema(&self) -> Value;

    /// Permission tier for this tool.
    fn tier(&self) -> ToolTier {
        ToolTier::Observe
    }

    /// Execute the tool with the given input.
    ///
    /// # Errors
    /// Returns an error if tool execution fails.
    fn execute(
        &self,
        ctx: &ToolContext<Ctx>,
        input: Value,
    ) -> impl Future<Output = Result<ToolResult>> + Send;
}

// ============================================================================
// AsyncTool Trait
// ============================================================================

/// A tool that performs long-running async operations.
///
/// `AsyncTool`s have two phases:
/// 1. `execute()` - Start the operation (lightweight, returns quickly)
/// 2. `check_status()` - Stream progress until completion
///
/// The actual work should happen externally (background task, external service)
/// and persist results to a durable store. The tool is just an orchestrator.
///
/// # Example
///
/// ```ignore
/// impl AsyncTool<MyCtx> for ExecutePixTransferTool {
///     type Name = PixToolName;
///     type Stage = PixTransferStage;
///
///     async fn execute(&self, ctx: &ToolContext<MyCtx>, input: Value) -> Result<ToolOutcome> {
///         let params = parse_input(&input)?;
///         let operation_id = ctx.app.pix_service.start_transfer(params).await?;
///         Ok(ToolOutcome::in_progress(
///             operation_id,
///             format!("PIX transfer of {} initiated", params.amount),
///         ))
///     }
///
///     fn check_status(&self, ctx: &ToolContext<MyCtx>, operation_id: &str)
///         -> impl Stream<Item = ToolStatus<PixTransferStage>> + Send
///     {
///         async_stream::stream! {
///             loop {
///                 let status = ctx.app.pix_service.get_status(operation_id).await;
///                 match status {
///                     PixStatus::Success { id } => {
///                         yield ToolStatus::Completed(ToolResult::success(id));
///                         break;
///                     }
///                     _ => yield ToolStatus::Progress { ... };
///                 }
///                 tokio::time::sleep(Duration::from_millis(500)).await;
///             }
///         }
///     }
/// }
/// ```
pub trait AsyncTool<Ctx>: Send + Sync {
    /// The type of name for this tool.
    type Name: ToolName;
    /// The type of progress stages for this tool.
    type Stage: ProgressStage;

    /// Returns the tool's strongly-typed name.
    fn name(&self) -> Self::Name;

    /// Human-readable display name for UI.
    fn display_name(&self) -> &'static str;

    /// Human-readable description of what the tool does.
    fn description(&self) -> &'static str;

    /// JSON schema for the tool's input parameters.
    fn input_schema(&self) -> Value;

    /// Permission tier for this tool.
    fn tier(&self) -> ToolTier {
        ToolTier::Observe
    }

    /// Execute the tool. Returns immediately with one of:
    /// - Success/Failed: Operation completed synchronously
    /// - `InProgress`: Operation started, use `check_status()` to stream updates
    ///
    /// # Errors
    /// Returns an error if tool execution fails.
    fn execute(
        &self,
        ctx: &ToolContext<Ctx>,
        input: Value,
    ) -> impl Future<Output = Result<ToolOutcome>> + Send;

    /// Stream status updates for an in-progress operation.
    /// Must yield until Completed or Failed.
    fn check_status(
        &self,
        ctx: &ToolContext<Ctx>,
        operation_id: &str,
    ) -> impl Stream<Item = ToolStatus<Self::Stage>> + Send;
}

// ============================================================================
// ListenExecuteTool Trait
// ============================================================================

/// A tool whose runtime has two phases:
/// 1. `listen()` - starts preparation and streams updates
/// 2. `execute()` - performs final execution after confirmation
///
/// This abstraction is useful when runtime state can expire or evolve before
/// execution (quotes, challenge windows, leases, approvals).
///
/// Ordering note: the agent loop consumes `listen()` updates before
/// `AgentHooks::pre_tool_use()` runs. Hooks can therefore block `execute()`, but
/// any side effects done during `listen()` have already happened.
pub trait ListenExecuteTool<Ctx>: Send + Sync {
    /// The type of name for this tool.
    type Name: ToolName;

    /// Returns the tool's strongly-typed name.
    fn name(&self) -> Self::Name;

    /// Human-readable display name for UI.
    fn display_name(&self) -> &'static str;

    /// Human-readable description of what the tool does.
    fn description(&self) -> &'static str;

    /// JSON schema for the tool's input parameters.
    fn input_schema(&self) -> Value;

    /// Permission tier for this tool.
    fn tier(&self) -> ToolTier {
        ToolTier::Confirm
    }

    /// Start and stream runtime preparation updates.
    fn listen(
        &self,
        ctx: &ToolContext<Ctx>,
        input: Value,
    ) -> impl Stream<Item = ListenToolUpdate> + Send;

    /// Execute using operation ID and optimistic concurrency revision.
    ///
    /// # Errors
    /// Returns an error if execution fails or revision is stale.
    fn execute(
        &self,
        ctx: &ToolContext<Ctx>,
        operation_id: &str,
        expected_revision: u64,
    ) -> impl Future<Output = Result<ToolResult>> + Send;

    /// Stop a listen operation (best effort).
    ///
    /// # Errors
    /// Returns an error if cancellation fails.
    fn cancel(
        &self,
        _ctx: &ToolContext<Ctx>,
        _operation_id: &str,
        _reason: ListenStopReason,
    ) -> impl Future<Output = Result<()>> + Send {
        async { Ok(()) }
    }
}

// ============================================================================
// Type-Erased Tool (for Registry)
// ============================================================================

/// Type-erased tool trait for registry storage.
///
/// This allows tools with different `Name` associated types to be stored
/// in the same registry by erasing the type information.
///
/// # Example
///
/// ```ignore
/// for tool in registry.all() {
///     println!("Tool: {} - {}", tool.name_str(), tool.description());
/// }
/// ```
#[async_trait]
pub trait ErasedTool<Ctx>: Send + Sync {
    /// Get the tool name as a string.
    fn name_str(&self) -> &str;
    /// Get a human-friendly display name for the tool.
    fn display_name(&self) -> &'static str;
    /// Get the tool description.
    fn description(&self) -> &'static str;
    /// Get the JSON schema for tool inputs.
    fn input_schema(&self) -> Value;
    /// Get the tool's permission tier.
    fn tier(&self) -> ToolTier;
    /// Execute the tool with the given input.
    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolResult>;
}

/// Wrapper that erases the Name associated type from a Tool.
struct ToolWrapper<T, Ctx>
where
    T: Tool<Ctx>,
{
    inner: T,
    name_cache: String,
    _marker: PhantomData<Ctx>,
}

impl<T, Ctx> ToolWrapper<T, Ctx>
where
    T: Tool<Ctx>,
{
    fn new(tool: T) -> Self {
        let name_cache = tool_name_to_string(&tool.name());
        Self {
            inner: tool,
            name_cache,
            _marker: PhantomData,
        }
    }
}

#[async_trait]
impl<T, Ctx> ErasedTool<Ctx> for ToolWrapper<T, Ctx>
where
    T: Tool<Ctx> + 'static,
    Ctx: Send + Sync + 'static,
{
    fn name_str(&self) -> &str {
        &self.name_cache
    }

    fn display_name(&self) -> &'static str {
        self.inner.display_name()
    }

    fn description(&self) -> &'static str {
        self.inner.description()
    }

    fn input_schema(&self) -> Value {
        self.inner.input_schema()
    }

    fn tier(&self) -> ToolTier {
        self.inner.tier()
    }

    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolResult> {
        self.inner.execute(ctx, input).await
    }
}

// ============================================================================
// Type-Erased AsyncTool (for Registry)
// ============================================================================

/// Type-erased async tool trait for registry storage.
///
/// This allows async tools with different `Name` and `Stage` associated types
/// to be stored in the same registry by erasing the type information.
#[async_trait]
pub trait ErasedAsyncTool<Ctx>: Send + Sync {
    /// Get the tool name as a string.
    fn name_str(&self) -> &str;
    /// Get a human-friendly display name for the tool.
    fn display_name(&self) -> &'static str;
    /// Get the tool description.
    fn description(&self) -> &'static str;
    /// Get the JSON schema for tool inputs.
    fn input_schema(&self) -> Value;
    /// Get the tool's permission tier.
    fn tier(&self) -> ToolTier;
    /// Execute the tool with the given input.
    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolOutcome>;
    /// Stream status updates for an in-progress operation (type-erased).
    fn check_status_stream<'a>(
        &'a self,
        ctx: &'a ToolContext<Ctx>,
        operation_id: &'a str,
    ) -> Pin<Box<dyn Stream<Item = ErasedToolStatus> + Send + 'a>>;
}

/// Wrapper that erases the Name and Stage associated types from an [`AsyncTool`].
struct AsyncToolWrapper<T, Ctx>
where
    T: AsyncTool<Ctx>,
{
    inner: T,
    name_cache: String,
    _marker: PhantomData<Ctx>,
}

impl<T, Ctx> AsyncToolWrapper<T, Ctx>
where
    T: AsyncTool<Ctx>,
{
    fn new(tool: T) -> Self {
        let name_cache = tool_name_to_string(&tool.name());
        Self {
            inner: tool,
            name_cache,
            _marker: PhantomData,
        }
    }
}

#[async_trait]
impl<T, Ctx> ErasedAsyncTool<Ctx> for AsyncToolWrapper<T, Ctx>
where
    T: AsyncTool<Ctx> + 'static,
    Ctx: Send + Sync + 'static,
{
    fn name_str(&self) -> &str {
        &self.name_cache
    }

    fn display_name(&self) -> &'static str {
        self.inner.display_name()
    }

    fn description(&self) -> &'static str {
        self.inner.description()
    }

    fn input_schema(&self) -> Value {
        self.inner.input_schema()
    }

    fn tier(&self) -> ToolTier {
        self.inner.tier()
    }

    async fn execute(&self, ctx: &ToolContext<Ctx>, input: Value) -> Result<ToolOutcome> {
        self.inner.execute(ctx, input).await
    }

    fn check_status_stream<'a>(
        &'a self,
        ctx: &'a ToolContext<Ctx>,
        operation_id: &'a str,
    ) -> Pin<Box<dyn Stream<Item = ErasedToolStatus> + Send + 'a>> {
        use futures::StreamExt;
        let stream = self.inner.check_status(ctx, operation_id);
        Box::pin(stream.map(ErasedToolStatus::from))
    }
}

// ============================================================================
// Type-Erased ListenExecuteTool (for Registry)
// ============================================================================

/// Type-erased listen/execute tool trait for registry storage.
#[async_trait]
pub trait ErasedListenTool<Ctx>: Send + Sync {
    /// Get the tool name as a string.
    fn name_str(&self) -> &str;
    /// Get a human-friendly display name for the tool.
    fn display_name(&self) -> &'static str;
    /// Get the tool description.
    fn description(&self) -> &'static str;
    /// Get the JSON schema for tool inputs.
    fn input_schema(&self) -> Value;
    /// Get the tool's permission tier.
    fn tier(&self) -> ToolTier;
    /// Start listen stream.
    fn listen_stream<'a>(
        &'a self,
        ctx: &'a ToolContext<Ctx>,
        input: Value,
    ) -> Pin<Box<dyn Stream<Item = ListenToolUpdate> + Send + 'a>>;
    /// Execute using a prepared operation.
    async fn execute(
        &self,
        ctx: &ToolContext<Ctx>,
        operation_id: &str,
        expected_revision: u64,
    ) -> Result<ToolResult>;
    /// Cancel operation.
    async fn cancel(
        &self,
        ctx: &ToolContext<Ctx>,
        operation_id: &str,
        reason: ListenStopReason,
    ) -> Result<()>;
}

/// Wrapper that erases the Name associated type from a [`ListenExecuteTool`].
struct ListenToolWrapper<T, Ctx>
where
    T: ListenExecuteTool<Ctx>,
{
    inner: T,
    name_cache: String,
    _marker: PhantomData<Ctx>,
}

impl<T, Ctx> ListenToolWrapper<T, Ctx>
where
    T: ListenExecuteTool<Ctx>,
{
    fn new(tool: T) -> Self {
        let name_cache = tool_name_to_string(&tool.name());
        Self {
            inner: tool,
            name_cache,
            _marker: PhantomData,
        }
    }
}

#[async_trait]
impl<T, Ctx> ErasedListenTool<Ctx> for ListenToolWrapper<T, Ctx>
where
    T: ListenExecuteTool<Ctx> + 'static,
    Ctx: Send + Sync + 'static,
{
    fn name_str(&self) -> &str {
        &self.name_cache
    }

    fn display_name(&self) -> &'static str {
        self.inner.display_name()
    }

    fn description(&self) -> &'static str {
        self.inner.description()
    }

    fn input_schema(&self) -> Value {
        self.inner.input_schema()
    }

    fn tier(&self) -> ToolTier {
        self.inner.tier()
    }

    fn listen_stream<'a>(
        &'a self,
        ctx: &'a ToolContext<Ctx>,
        input: Value,
    ) -> Pin<Box<dyn Stream<Item = ListenToolUpdate> + Send + 'a>> {
        let stream = self.inner.listen(ctx, input);
        Box::pin(stream)
    }

    async fn execute(
        &self,
        ctx: &ToolContext<Ctx>,
        operation_id: &str,
        expected_revision: u64,
    ) -> Result<ToolResult> {
        self.inner
            .execute(ctx, operation_id, expected_revision)
            .await
    }

    async fn cancel(
        &self,
        ctx: &ToolContext<Ctx>,
        operation_id: &str,
        reason: ListenStopReason,
    ) -> Result<()> {
        self.inner.cancel(ctx, operation_id, reason).await
    }
}

/// Registry of available tools.
///
/// Tools are stored with their names erased to allow different `Name` types
/// in the same registry. The registry uses string-based lookup for LLM
/// compatibility.
///
/// Supports both synchronous [`Tool`]s and asynchronous [`AsyncTool`]s.
pub struct ToolRegistry<Ctx> {
    tools: HashMap<String, Arc<dyn ErasedTool<Ctx>>>,
    async_tools: HashMap<String, Arc<dyn ErasedAsyncTool<Ctx>>>,
    listen_tools: HashMap<String, Arc<dyn ErasedListenTool<Ctx>>>,
}

impl<Ctx> Clone for ToolRegistry<Ctx> {
    fn clone(&self) -> Self {
        Self {
            tools: self.tools.clone(),
            async_tools: self.async_tools.clone(),
            listen_tools: self.listen_tools.clone(),
        }
    }
}

impl<Ctx: Send + Sync + 'static> Default for ToolRegistry<Ctx> {
    fn default() -> Self {
        Self::new()
    }
}

impl<Ctx: Send + Sync + 'static> ToolRegistry<Ctx> {
    #[must_use]
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
            async_tools: HashMap::new(),
            listen_tools: HashMap::new(),
        }
    }

    /// Register a synchronous tool in the registry.
    ///
    /// The tool's name is converted to a string via serde serialization
    /// and used as the lookup key.
    pub fn register<T>(&mut self, tool: T) -> &mut Self
    where
        T: Tool<Ctx> + 'static,
    {
        let wrapper = ToolWrapper::new(tool);
        let name = wrapper.name_str().to_string();
        self.tools.insert(name, Arc::new(wrapper));
        self
    }

    /// Register an async tool in the registry.
    ///
    /// Async tools have two phases: execute (lightweight, starts operation)
    /// and `check_status` (streams progress until completion).
    pub fn register_async<T>(&mut self, tool: T) -> &mut Self
    where
        T: AsyncTool<Ctx> + 'static,
    {
        let wrapper = AsyncToolWrapper::new(tool);
        let name = wrapper.name_str().to_string();
        self.async_tools.insert(name, Arc::new(wrapper));
        self
    }

    /// Register a listen/execute tool in the registry.
    ///
    /// Listen/execute tools start by streaming updates via `listen()`, then run
    /// final execution with `execute()` once confirmed.
    pub fn register_listen<T>(&mut self, tool: T) -> &mut Self
    where
        T: ListenExecuteTool<Ctx> + 'static,
    {
        let wrapper = ListenToolWrapper::new(tool);
        let name = wrapper.name_str().to_string();
        self.listen_tools.insert(name, Arc::new(wrapper));
        self
    }

    /// Get a synchronous tool by name.
    #[must_use]
    pub fn get(&self, name: &str) -> Option<&Arc<dyn ErasedTool<Ctx>>> {
        self.tools.get(name)
    }

    /// Get an async tool by name.
    #[must_use]
    pub fn get_async(&self, name: &str) -> Option<&Arc<dyn ErasedAsyncTool<Ctx>>> {
        self.async_tools.get(name)
    }

    /// Get a listen/execute tool by name.
    #[must_use]
    pub fn get_listen(&self, name: &str) -> Option<&Arc<dyn ErasedListenTool<Ctx>>> {
        self.listen_tools.get(name)
    }

    /// Check if a tool name refers to an async tool.
    #[must_use]
    pub fn is_async(&self, name: &str) -> bool {
        self.async_tools.contains_key(name)
    }

    /// Check if a tool name refers to a listen/execute tool.
    #[must_use]
    pub fn is_listen(&self, name: &str) -> bool {
        self.listen_tools.contains_key(name)
    }

    /// Get all registered synchronous tools.
    pub fn all(&self) -> impl Iterator<Item = &Arc<dyn ErasedTool<Ctx>>> {
        self.tools.values()
    }

    /// Get all registered async tools.
    pub fn all_async(&self) -> impl Iterator<Item = &Arc<dyn ErasedAsyncTool<Ctx>>> {
        self.async_tools.values()
    }

    /// Get all registered listen/execute tools.
    pub fn all_listen(&self) -> impl Iterator<Item = &Arc<dyn ErasedListenTool<Ctx>>> {
        self.listen_tools.values()
    }

    /// Get the number of registered tools (sync + async).
    #[must_use]
    pub fn len(&self) -> usize {
        self.tools.len() + self.async_tools.len() + self.listen_tools.len()
    }

    /// Check if the registry is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.tools.is_empty() && self.async_tools.is_empty() && self.listen_tools.is_empty()
    }

    /// Filter tools by a predicate.
    ///
    /// Removes tools for which the predicate returns false.
    /// The predicate receives the tool name.
    /// Applies to both sync and async tools.
    ///
    /// # Example
    ///
    /// ```ignore
    /// registry.filter(|name| name != "bash");
    /// ```
    pub fn filter<F>(&mut self, predicate: F)
    where
        F: Fn(&str) -> bool,
    {
        self.tools.retain(|name, _| predicate(name));
        self.async_tools.retain(|name, _| predicate(name));
        self.listen_tools.retain(|name, _| predicate(name));
    }

    /// Convert all tools (sync + async) to LLM tool definitions.
    #[must_use]
    pub fn to_llm_tools(&self) -> Vec<llm::Tool> {
        let mut tools: Vec<_> = self
            .tools
            .values()
            .map(|tool| llm::Tool {
                name: tool.name_str().to_string(),
                description: tool.description().to_string(),
                input_schema: tool.input_schema(),
            })
            .collect();

        tools.extend(self.async_tools.values().map(|tool| llm::Tool {
            name: tool.name_str().to_string(),
            description: tool.description().to_string(),
            input_schema: tool.input_schema(),
        }));

        tools.extend(self.listen_tools.values().map(|tool| llm::Tool {
            name: tool.name_str().to_string(),
            description: tool.description().to_string(),
            input_schema: tool.input_schema(),
        }));

        tools
    }
}

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

    // Test tool name enum for tests
    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
    #[serde(rename_all = "snake_case")]
    enum TestToolName {
        MockTool,
        AnotherTool,
    }

    impl ToolName for TestToolName {}

    struct MockTool;

    impl Tool<()> for MockTool {
        type Name = TestToolName;

        fn name(&self) -> TestToolName {
            TestToolName::MockTool
        }

        fn display_name(&self) -> &'static str {
            "Mock Tool"
        }

        fn description(&self) -> &'static str {
            "A mock tool for testing"
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({
                "type": "object",
                "properties": {
                    "message": { "type": "string" }
                }
            })
        }

        async fn execute(&self, _ctx: &ToolContext<()>, input: Value) -> Result<ToolResult> {
            let message = input
                .get("message")
                .and_then(|v| v.as_str())
                .unwrap_or("no message");
            Ok(ToolResult::success(format!("Received: {message}")))
        }
    }

    #[test]
    fn test_tool_name_serialization() {
        let name = TestToolName::MockTool;
        assert_eq!(tool_name_to_string(&name), "mock_tool");

        let parsed: TestToolName = tool_name_from_str("mock_tool").unwrap();
        assert_eq!(parsed, TestToolName::MockTool);
    }

    #[test]
    fn test_dynamic_tool_name() {
        let name = DynamicToolName::new("my_mcp_tool");
        assert_eq!(tool_name_to_string(&name), "my_mcp_tool");
        assert_eq!(name.as_str(), "my_mcp_tool");
    }

    #[test]
    fn test_tool_registry() {
        let mut registry = ToolRegistry::new();
        registry.register(MockTool);

        assert_eq!(registry.len(), 1);
        assert!(registry.get("mock_tool").is_some());
        assert!(registry.get("nonexistent").is_none());
    }

    #[test]
    fn test_to_llm_tools() {
        let mut registry = ToolRegistry::new();
        registry.register(MockTool);

        let llm_tools = registry.to_llm_tools();
        assert_eq!(llm_tools.len(), 1);
        assert_eq!(llm_tools[0].name, "mock_tool");
    }

    struct AnotherTool;

    impl Tool<()> for AnotherTool {
        type Name = TestToolName;

        fn name(&self) -> TestToolName {
            TestToolName::AnotherTool
        }

        fn display_name(&self) -> &'static str {
            "Another Tool"
        }

        fn description(&self) -> &'static str {
            "Another tool for testing"
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({ "type": "object" })
        }

        async fn execute(&self, _ctx: &ToolContext<()>, _input: Value) -> Result<ToolResult> {
            Ok(ToolResult::success("Done"))
        }
    }

    #[test]
    fn test_filter_tools() {
        let mut registry = ToolRegistry::new();
        registry.register(MockTool);
        registry.register(AnotherTool);

        assert_eq!(registry.len(), 2);

        // Filter out mock_tool
        registry.filter(|name| name != "mock_tool");

        assert_eq!(registry.len(), 1);
        assert!(registry.get("mock_tool").is_none());
        assert!(registry.get("another_tool").is_some());
    }

    #[test]
    fn test_filter_tools_keep_all() {
        let mut registry = ToolRegistry::new();
        registry.register(MockTool);
        registry.register(AnotherTool);

        registry.filter(|_| true);

        assert_eq!(registry.len(), 2);
    }

    #[test]
    fn test_filter_tools_remove_all() {
        let mut registry = ToolRegistry::new();
        registry.register(MockTool);
        registry.register(AnotherTool);

        registry.filter(|_| false);

        assert!(registry.is_empty());
    }

    #[test]
    fn test_display_name() {
        let mut registry = ToolRegistry::new();
        registry.register(MockTool);

        let tool = registry.get("mock_tool").unwrap();
        assert_eq!(tool.display_name(), "Mock Tool");
    }

    struct ListenMockTool;

    impl ListenExecuteTool<()> for ListenMockTool {
        type Name = TestToolName;

        fn name(&self) -> TestToolName {
            TestToolName::MockTool
        }

        fn display_name(&self) -> &'static str {
            "Listen Mock Tool"
        }

        fn description(&self) -> &'static str {
            "A listen/execute mock tool for testing"
        }

        fn input_schema(&self) -> Value {
            serde_json::json!({ "type": "object" })
        }

        fn listen(
            &self,
            _ctx: &ToolContext<()>,
            _input: Value,
        ) -> impl futures::Stream<Item = ListenToolUpdate> + Send {
            futures::stream::iter(vec![ListenToolUpdate::Ready {
                operation_id: "op_1".to_string(),
                revision: 1,
                message: "ready".to_string(),
                snapshot: serde_json::json!({"ok": true}),
                expires_at: None,
            }])
        }

        async fn execute(
            &self,
            _ctx: &ToolContext<()>,
            _operation_id: &str,
            _expected_revision: u64,
        ) -> Result<ToolResult> {
            Ok(ToolResult::success("Executed"))
        }
    }

    #[test]
    fn test_listen_tool_registry() {
        let mut registry = ToolRegistry::new();
        registry.register_listen(ListenMockTool);

        assert_eq!(registry.len(), 1);
        assert!(registry.get_listen("mock_tool").is_some());
        assert!(registry.is_listen("mock_tool"));
    }
}