use tower_mcp::{McpRouter, ToolBuilder, CallToolResult};
use schemars::JsonSchema;
use serde::Deserialize;

// Define your input type - schema is auto-generated
#[derive(Debug, Deserialize, JsonSchema)]
struct GreetInput {
    name: String,
}

// Build a tool with type-safe handler
let greet = ToolBuilder::new("greet")
    .title("Greet")
    .description("Greet someone by name")
    .handler(|input: GreetInput| async move {
        Ok(CallToolResult::text(format!("Hello, {}!", input.name)))
    })
    .build();

// Create router with tools
let router = McpRouter::new()
    .server_info("my-server", "1.0.0")
    .instructions("This server provides greeting functionality")
    .tool(greet);

// The router implements tower::Service and can be composed with middleware

use tower_mcp::{ToolBuilder, CallToolResult};
use schemars::JsonSchema;
use serde::Deserialize;

#[derive(Debug, Deserialize, JsonSchema)]
struct AddInput {
    a: i64,
    b: i64,
}

let add = ToolBuilder::new("add")
    .description("Add two numbers")
    .read_only()  // Hint: this tool doesn't modify state
    .handler(|input: AddInput| async move {
        Ok(CallToolResult::text(format!("{}", input.a + input.b)))
    })
    .build();

use tower_mcp::{tool_fn, prompt_fn, resource_fn, resource_template_fn};
use tower_mcp::{CallToolResult, McpRouter};
use tower_mcp::protocol::{GetPromptResult, ReadResourceResult};

#[derive(Debug, Deserialize, JsonSchema)]
struct AddInput { a: i64, b: i64 }

#[tool_fn(description = "Add two numbers")]
async fn add(input: AddInput) -> Result<CallToolResult, tower_mcp::Error> {
    Ok(CallToolResult::text(format!("{}", input.a + input.b)))
}

#[prompt_fn(description = "Greet someone", args(name = "Name to greet"))]
async fn greet(args: HashMap<String, String>) -> Result<GetPromptResult, tower_mcp::Error> {
    let name = args.get("name").cloned().unwrap_or_default();
    Ok(GetPromptResult::user_message(format!("Hello, {name}!")))
}

#[resource_fn(uri = "app://config", description = "App configuration")]
async fn config() -> Result<ReadResourceResult, tower_mcp::Error> {
    Ok(ReadResourceResult::text("app://config", "debug=true"))
}

// Each macro generates a constructor: add_tool(), greet_prompt(), config_resource()
let router = McpRouter::new()
    .server_info("my-server", "1.0.0")
    .tool(add_tool())
    .prompt(greet_prompt())
    .resource(config_resource());

use tower_mcp::tool::McpTool;
use tower_mcp::{Result, CallToolResult};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::sync::Arc;

struct Calculator {
    precision: u32,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct CalcInput {
    expression: String,
}

impl McpTool for Calculator {
    const NAME: &'static str = "calculate";
    const DESCRIPTION: &'static str = "Evaluate a mathematical expression";

    type Input = CalcInput;
    type Output = f64;

    async fn call(&self, input: Self::Input) -> Result<Self::Output> {
        // Your calculation logic here
        Ok(42.0)
    }
}

// Convert to Tool and register
let calc = Calculator { precision: 10 };
let router = McpRouter::new().tool(calc.into_tool());

use std::sync::Arc;
use tower_mcp::{ToolBuilder, CallToolResult};
use tower_mcp::extract::{State, Context, Json};

#[derive(Clone)]
struct AppState { db_url: String }

let state = Arc::new(AppState { db_url: "postgres://...".into() });

let search = ToolBuilder::new("search")
    .description("Search with progress updates")
    .extractor_handler(state, |
        State(app): State<Arc<AppState>>,
        ctx: Context,
        Json(input): Json<SearchInput>,
    | async move {
        // Report progress
        ctx.report_progress(0.5, Some(1.0), Some("Searching...")).await;
        // Use state
        let results = format!("Searched {} for: {}", app.db_url, input.query);
        Ok(CallToolResult::text(results))
    })
    .build();

use tower_mcp::ResourceBuilder;

// Static resource with inline content
let config = ResourceBuilder::new("file:///config.json")
    .name("Configuration")
    .description("Server configuration")
    .json(serde_json::json!({
        "version": "1.0.0",
        "debug": true
    }))
    .build();

// Dynamic resource with handler
let status = ResourceBuilder::new("app:///status")
    .name("Server Status")
    .description("Current server status")
    .handler(|| async {
        Ok("Running".to_string())
    })
    .build();

let router = McpRouter::new()
    .resource(config)
    .resource(status);

use tower_mcp::{PromptBuilder, GetPromptResult};

let greet = PromptBuilder::new("greet")
    .description("Generate a greeting")
    .required_arg("name", "Name to greet")
    .optional_arg("style", "Greeting style (formal/casual)")
    .handler(|args| async move {
        let name = args.get("name").map(|s| s.as_str()).unwrap_or("World");
        let style = args.get("style").map(|s| s.as_str()).unwrap_or("casual");

        let text = match style {
            "formal" => format!("Good day, {}. How may I assist you?", name),
            _ => format!("Hey {}!", name),
        };

        // Builder handles message construction
        Ok(GetPromptResult::builder()
            .description("A friendly greeting")
            .user(text)
            .build())
    })
    .build();

let router = McpRouter::new().prompt(greet);

// Merge routers (combines all tools/resources/prompts)
let api_router = McpRouter::new()
    .tool(search_tool)
    .tool(fetch_tool);

let admin_router = McpRouter::new()
    .tool(reset_tool)
    .tool(stats_tool);

let combined = McpRouter::new()
    .merge(api_router)
    .merge(admin_router);

// Nest with prefix (adds prefix to all tool names)
let v1 = McpRouter::new().tool(legacy_tool);
let v2 = McpRouter::new().tool(new_tool);

let versioned = McpRouter::new()
    .nest("v1", v1)   // Tools become "v1_legacy_tool"
    .nest("v2", v2);  // Tools become "v2_new_tool"

use tower_mcp::proxy::McpProxy;
use tower_mcp::client::StdioClientTransport;

let proxy = McpProxy::builder("my-proxy", "1.0.0")
    .backend("db", StdioClientTransport::spawn("db-server", &[]).await?)
    .await
    .backend("fs", StdioClientTransport::spawn("fs-server", &[]).await?)
    .await
    .build()
    .await?;

// Tools become db_query, fs_read, etc.
// Serve over any transport.
StdioTransport::new(proxy).run().await?;

use std::time::Duration;
use tower::timeout::TimeoutLayer;

let proxy = McpProxy::builder("proxy", "1.0.0")
    .backend("fast", cache_transport).await
    .backend_layer(TimeoutLayer::new(Duration::from_secs(2)))
    .backend("slow", llm_transport).await
    .backend_layer(TimeoutLayer::new(Duration::from_secs(60)))
    .build().await?;

use std::sync::Arc;
use tower_mcp::extract::Extension;

#[derive(Clone)]
struct AppState {
    db: DatabasePool,
    config: Config,
}

let state = Arc::new(AppState { /* ... */ });

// Tools access state via Extension<T> extractor
let tool = ToolBuilder::new("query")
    .extractor_handler(
        (),
        |Extension(app): Extension<Arc<AppState>>, Json(input): Json<QueryInput>| async move {
            let result = app.db.query(&input.sql).await?;
            Ok(CallToolResult::text(result))
        },
    )
    .build();

let router = McpRouter::new()
    .with_state(state)  // Makes AppState available to all handlers
    .tool(tool);

use tower_mcp::{McpRouter, StdioTransport};

let router = McpRouter::new()
    .server_info("my-server", "1.0.0")
    .tool(my_tool);

// Serve over stdin/stdout
StdioTransport::new(router).serve().await?;

use tower_mcp::{McpRouter, HttpTransport};

let router = McpRouter::new()
    .server_info("my-server", "1.0.0")
    .tool(my_tool);

let transport = HttpTransport::new(router);
let app = transport.into_router();

// Serve with axum
let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await?;
axum::serve(listener, app).await?;

use tower_mcp::auth::extract_api_key;
use axum::middleware;

// Add auth layer to the HTTP transport
let app = transport.into_router()
    .layer(middleware::from_fn(auth_middleware));

use tower::ServiceBuilder;
use tower_mcp::middleware::{AuditLayer, McpTracingLayer};

let transport = StdioTransport::new(router)
    .layer(
        ServiceBuilder::new()
            .layer(McpTracingLayer::new())
            .layer(AuditLayer::new())
            .into_inner(),
    );

use tower_mcp::TestClient;
use serde_json::json;

let mut client = TestClient::from_router(router);
client.initialize().await;

// List and call tools
let tools = client.list_tools().await;
assert_eq!(tools.len(), 1);

let result = client.call_tool("greet", json!({"name": "World"})).await;
assert_eq!(result.all_text(), "Hello, World!");

// Typed deserialization
let stats: ServerStats = client.call_tool_typed("stats", json!({})).await;

// Assert expected errors
let err = client.call_tool_expect_error("missing", json!({})).await;

use tower_mcp::CapabilityFilter;

// Hide write tools from sessions that aren't authorized
let router = McpRouter::new()
    .tool(read_tool)
    .tool(write_tool)
    .tool_filter(CapabilityFilter::write_guard(|session| {
        session.get::<UserRole>()
            .map(|r| r.is_admin())
            .unwrap_or(false)
    }));

let router = McpRouter::new()
    .resource(public_resource)
    .resource(internal_resource)
    .resource_filter(CapabilityFilter::new(|session, resource: &Resource| {
        !resource.name().contains("internal") || session.get::<AdminClaim>().is_some()
    }));