reasonkit-core 0.1.8

//! Docset Tools Benchmarks
//!
//! Benchmarks for the docset MCP tools implementation in ReasonKit.

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use reasonkit::mcp::docset_tools::*;
use reasonkit::mcp::tools::Tool;
use reasonkit::retrieval::HybridRetriever;
use serde_json::json;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::runtime::Runtime;

fn tool_by_name<'a>(tools: &'a [Tool], name: &str) -> &'a Tool {
    tools
        .iter()
        .find(|tool| tool.name == name)
        .unwrap_or_else(|| panic!("Expected tool '{name}' to exist"))
}

fn get_string_arg<'a>(args: &'a HashMap<String, serde_json::Value>, key: &str) -> Option<&'a str> {
    args.get(key).and_then(|value| value.as_str())
}

/// Benchmark tool definition creation
fn bench_tool_definitions_creation(c: &mut Criterion) {
    c.bench_function("docset_tool_definitions_creation", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            black_box(tools);
        })
    });
}

/// Benchmark argument extraction
fn bench_argument_extraction(c: &mut Criterion) {
    c.bench_function("docset_argument_extraction", |b| {
        // Setup test data
        let mut args = HashMap::new();
        args.insert("doc_id".to_string(), json!("test-document-id"));
        args.insert("chunk_id".to_string(), json!("test-chunk-id"));

        b.iter(|| {
            let result = get_string_arg(&args, "doc_id");
            black_box(result);
        })
    });
}

/// Benchmark tool name validation
fn bench_tool_name_validation(c: &mut Criterion) {
    c.bench_function("docset_tool_name_validation", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
            black_box(names.contains(&"rkmem_docs_search"));
        })
    });
}

/// Benchmark JSON schema validation
fn bench_json_schema_validation(c: &mut Criterion) {
    let tools = DocsetHandler::tool_definitions();
    let tool = tool_by_name(&tools, "rkmem_docs_search").clone();

    c.bench_function("docset_json_schema_validation", move |b| {
        b.iter(|| black_box(tool.input_schema.is_object()))
    });
}

/// Benchmark tool definitions access
fn bench_tool_definitions_access(c: &mut Criterion) {
    c.bench_function("docset_tool_definitions_access", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            for tool in tools {
                black_box(&tool.name);
                black_box(&tool.description);
                black_box(&tool.input_schema);
            }
        })
    });
}

/// Benchmark individual tool creation
fn bench_individual_tool_creation(c: &mut Criterion) {
    c.bench_function("docset_individual_tool_creation", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let get_chunk = tool_by_name(&tools, "rkmem_docs_get_chunk");
            let get_neighbors = tool_by_name(&tools, "rkmem_docs_get_neighbors");
            let search = tool_by_name(&tools, "rkmem_docs_search");
            let ingest = tool_by_name(&tools, "rkmem_docs_ingest");
            let list = tool_by_name(&tools, "rkmem_docs_list");

            black_box(get_chunk);
            black_box(get_neighbors);
            black_box(search);
            black_box(ingest);
            black_box(list);
        })
    });
}

/// Benchmark handler creation
fn bench_handler_creation(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    c.bench_function("docset_handler_creation", |b| {
        b.to_async(&rt).iter(|| async {
            let retriever = HybridRetriever::in_memory().expect("Failed to create retriever");
            let handler = DocsetHandler::new(Arc::new(retriever));
            black_box(handler);
        })
    });
}

/// Benchmark tool calling with missing arguments (error case)
fn bench_tool_calling_error_case(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();

    c.bench_function("docset_tool_calling_error_case", |b| {
        b.to_async(&rt).iter(|| async {
            let retriever = HybridRetriever::in_memory().expect("Failed to create retriever");
            let handler = DocsetHandler::new(Arc::new(retriever));

            let args = HashMap::new(); // Empty args should cause error
            let result = handler.call_tool("rkmem_docs_get_chunk", args).await;
            let _ = black_box(result);
        })
    });
}

/// Benchmark extract_required_string with various input sizes
fn bench_argument_extraction_various_sizes(c: &mut Criterion) {
    c.bench_function("docset_argument_extraction_short", |b| {
        let mut args = HashMap::new();
        args.insert("key".to_string(), json!("short"));
        b.iter(|| {
            let result = get_string_arg(&args, "key");
            black_box(result);
        })
    });

    c.bench_function("docset_argument_extraction_medium", |b| {
        let mut args = HashMap::new();
        args.insert(
            "key".to_string(),
            json!("medium length string for testing purposes"),
        );
        b.iter(|| {
            let result = get_string_arg(&args, "key");
            black_box(result);
        })
    });

    c.bench_function("docset_argument_extraction_long", |b| {
        let mut args = HashMap::new();
        args.insert("key".to_string(), json!("very long string that is meant to simulate a large input for performance testing of the extraction function to ensure it handles various sizes efficiently without significant performance degradation"));
        b.iter(|| {
            let result = get_string_arg(&args, "key");
            black_box(result);
        })
    });
}

/// Benchmark schema validation for different tools
fn bench_schema_validation_all_tools(c: &mut Criterion) {
    c.bench_function("docset_schema_validation_all_tools", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            for tool in tools {
                black_box(tool.input_schema.is_object());
            }
        })
    });
}

/// Benchmark tool name lookup performance
fn bench_tool_name_lookup(c: &mut Criterion) {
    c.bench_function("docset_tool_name_lookup_existing", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
            let exists = names.contains(&"rkmem_docs_search");
            black_box(exists);
        })
    });

    c.bench_function("docset_tool_name_lookup_nonexistent", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
            let exists = names.contains(&"nonexistent_tool_name");
            black_box(exists);
        })
    });
}

/// Benchmark memory allocation for tool definitions
fn bench_tool_definitions_memory_allocation(c: &mut Criterion) {
    c.bench_function("docset_tool_definitions_memory_allocation", |b| {
        b.iter(|| {
            // This tests the memory allocation overhead of creating tool definitions
            let tools = DocsetHandler::tool_definitions();
            let mut total_chars = 0;

            for tool in &tools {
                total_chars += tool.name.len();
                if let Some(desc) = &tool.description {
                    total_chars += desc.len();
                }
                // Add JSON schema size approximation
                total_chars += tool.input_schema.to_string().len();
            }

            black_box(total_chars);
        })
    });
}

/// Benchmark concurrent tool definition access
fn bench_concurrent_tool_access(c: &mut Criterion) {
    c.bench_function("docset_concurrent_tool_access", |b| {
        b.iter(|| {
            // Simulate concurrent access by accessing multiple times
            let tools1 = DocsetHandler::tool_definitions();
            let tools2 = DocsetHandler::tool_definitions();
            let tools3 = DocsetHandler::tool_definitions();

            black_box(tools1.len());
            black_box(tools2.len());
            black_box(tools3.len());
        })
    });
}

/// Benchmark error message generation
fn bench_error_message_generation(c: &mut Criterion) {
    c.bench_function("docset_error_message_generation", |b| {
        b.iter(|| {
            // Test various error message generation scenarios
            let error1 = format!("Missing required argument: {}", "doc_id");
            let error2 = format!("Invalid {} format", "UUID");
            let error3 = format!("{} not found in {}", "Chunk", "Document");

            black_box(error1);
            black_box(error2);
            black_box(error3);
        })
    });
}

/// Benchmark JSON serialization of results
fn bench_json_serialization_performance(c: &mut Criterion) {
    c.bench_function("docset_json_serialization_small", |b| {
        b.iter(|| {
            let small_result = json!({
                "status": "success",
                "message": "test"
            });
            let serialized = serde_json::to_string(&small_result).unwrap();
            black_box(serialized);
        })
    });

    c.bench_function("docset_json_serialization_medium", |b| {
        b.iter(|| {
            let medium_result = json!({
                "status": "success",
                "doc_id": "123e4567-e89b-12d3-a456-426614174000",
                "chunk_id": "abcdef12-3456-7890-abcd-ef1234567890",
                "text": "This is a sample text for testing JSON serialization performance with medium sized content.",
                "section": "Test Section",
                "index": 0,
                "start_char": 0,
                "end_char": 85
            });
            let serialized = serde_json::to_string(&medium_result).unwrap();
            black_box(serialized);
        })
    });
}

/// Benchmark tool handler method resolution
fn bench_tool_handler_method_resolution(c: &mut Criterion) {
    c.bench_function("docset_tool_handler_method_resolution", |b| {
        b.iter(|| {
            // Test the match statement performance for tool resolution
            let tool_names = [
                "rkmem_docs_get_chunk",
                "rkmem_docs_get_neighbors",
                "rkmem_docs_search",
                "rkmem_docs_ingest",
                "rkmem_docs_list",
                "unknown_tool",
            ];

            for &name in &tool_names {
                match name {
                    "rkmem_docs_get_chunk" => black_box(1),
                    "rkmem_docs_get_neighbors" => black_box(2),
                    "rkmem_docs_search" => black_box(3),
                    "rkmem_docs_ingest" => black_box(4),
                    "rkmem_docs_list" => black_box(5),
                    _ => black_box(0), // Unknown tool
                };
            }
        })
    });
}

/// Benchmark tool definitions count validation
fn bench_tool_definitions_count_validation(c: &mut Criterion) {
    c.bench_function("docset_tool_definitions_count_validation", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let count = tools.len();
            black_box(count == 5); // Should always be 5 tools
        })
    });
}

/// Benchmark schema property access
fn bench_schema_property_access(c: &mut Criterion) {
    let tools = DocsetHandler::tool_definitions();
    let tool = tool_by_name(&tools, "rkmem_docs_search").clone();

    c.bench_function("docset_schema_property_access", move |b| {
        b.iter(|| {
            let schema_obj = tool.input_schema.as_object().unwrap();

            // Access various schema properties
            let properties = schema_obj.get("properties");
            let required = schema_obj.get("required");
            let additional_props = schema_obj.get("additionalProperties");

            black_box(properties);
            black_box(required);
            black_box(additional_props);
        })
    });
}

/// Benchmark tool description access patterns
fn bench_tool_description_access(c: &mut Criterion) {
    c.bench_function("docset_tool_description_access", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let mut total_length = 0;

            for tool in tools {
                if let Some(description) = tool.description {
                    total_length += description.len();
                }
            }

            black_box(total_length);
        })
    });
}

/// Benchmark argument extraction error cases
fn bench_argument_extraction_error_cases(c: &mut Criterion) {
    c.bench_function("docset_argument_extraction_missing_key", |b| {
        let args = HashMap::new(); // Empty map
        b.iter(|| {
            let result = get_string_arg(&args, "nonexistent_key");
            black_box(result.is_none());
        })
    });

    c.bench_function("docset_argument_extraction_wrong_type", |b| {
        let mut args = HashMap::new();
        args.insert("key".to_string(), json!(123)); // Not a string
        b.iter(|| {
            let result = get_string_arg(&args, "key");
            black_box(result.is_none());
        })
    });
}

/// Benchmark tool name iteration
fn bench_tool_name_iteration(c: &mut Criterion) {
    c.bench_function("docset_tool_name_iteration", |b| {
        b.iter(|| {
            let tools = DocsetHandler::tool_definitions();
            let mut names = Vec::new();

            for tool in tools {
                names.push(tool.name);
            }

            black_box(names);
        })
    });
}

criterion_group!(
    benches,
    bench_tool_definitions_creation,
    bench_argument_extraction,
    bench_tool_name_validation,
    bench_json_schema_validation,
    bench_tool_definitions_access,
    bench_individual_tool_creation,
    bench_handler_creation,
    bench_tool_calling_error_case,
    bench_argument_extraction_various_sizes,
    bench_schema_validation_all_tools,
    bench_tool_name_lookup,
    bench_tool_definitions_memory_allocation,
    bench_concurrent_tool_access,
    bench_error_message_generation,
    bench_json_serialization_performance,
    bench_tool_handler_method_resolution,
    bench_tool_definitions_count_validation,
    bench_schema_property_access,
    bench_tool_description_access,
    bench_argument_extraction_error_cases,
    bench_tool_name_iteration
);

criterion_main!(benches);