codex-memory 3.0.15

use crate::common::TestDatabaseManager;
use anyhow::Result;
use codex_memory::mcp_server::MCPHandlers;
use codex_memory::Storage;
use serde_json::json;
use serial_test::serial;
use std::sync::Arc;

#[tokio::test]
#[serial]
async fn test_mcp_search_basic() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Store test memories first
    let store_params1 = json!({
        "content": "Rust programming language fundamentals",
        "context": "Programming tutorial",
        "summary": "Basic concepts in Rust programming",
        "tags": ["rust", "programming", "fundamentals"]
    });

    let store_params2 = json!({
        "content": "Advanced Python data structures",
        "context": "Data structures guide",
        "summary": "Complex Python data handling techniques",
        "tags": ["python", "data-structures", "advanced"]
    });

    handlers
        .handle_tool_call("store_memory", store_params1)
        .await?;
    handlers
        .handle_tool_call("store_memory", store_params2)
        .await?;

    // Perform basic search
    let search_params = json!({
        "query": "programming language"
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await?;

    // Validate response structure
    assert!(result.is_object(), "Search result should be an object");
    assert!(result["results"].is_array(), "Should have results array");
    assert!(
        result["search_metadata"].is_object(),
        "Should have search metadata"
    );

    // Check metadata
    let metadata = &result["search_metadata"];
    assert_eq!(metadata["query"], "programming language");
    assert!(
        metadata["search_time_ms"].is_number(),
        "Should include search time"
    );
    assert!(
        metadata["total_results"].is_number(),
        "Should include result count"
    );

    // Results should have proper structure
    if let Some(results) = result["results"].as_array() {
        if !results.is_empty() {
            let first_result = &results[0];
            assert!(first_result["id"].is_string(), "Result should have ID");
            assert!(
                first_result["content"].is_string(),
                "Result should have content"
            );
            assert!(
                first_result["context"].is_string(),
                "Result should have context"
            );
            assert!(
                first_result["summary"].is_string(),
                "Result should have summary"
            );
            assert!(first_result["tags"].is_array(), "Result should have tags");
            assert!(
                first_result["combined_score"].is_number(),
                "Result should have combined score"
            );
        }
    }

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_with_all_parameters() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Store test memory
    let store_params = json!({
        "content": "Machine learning model training techniques",
        "context": "ML training guide",
        "summary": "Methods for training ML models effectively",
        "tags": ["machine-learning", "training", "models"]
    });

    handlers
        .handle_tool_call("store_memory", store_params)
        .await?;

    // Search with all parameters
    let search_params = json!({
        "query": "machine learning models",
        "tag_filter": ["machine-learning"],
        "use_tag_embedding": false,
        "use_content_embedding": false,
        "similarity_threshold": 0.3,
        "max_results": 5,
        "search_strategy": "hybrid",
        "boost_recent": true,
        "tag_weight": 0.3,
        "content_weight": 0.7
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await?;

    // Validate that parameters were respected
    let metadata = &result["search_metadata"];
    assert_eq!(metadata["similarity_threshold"], 0.3);
    assert_eq!(metadata["max_results"], 5);
    assert_eq!(metadata["search_strategy"], "hybrid");
    assert_eq!(metadata["boost_recent"], true);
    assert_eq!(metadata["tag_weight"], 0.3);
    assert_eq!(metadata["content_weight"], 0.7);
    assert_eq!(metadata["use_tag_embedding"], false);
    assert_eq!(metadata["use_content_embedding"], false);

    // Check tag filter
    if let Some(tag_filter) = metadata["tag_filter"].as_array() {
        assert_eq!(tag_filter.len(), 1);
        assert_eq!(tag_filter[0], "machine-learning");
    }

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_parameter_validation() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Test with invalid similarity threshold (should be clamped)
    let search_params = json!({
        "query": "test query",
        "similarity_threshold": 1.5  // Invalid: > 1.0
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await?;
    let metadata = &result["search_metadata"];
    assert_eq!(
        metadata["similarity_threshold"], 1.0,
        "Should clamp to max value"
    );

    // Test with invalid max_results (should be clamped)
    let search_params2 = json!({
        "query": "test query",
        "max_results": 200  // Invalid: > 100
    });

    let result2 = handlers
        .handle_tool_call("search_memory", search_params2)
        .await?;
    let metadata2 = &result2["search_metadata"];
    assert_eq!(metadata2["max_results"], 100, "Should clamp to max value");

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_missing_query() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Search without query parameter
    let search_params = json!({
        "similarity_threshold": 0.7
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await;

    // Should return an error
    assert!(result.is_err(), "Should fail with missing query parameter");

    if let Err(error) = result {
        assert!(
            error.to_string().contains("Missing query parameter"),
            "Error should mention missing query parameter"
        );
    }

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_default_parameters() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Store test memory
    let store_params = json!({
        "content": "Test content for default parameters",
        "context": "Test context",
        "summary": "Testing default search parameters",
        "tags": ["test", "defaults"]
    });

    handlers
        .handle_tool_call("store_memory", store_params)
        .await?;

    // Search with minimal parameters (should use defaults)
    let search_params = json!({
        "query": "test content"
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await?;

    // Check that defaults were applied correctly
    let metadata = &result["search_metadata"];
    assert_eq!(
        metadata["similarity_threshold"], 0.7,
        "Should use default similarity threshold"
    );
    assert_eq!(
        metadata["max_results"], 10,
        "Should use default max results"
    );
    assert_eq!(
        metadata["search_strategy"], "hybrid",
        "Should use default search strategy"
    );
    assert_eq!(
        metadata["boost_recent"], false,
        "Should use default boost_recent"
    );
    assert_eq!(metadata["tag_weight"], 0.4, "Should use default tag weight");
    assert_eq!(
        metadata["content_weight"], 0.6,
        "Should use default content weight"
    );
    assert_eq!(
        metadata["use_tag_embedding"], true,
        "Should use default tag embedding"
    );
    assert_eq!(
        metadata["use_content_embedding"], true,
        "Should use default content embedding"
    );
    assert!(
        metadata["tag_filter"].is_null(),
        "Should have null tag filter by default"
    );

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_different_strategies() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Store test memory
    let store_params = json!({
        "content": "Distributed systems architecture patterns",
        "context": "System design context",
        "summary": "Common patterns in distributed systems",
        "tags": ["distributed-systems", "architecture", "patterns"]
    });

    handlers
        .handle_tool_call("store_memory", store_params)
        .await?;

    let strategies = ["tags_first", "content_first", "hybrid"];

    for strategy in &strategies {
        let search_params = json!({
            "query": "distributed systems",
            "search_strategy": strategy,
            "use_tag_embedding": false,
            "use_content_embedding": false
        });

        let result = handlers
            .handle_tool_call("search_memory", search_params)
            .await?;

        // Each strategy should work without error
        assert!(result.is_array(), "Result should be valid array");
    }

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_empty_database() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Search in empty database
    let search_params = json!({
        "query": "anything"
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await?;

    // Should return empty results but valid structure
    assert!(result["results"].is_array(), "Should have results array");
    let results = result["results"].as_array().unwrap();
    assert!(
        results.is_empty(),
        "Should have no results in empty database"
    );

    let metadata = &result["search_metadata"];
    assert_eq!(metadata["total_results"], 0, "Should report zero results");
    assert!(
        metadata["search_time_ms"].is_number(),
        "Should still report search time"
    );

    manager.cleanup().await?;
    Ok(())
}

#[tokio::test]
#[serial]
async fn test_mcp_search_result_scoring() -> Result<()> {
    let mut manager = TestDatabaseManager::new()?;
    let pool = manager.setup_test_database().await?;
    let storage = Arc::new(Storage::new(pool));
    let handlers = MCPHandlers::new(storage);

    // Store memories with different relevance levels
    let high_relevance = json!({
        "content": "Advanced neural network architectures for deep learning",
        "context": "Deep learning research",
        "summary": "Comprehensive guide to neural network design",
        "tags": ["neural-networks", "deep-learning", "architecture"]
    });

    let medium_relevance = json!({
        "content": "Machine learning model evaluation techniques",
        "context": "ML evaluation guide",
        "summary": "Methods for evaluating ML model performance",
        "tags": ["machine-learning", "evaluation", "models"]
    });

    handlers
        .handle_tool_call("store_memory", high_relevance)
        .await?;
    handlers
        .handle_tool_call("store_memory", medium_relevance)
        .await?;

    // Search for neural networks (should prefer first memory)
    let search_params = json!({
        "query": "neural networks deep learning",
        "similarity_threshold": 0.1,
        "use_tag_embedding": false,
        "use_content_embedding": false
    });

    let result = handlers
        .handle_tool_call("search_memory", search_params)
        .await?;

    let results = result["results"].as_array().unwrap();
    if results.len() >= 2 {
        let first_score = results[0]["combined_score"].as_f64().unwrap();
        let second_score = results[1]["combined_score"].as_f64().unwrap();

        assert!(
            first_score >= second_score,
            "Results should be ordered by descending score: {} vs {}",
            first_score,
            second_score
        );
    }

    // All scores should be valid
    for result in results {
        let score = result["combined_score"].as_f64().unwrap();
        assert!(score >= 0.0, "Scores should be non-negative");
        assert!(
            score <= 1.0 || score > 1.0,
            "Scores should be in valid range"
        ); // Allow for boosting
    }

    manager.cleanup().await?;
    Ok(())
}