//! `Tool` trait + object-safe `AnyTool` + `Erased<T>` adapter.

use crate::tool::{ToolContext, ToolEventSink, ToolKind};
use async_trait::async_trait;
use schemars::{JsonSchema, schema_for};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};

/// Streaming-aware tool input.
///
/// Tools opt into streaming via [`Tool::supports_input_streaming`]. If they
/// do not, the harness buffers and always provides [`ToolInput::Final`].
pub enum ToolInput<T> {
    Final(T),
    Partial {
        partial: mpsc::UnboundedReceiver<serde_json::Value>,
        final_input: oneshot::Receiver<T>,
    },
}

/// JSON-shaped variant used by the object-safe `AnyTool` trait.
pub enum AnyToolInput {
    Final(serde_json::Value),
    Partial {
        partial: mpsc::UnboundedReceiver<serde_json::Value>,
        final_input: oneshot::Receiver<serde_json::Value>,
    },
}

/// Strongly-typed tool implementation.
#[async_trait]
pub trait Tool: Send + Sync + 'static {
    type Input: for<'de> Deserialize<'de> + Serialize + JsonSchema + Send + 'static;
    type Output: Serialize + Send + 'static;

    const NAME: &'static str;

    fn kind() -> ToolKind;

    fn input_schema() -> serde_json::Value {
        serde_json::to_value(schema_for!(Self::Input)).expect("schema_for must serialise")
    }

    fn supports_input_streaming() -> bool { false }

    async fn run(
        self: Arc<Self>,
        input: ToolInput<Self::Input>,
        events: ToolEventSink,
        ctx: &ToolContext,
    ) -> Result<Self::Output, Self::Output>;

    fn erase(self: Arc<Self>) -> Arc<dyn AnyTool> where Self: Sized { Arc::new(Erased(self)) }
}

/// Object-safe variant. The registry stores `Arc<dyn AnyTool>`.
#[async_trait]
pub trait AnyTool: Send + Sync + 'static {
    fn name(&self) -> &'static str;
    fn kind(&self) -> ToolKind;
    fn input_schema(&self) -> serde_json::Value;
    fn supports_input_streaming(&self) -> bool;

    async fn run(
        self: Arc<Self>,
        input: AnyToolInput,
        events: ToolEventSink,
        ctx: &ToolContext,
    ) -> Result<serde_json::Value, serde_json::Value>;
}

/// Adapter from a typed `Tool` to `AnyTool`.
pub struct Erased<T: Tool>(pub Arc<T>);

#[async_trait]
impl<T: Tool> AnyTool for Erased<T> {
    fn name(&self) -> &'static str { T::NAME }
    fn kind(&self) -> ToolKind { T::kind() }
    fn input_schema(&self) -> serde_json::Value { T::input_schema() }
    fn supports_input_streaming(&self) -> bool { T::supports_input_streaming() }

    async fn run(
        self: Arc<Self>,
        input: AnyToolInput,
        events: ToolEventSink,
        ctx: &ToolContext,
    ) -> Result<serde_json::Value, serde_json::Value> {
        let typed = match input {
            AnyToolInput::Final(v) => {
                let parsed: T::Input = serde_json::from_value(v).map_err(|e| {
                    serde_json::json!({ "error": format!("invalid input: {e}") })
                })?;
                ToolInput::Final(parsed)
            }
            AnyToolInput::Partial { partial, final_input: _ } => {
                // For v1: tools that opt into streaming receive Partial; the
                // typed-final-input wiring is added when a streaming tool ships.
                // Until then, only Final is fed by the dispatcher.
                let _ = partial;
                return Err(serde_json::json!({
                    "error": "streaming inputs are not yet wired in v1 dispatcher"
                }));
            }
        };
        let inner = self.0.clone();
        let result = inner.run(typed, events, ctx).await;
        match result {
            Ok(o)  => Ok(serde_json::to_value(o).unwrap_or(serde_json::Value::Null)),
            Err(o) => Err(serde_json::to_value(o).unwrap_or(serde_json::Value::Null)),
        }
    }
}

#[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::{ClientKind, ProtocolKind};
    use std::path::PathBuf;

    #[derive(Serialize, Deserialize, JsonSchema)]
    struct EchoInput { msg: String }

    #[derive(Serialize, Deserialize)]
    struct EchoOutput { echoed: String }

    struct EchoTool;

    #[async_trait]
    impl Tool for EchoTool {
        type Input = EchoInput;
        type Output = EchoOutput;
        const NAME: &'static str = "echo";
        fn kind() -> ToolKind { ToolKind::Other }

        async fn run(
            self: Arc<Self>,
            input: ToolInput<Self::Input>,
            _events: ToolEventSink,
            _ctx: &ToolContext,
        ) -> Result<Self::Output, Self::Output> {
            let i = match input { ToolInput::Final(i) => i, _ => unreachable!() };
            Ok(EchoOutput { echoed: i.msg })
        }
    }

    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,
        )
    }

    #[tokio::test]
    async fn typed_tool_runs_and_returns_output() {
        let tool: Arc<EchoTool> = Arc::new(EchoTool);
        let any: Arc<dyn AnyTool> = tool.erase();
        let (sink, _rx) = ToolEventSink::new();

        let result = any.run(
            AnyToolInput::Final(serde_json::json!({ "msg": "hello" })),
            sink,
            &ctx(),
        ).await.unwrap();

        assert_eq!(result["echoed"], "hello");
    }

    #[tokio::test]
    async fn invalid_input_returns_structured_error() {
        let tool: Arc<EchoTool> = Arc::new(EchoTool);
        let any: Arc<dyn AnyTool> = tool.erase();
        let (sink, _rx) = ToolEventSink::new();

        let err = any.run(
            AnyToolInput::Final(serde_json::json!({ "wrong": 1 })),
            sink,
            &ctx(),
        ).await.unwrap_err();

        assert!(err["error"].as_str().unwrap().contains("invalid input"));
    }

    #[test]
    fn name_and_kind_round_trip_through_erase() {
        let tool: Arc<EchoTool> = Arc::new(EchoTool);
        let any: Arc<dyn AnyTool> = tool.erase();
        assert_eq!(any.name(), "echo");
        assert_eq!(any.kind(), ToolKind::Other);
    }
}
