ipfrs-tensorlogic 0.2.0

//! Integration tests for distributed reasoning
//!
//! These tests verify the complete distributed reasoning pipeline including:
//! - Remote knowledge retrieval
//! - Distributed goal resolution
//! - Recursive query handling with tabling
//! - Proof assembly from distributed fragments

use ipfrs_tensorlogic::{
    Constant, DistributedGoalResolver, DistributedProofAssembler, DistributedReasoner,
    FactDiscoveryRequest, GoalResolutionRequest, IncrementalLoadRequest, KnowledgeBase,
    MockRemoteKnowledgeProvider, Predicate, QueryRequest, RemoteKnowledgeProvider, Rule,
    Substitution, TabledInferenceEngine, Term,
};
use std::collections::HashSet;
use std::sync::Arc;

#[tokio::test]
async fn test_local_and_remote_resolution() {
    // Create local knowledge base
    let mut local_kb = KnowledgeBase::new();
    local_kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Const(Constant::String("bob".to_string())),
        ],
    ));

    // Create remote knowledge base (simulated)
    let mut remote_kb = KnowledgeBase::new();
    remote_kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("bob".to_string())),
            Term::Const(Constant::String("charlie".to_string())),
        ],
    ));

    // Create resolver with local KB
    let mut resolver = DistributedGoalResolver::new(Arc::new(local_kb));

    // Test local resolution
    let goal = Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Var("X".to_string()),
        ],
    );

    let solutions = resolver.resolve(&goal, &Substitution::new()).await.unwrap();
    assert_eq!(solutions.len(), 1);
    assert_eq!(
        solutions[0].get("X"),
        Some(&Term::Const(Constant::String("bob".to_string())))
    );

    // Add remote provider and test remote resolution
    let provider = Arc::new(MockRemoteKnowledgeProvider::new(Arc::new(remote_kb)));
    resolver = resolver.with_provider(provider);

    let remote_goal = Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("bob".to_string())),
            Term::Var("Y".to_string()),
        ],
    );

    let remote_solutions = resolver
        .resolve(&remote_goal, &Substitution::new())
        .await
        .unwrap();
    assert!(!remote_solutions.is_empty());
}

#[tokio::test]
async fn test_fact_prefetching() {
    // Create remote knowledge base with multiple facts
    let mut remote_kb = KnowledgeBase::new();
    for i in 0..5 {
        remote_kb.add_fact(Predicate::new(
            "number".to_string(),
            vec![Term::Const(Constant::Int(i))],
        ));
    }

    let provider = Arc::new(MockRemoteKnowledgeProvider::new(Arc::new(remote_kb)));
    let mut resolver =
        DistributedGoalResolver::new(Arc::new(KnowledgeBase::new())).with_provider(provider);

    // Prefetch facts
    let count = resolver.prefetch_facts("number").await.unwrap();
    assert_eq!(count, 5);

    // Verify cached facts
    let cached = resolver.get_cached_facts("number");
    assert!(cached.is_some());
    assert_eq!(cached.unwrap().len(), 5);
}

#[tokio::test]
async fn test_query_request_response() {
    let mut kb = KnowledgeBase::new();
    kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Const(Constant::String("bob".to_string())),
        ],
    ));

    let provider = MockRemoteKnowledgeProvider::new(Arc::new(kb));

    let request = QueryRequest {
        predicate_name: "parent".to_string(),
        ground_args: vec![],
        max_results: 10,
        max_depth: 5,
        request_id: "test_123".to_string(),
    };

    let response = provider.query_predicate(request).await.unwrap();
    assert_eq!(response.predicates.len(), 1);
    assert_eq!(response.peer_id, "mock_peer");
    assert!(!response.has_more);
}

#[tokio::test]
async fn test_fact_discovery_multi_hop() {
    let mut kb = KnowledgeBase::new();

    // Add facts for different entities
    kb.add_fact(Predicate::new(
        "city".to_string(),
        vec![Term::Const(Constant::String("Tokyo".to_string()))],
    ));
    kb.add_fact(Predicate::new(
        "city".to_string(),
        vec![Term::Const(Constant::String("Paris".to_string()))],
    ));
    kb.add_fact(Predicate::new(
        "city".to_string(),
        vec![Term::Const(Constant::String("London".to_string()))],
    ));

    let provider = MockRemoteKnowledgeProvider::new(Arc::new(kb));

    let request = FactDiscoveryRequest {
        predicate_name: "city".to_string(),
        arg_patterns: vec![],
        max_hops: 3,
        ttl: 30,
        exclude_peers: HashSet::new(),
    };

    let response = provider.discover_facts(request).await.unwrap();
    assert_eq!(response.facts.len(), 3);
    assert_eq!(response.peers_queried, 1);

    // All facts should be at hop 0 (from the same peer)
    for hop in response.hops.values() {
        assert_eq!(*hop, 0);
    }
}

#[tokio::test]
async fn test_incremental_loading_pagination() {
    let mut kb = KnowledgeBase::new();

    // Add 20 facts
    for i in 0..20 {
        kb.add_fact(Predicate::new(
            "item".to_string(),
            vec![Term::Const(Constant::Int(i))],
        ));
    }

    let provider = MockRemoteKnowledgeProvider::new(Arc::new(kb));

    // Load first batch
    let request1 = IncrementalLoadRequest {
        predicate_name: "item".to_string(),
        batch_size: 5,
        offset: 0,
        filter: None,
    };

    let response1 = provider.load_incremental(request1).await.unwrap();
    assert_eq!(response1.batch.len(), 5);
    assert_eq!(response1.total_count, 20);
    assert!(!response1.is_last);
    assert_eq!(response1.next_offset, Some(5));

    // Load second batch
    let request2 = IncrementalLoadRequest {
        predicate_name: "item".to_string(),
        batch_size: 5,
        offset: 5,
        filter: None,
    };

    let response2 = provider.load_incremental(request2).await.unwrap();
    assert_eq!(response2.batch.len(), 5);
    assert_eq!(response2.next_offset, Some(10));

    // Load last batch
    let request3 = IncrementalLoadRequest {
        predicate_name: "item".to_string(),
        batch_size: 5,
        offset: 15,
        filter: None,
    };

    let response3 = provider.load_incremental(request3).await.unwrap();
    assert_eq!(response3.batch.len(), 5);
    assert!(response3.is_last);
    assert_eq!(response3.next_offset, None);
}

#[tokio::test]
async fn test_goal_resolution_with_proof() {
    let mut kb = KnowledgeBase::new();
    kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Const(Constant::String("bob".to_string())),
        ],
    ));

    let provider = MockRemoteKnowledgeProvider::new(Arc::new(kb));

    let goal = Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Var("X".to_string()),
        ],
    );

    let request = GoalResolutionRequest {
        goal,
        substitution: std::collections::HashMap::new(),
        depth: 0,
        requester: "test".to_string(),
        request_id: "test_456".to_string(),
    };

    let response = provider.resolve_goal(request).await.unwrap();
    assert!(response.solved);
    assert_eq!(response.solutions.len(), 1);
    assert!(response.proof.is_some());

    // Verify the proof
    let proof = response.proof.unwrap();
    assert!(proof.is_fact());
}

#[tokio::test]
async fn test_distributed_reasoner_with_cache() {
    let mut kb = KnowledgeBase::new();
    kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Const(Constant::String("bob".to_string())),
        ],
    ));

    // Create cache manager
    let cache_manager = Arc::new(ipfrs_tensorlogic::CacheManager::new());

    let reasoner = DistributedReasoner::with_cache(cache_manager.clone()).unwrap();

    let goal = Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Var("X".to_string()),
        ],
    );

    // First query - should cache
    let solutions1 = reasoner.query(&goal, &kb).await.unwrap();
    assert_eq!(solutions1.len(), 1);

    // Second query - should hit cache
    let solutions2 = reasoner.query(&goal, &kb).await.unwrap();
    assert_eq!(solutions2.len(), 1);

    // Verify cache stats
    let stats = reasoner.cache_stats().unwrap();
    assert!(stats.query_stats.hits >= 1);
}

#[test]
fn test_tabled_inference_recursive() {
    let mut kb = KnowledgeBase::new();

    // Add parent facts
    kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Const(Constant::String("bob".to_string())),
        ],
    ));
    kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("bob".to_string())),
            Term::Const(Constant::String("charlie".to_string())),
        ],
    ));

    // Add base rule: ancestor(X, Y) :- parent(X, Y)
    kb.add_rule(Rule::new(
        Predicate::new(
            "ancestor".to_string(),
            vec![Term::Var("X".to_string()), Term::Var("Y".to_string())],
        ),
        vec![Predicate::new(
            "parent".to_string(),
            vec![Term::Var("X".to_string()), Term::Var("Y".to_string())],
        )],
    ));

    // Add recursive rule: ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z)
    kb.add_rule(Rule::new(
        Predicate::new(
            "ancestor".to_string(),
            vec![Term::Var("X".to_string()), Term::Var("Z".to_string())],
        ),
        vec![
            Predicate::new(
                "parent".to_string(),
                vec![Term::Var("X".to_string()), Term::Var("Y".to_string())],
            ),
            Predicate::new(
                "ancestor".to_string(),
                vec![Term::Var("Y".to_string()), Term::Var("Z".to_string())],
            ),
        ],
    ));

    // Query for all ancestors of alice
    let goal = Predicate::new(
        "ancestor".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Var("Z".to_string()),
        ],
    );

    let engine = TabledInferenceEngine::new();
    let solutions = engine.query(&goal, &kb).unwrap();

    // Should find at least bob as an ancestor
    assert!(!solutions.is_empty(), "Should find at least one ancestor");

    // The tabled engine should find bob through the base rule
    // Note: The implementation may vary in how it returns results
    assert!(!solutions.is_empty(), "Should find at least bob");
}

#[tokio::test]
async fn test_distributed_proof_assembly() {
    let mut kb = KnowledgeBase::new();
    kb.add_fact(Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Const(Constant::String("bob".to_string())),
        ],
    ));

    let provider = Arc::new(MockRemoteKnowledgeProvider::new(Arc::new(kb)));
    let mut assembler = DistributedProofAssembler::new(provider);

    let goal = Predicate::new(
        "parent".to_string(),
        vec![
            Term::Const(Constant::String("alice".to_string())),
            Term::Var("X".to_string()),
        ],
    );

    let proof = assembler.assemble_proof(&goal).await.unwrap();
    assert!(proof.is_some());

    let proof = proof.unwrap();
    assert!(proof.is_fact());
    assert_eq!(proof.goal.name, "parent");
}

#[tokio::test]
async fn test_concurrent_goal_resolution() {
    let mut kb = KnowledgeBase::new();

    // Add multiple facts
    for i in 0..10 {
        kb.add_fact(Predicate::new(
            "number".to_string(),
            vec![Term::Const(Constant::Int(i))],
        ));
    }

    let provider = Arc::new(MockRemoteKnowledgeProvider::new(Arc::new(kb)));

    // Create multiple concurrent goals
    let mut handles = vec![];
    for i in 0..5 {
        let provider_clone = provider.clone();
        let goal = Predicate::new("number".to_string(), vec![Term::Const(Constant::Int(i))]);

        let handle = tokio::spawn(async move {
            let mut resolver = DistributedGoalResolver::new(Arc::new(KnowledgeBase::new()))
                .with_provider(provider_clone);
            resolver.resolve(&goal, &Substitution::new()).await
        });

        handles.push(handle);
    }

    // Wait for all to complete
    for handle in handles {
        let result = handle.await.unwrap();
        assert!(result.is_ok());
    }
}