ipfrs_tensorlogic/remote_reasoning/

mod.rs

//! Remote Knowledge Retrieval and Distributed Reasoning
//!
//! This module provides protocols and interfaces for distributed reasoning
//! across a network of IPFS nodes. It defines the abstractions needed for:
//!
//! - Remote predicate lookup
//! - Fact discovery from network peers
//! - Incremental fact loading
//! - Distributed goal resolution
//! - Proof assembly from distributed fragments
//!
//! # Architecture
//!
//! The remote reasoning system is designed to work with ipfrs-network once
//! integrated. The traits defined here provide the interface that network
//! implementations will satisfy.
//!
//! ## Example
//!
//! ```ignore
//! use ipfrs_tensorlogic::{RemoteKnowledgeProvider, QueryRequest};
//!
//! // Once ipfrs-network is integrated:
//! let provider = NetworkKnowledgeProvider::new(network_client);
//! let results = provider.query_predicate("parent", vec!["Alice"]).await?;
//! ```

pub mod session;

pub use session::{
    DistributedInferenceSession, DistributedReasonerConfig, DistributedReasonerV2,
    InferenceRequest, InferenceResponse, InferenceResultStream, PartialResult, ReasoningError,
    RemoteResult, SessionMetrics, SessionStats,
};

use crate::ir::{KnowledgeBase, Predicate, Rule, Term};
use crate::proof_storage::{ProofFragment, ProofFragmentRef};
use crate::reasoning::{Proof, Substitution};
use async_trait::async_trait;
use ipfrs_core::Cid;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use thiserror::Error;

/// Errors that can occur during remote reasoning
#[derive(Debug, Error)]
pub enum RemoteReasoningError {
    #[error("Network error: {0}")]
    NetworkError(String),

    #[error("Timeout waiting for remote response")]
    Timeout,

    #[error("Invalid response from peer: {0}")]
    InvalidResponse(String),

    #[error("Peer not found: {0}")]
    PeerNotFound(String),

    #[error("No peers available for query")]
    NoPeersAvailable,

    #[error("Serialization error: {0}")]
    SerializationError(String),

    #[error("Remote query failed: {0}")]
    QueryFailed(String),
}

/// Query request for remote predicate lookup
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryRequest {
    /// Predicate name to query
    pub predicate_name: String,

    /// Ground arguments (constants only)
    pub ground_args: Vec<String>,

    /// Maximum number of results to return
    pub max_results: usize,

    /// Query depth limit
    pub max_depth: usize,

    /// Request ID for tracking
    pub request_id: String,
}

/// Query response containing facts from remote peer
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResponse {
    /// Request ID this response is for
    pub request_id: String,

    /// Predicates matching the query
    pub predicates: Vec<Predicate>,

    /// Rules matching the query
    pub rules: Vec<Rule>,

    /// Proof fragments (if proofs requested)
    pub proof_fragments: Vec<ProofFragmentRef>,

    /// Peer ID that responded
    pub peer_id: String,

    /// Whether more results are available
    pub has_more: bool,

    /// Continuation token for pagination
    pub continuation_token: Option<String>,
}

/// Fact discovery request for network-wide search
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FactDiscoveryRequest {
    /// Predicate name to discover
    pub predicate_name: String,

    /// Optional argument patterns (None = wildcard)
    pub arg_patterns: Vec<Option<String>>,

    /// Maximum hops for multi-hop search
    pub max_hops: usize,

    /// TTL for the request
    pub ttl: u32,

    /// Exclude peers already queried
    pub exclude_peers: HashSet<String>,
}

/// Fact discovery response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FactDiscoveryResponse {
    /// Discovered facts
    pub facts: Vec<Predicate>,

    /// Peer ID that provided each fact
    pub sources: HashMap<usize, String>, // fact index -> peer ID

    /// Number of peers queried
    pub peers_queried: usize,

    /// Hops taken to find facts
    pub hops: HashMap<usize, usize>, // fact index -> hops
}

/// Incremental loading request for streaming facts
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IncrementalLoadRequest {
    /// Predicate name to load
    pub predicate_name: String,

    /// Batch size for incremental loading
    pub batch_size: usize,

    /// Offset for pagination
    pub offset: usize,

    /// Filter criteria (optional)
    pub filter: Option<HashMap<String, String>>,
}

/// Incremental loading response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IncrementalLoadResponse {
    /// Batch of predicates
    pub batch: Vec<Predicate>,

    /// Total count available
    pub total_count: usize,

    /// Next offset for continuation
    pub next_offset: Option<usize>,

    /// Whether this is the last batch
    pub is_last: bool,
}

/// Goal resolution request for distributed solving
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GoalResolutionRequest {
    /// Goal to solve
    pub goal: Predicate,

    /// Current substitution
    pub substitution: HashMap<String, Term>,

    /// Depth in the proof tree
    pub depth: usize,

    /// Requesting peer ID
    pub requester: String,

    /// Request ID for tracking
    pub request_id: String,
}

/// Goal resolution response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GoalResolutionResponse {
    /// Request ID this response is for
    pub request_id: String,

    /// Whether the goal was solved
    pub solved: bool,

    /// Substitutions that solve the goal
    pub solutions: Vec<HashMap<String, Term>>,

    /// Proof (if requested)
    pub proof: Option<Proof>,

    /// Proof fragments for assembly
    pub proof_fragments: Vec<ProofFragmentRef>,
}

/// Trait for remote knowledge retrieval
#[async_trait]
pub trait RemoteKnowledgeProvider: Send + Sync {
    /// Query a predicate from remote peers
    async fn query_predicate(
        &self,
        request: QueryRequest,
    ) -> Result<QueryResponse, RemoteReasoningError>;

    /// Discover facts across the network
    async fn discover_facts(
        &self,
        request: FactDiscoveryRequest,
    ) -> Result<FactDiscoveryResponse, RemoteReasoningError>;

    /// Load facts incrementally
    async fn load_incremental(
        &self,
        request: IncrementalLoadRequest,
    ) -> Result<IncrementalLoadResponse, RemoteReasoningError>;

    /// Resolve a goal using remote peers
    async fn resolve_goal(
        &self,
        request: GoalResolutionRequest,
    ) -> Result<GoalResolutionResponse, RemoteReasoningError>;

    /// Get available peers for querying
    async fn get_available_peers(&self) -> Result<Vec<String>, RemoteReasoningError>;
}

/// Distributed goal resolver
pub struct DistributedGoalResolver {
    /// Local knowledge base
    local_kb: Arc<KnowledgeBase>,

    /// Remote knowledge provider
    remote_provider: Option<Arc<dyn RemoteKnowledgeProvider>>,

    /// Maximum depth for distributed resolution
    max_depth: usize,

    /// Timeout for remote queries (milliseconds)
    timeout_ms: u64,

    /// Cache for remote facts
    remote_fact_cache: HashMap<String, Vec<Predicate>>,
}

impl DistributedGoalResolver {
    /// Create a new distributed goal resolver
    pub fn new(local_kb: Arc<KnowledgeBase>) -> Self {
        Self {
            local_kb,
            remote_provider: None,
            max_depth: 10,
            timeout_ms: 5000,
            remote_fact_cache: HashMap::new(),
        }
    }

    /// Set the remote knowledge provider
    pub fn with_provider(mut self, provider: Arc<dyn RemoteKnowledgeProvider>) -> Self {
        self.remote_provider = Some(provider);
        self
    }

    /// Set maximum depth
    pub fn with_max_depth(mut self, max_depth: usize) -> Self {
        self.max_depth = max_depth;
        self
    }

    /// Set timeout in milliseconds
    pub fn with_timeout(mut self, timeout_ms: u64) -> Self {
        self.timeout_ms = timeout_ms;
        self
    }

    /// Resolve a goal using both local and remote knowledge
    pub async fn resolve(
        &mut self,
        goal: &Predicate,
        substitution: &Substitution,
    ) -> Result<Vec<Substitution>, RemoteReasoningError> {
        // First, try local resolution
        let local_solutions = self.resolve_local(goal, substitution);

        if !local_solutions.is_empty() {
            return Ok(local_solutions);
        }

        // If no local solutions and remote provider available, try remote
        if let Some(provider) = self.remote_provider.clone() {
            let remote_solutions = self.resolve_remote(goal, substitution, &provider).await?;
            Ok(remote_solutions)
        } else {
            Ok(Vec::new())
        }
    }

    /// Resolve a goal locally
    fn resolve_local(&self, goal: &Predicate, _substitution: &Substitution) -> Vec<Substitution> {
        // Check if goal matches any facts in local KB
        let facts = self.local_kb.get_predicates(&goal.name);

        let mut solutions = Vec::new();
        for fact in facts {
            if let Some(subst) =
                crate::reasoning::unify_predicates(goal, fact, &Substitution::new())
            {
                solutions.push(subst);
            }
        }

        solutions
    }

    /// Resolve a goal remotely
    async fn resolve_remote(
        &mut self,
        goal: &Predicate,
        substitution: &Substitution,
        provider: &Arc<dyn RemoteKnowledgeProvider>,
    ) -> Result<Vec<Substitution>, RemoteReasoningError> {
        // Create goal resolution request
        let request = GoalResolutionRequest {
            goal: goal.clone(),
            substitution: substitution.clone(),
            depth: 0,
            requester: "local".to_string(),
            request_id: uuid::Uuid::new_v4().to_string(),
        };

        // Query remote peers
        let response = provider.resolve_goal(request).await?;

        // Convert solutions
        Ok(response.solutions)
    }

    /// Prefetch facts for a predicate from remote peers
    pub async fn prefetch_facts(
        &mut self,
        predicate_name: &str,
    ) -> Result<usize, RemoteReasoningError> {
        let Some(provider) = &self.remote_provider else {
            return Ok(0);
        };

        // Create discovery request
        let request = FactDiscoveryRequest {
            predicate_name: predicate_name.to_string(),
            arg_patterns: Vec::new(),
            max_hops: 3,
            ttl: 30,
            exclude_peers: HashSet::new(),
        };

        // Discover facts
        let response = provider.discover_facts(request).await?;

        // Cache the facts
        let count = response.facts.len();
        self.remote_fact_cache
            .insert(predicate_name.to_string(), response.facts);

        Ok(count)
    }

    /// Get cached remote facts
    pub fn get_cached_facts(&self, predicate_name: &str) -> Option<&[Predicate]> {
        self.remote_fact_cache
            .get(predicate_name)
            .map(|v| v.as_slice())
    }

    /// Clear the remote fact cache
    pub fn clear_cache(&mut self) {
        self.remote_fact_cache.clear();
    }
}

/// Proof assembler for distributed proofs
pub struct DistributedProofAssembler {
    /// Remote knowledge provider
    remote_provider: Arc<dyn RemoteKnowledgeProvider>,

    /// Cache of proof fragments
    fragment_cache: HashMap<Cid, ProofFragment>,

    /// Maximum proof depth
    #[allow(dead_code)]
    max_depth: usize,
}

impl DistributedProofAssembler {
    /// Create a new distributed proof assembler
    pub fn new(remote_provider: Arc<dyn RemoteKnowledgeProvider>) -> Self {
        Self {
            remote_provider,
            fragment_cache: HashMap::new(),
            max_depth: 100,
        }
    }

    /// Assemble a proof from distributed fragments
    pub async fn assemble_proof(
        &mut self,
        goal: &Predicate,
    ) -> Result<Option<Proof>, RemoteReasoningError> {
        // Request goal resolution with proof
        let request = GoalResolutionRequest {
            goal: goal.clone(),
            substitution: HashMap::new(),
            depth: 0,
            requester: "local".to_string(),
            request_id: uuid::Uuid::new_v4().to_string(),
        };

        let response = self.remote_provider.resolve_goal(request).await?;

        if response.solved {
            Ok(response.proof)
        } else {
            Ok(None)
        }
    }

    /// Fetch a proof fragment from the network
    pub async fn fetch_fragment(
        &mut self,
        cid: Cid,
    ) -> Result<ProofFragment, RemoteReasoningError> {
        // Check cache first
        if let Some(fragment) = self.fragment_cache.get(&cid) {
            return Ok(fragment.clone());
        }

        // In a real implementation, this would fetch from IPFS
        // For now, return an error
        Err(RemoteReasoningError::NetworkError(
            "Fragment fetch not yet implemented".to_string(),
        ))
    }
}

/// Mock implementation for testing (will be replaced with network implementation)
pub struct MockRemoteKnowledgeProvider {
    /// Mock knowledge base
    mock_kb: Arc<KnowledgeBase>,
}

impl MockRemoteKnowledgeProvider {
    /// Create a new mock provider
    pub fn new(mock_kb: Arc<KnowledgeBase>) -> Self {
        Self { mock_kb }
    }
}

#[async_trait]
impl RemoteKnowledgeProvider for MockRemoteKnowledgeProvider {
    async fn query_predicate(
        &self,
        request: QueryRequest,
    ) -> Result<QueryResponse, RemoteReasoningError> {
        let predicates = self
            .mock_kb
            .get_predicates(&request.predicate_name)
            .into_iter()
            .take(request.max_results)
            .cloned()
            .collect();

        let rules = self
            .mock_kb
            .get_rules(&request.predicate_name)
            .into_iter()
            .take(request.max_results)
            .cloned()
            .collect();

        Ok(QueryResponse {
            request_id: request.request_id,
            predicates,
            rules,
            proof_fragments: Vec::new(),
            peer_id: "mock_peer".to_string(),
            has_more: false,
            continuation_token: None,
        })
    }

    async fn discover_facts(
        &self,
        request: FactDiscoveryRequest,
    ) -> Result<FactDiscoveryResponse, RemoteReasoningError> {
        let facts: Vec<Predicate> = self
            .mock_kb
            .get_predicates(&request.predicate_name)
            .into_iter()
            .cloned()
            .collect();

        let sources: HashMap<usize, String> = (0..facts.len())
            .map(|i| (i, "mock_peer".to_string()))
            .collect();

        let hops: HashMap<usize, usize> = (0..facts.len()).map(|i| (i, 0)).collect();

        Ok(FactDiscoveryResponse {
            facts,
            sources,
            peers_queried: 1,
            hops,
        })
    }

    async fn load_incremental(
        &self,
        request: IncrementalLoadRequest,
    ) -> Result<IncrementalLoadResponse, RemoteReasoningError> {
        let all_facts: Vec<Predicate> = self
            .mock_kb
            .get_predicates(&request.predicate_name)
            .into_iter()
            .cloned()
            .collect();

        let total_count = all_facts.len();
        let start = request.offset;
        let end = (start + request.batch_size).min(total_count);

        let batch = all_facts[start..end].to_vec();
        let is_last = end >= total_count;
        let next_offset = if is_last { None } else { Some(end) };

        Ok(IncrementalLoadResponse {
            batch,
            total_count,
            next_offset,
            is_last,
        })
    }

    async fn resolve_goal(
        &self,
        request: GoalResolutionRequest,
    ) -> Result<GoalResolutionResponse, RemoteReasoningError> {
        // Simple fact matching
        let facts = self.mock_kb.get_predicates(&request.goal.name);
        let mut solutions = Vec::new();

        for fact in facts {
            if let Some(subst) =
                crate::reasoning::unify_predicates(&request.goal, fact, &Substitution::new())
            {
                solutions.push(subst);
            }
        }

        let solved = !solutions.is_empty();
        let proof = if solved {
            Some(Proof::fact(request.goal.clone()))
        } else {
            None
        };

        Ok(GoalResolutionResponse {
            request_id: request.request_id,
            solved,
            solutions,
            proof,
            proof_fragments: Vec::new(),
        })
    }

    async fn get_available_peers(&self) -> Result<Vec<String>, RemoteReasoningError> {
        Ok(vec!["mock_peer".to_string()])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ir::Constant;

    #[tokio::test]
    async fn test_query_request_serialization() {
        let request = QueryRequest {
            predicate_name: "parent".to_string(),
            ground_args: vec!["Alice".to_string()],
            max_results: 10,
            max_depth: 5,
            request_id: "test_123".to_string(),
        };

        let json = serde_json::to_string(&request).expect("test: should succeed");
        let decoded: QueryRequest = serde_json::from_str(&json).expect("test: should succeed");

        assert_eq!(request.predicate_name, decoded.predicate_name);
        assert_eq!(request.ground_args, decoded.ground_args);
    }

    #[tokio::test]
    async fn test_mock_provider_query() {
        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
            .expect("test: should succeed");
        assert_eq!(response.predicates.len(), 1);
        assert_eq!(response.predicates[0].name, "parent");
    }

    #[tokio::test]
    async fn test_distributed_resolver() {
        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())),
            ],
        ));

        let mut resolver = DistributedGoalResolver::new(Arc::new(local_kb));

        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
            .expect("test: should succeed");
        assert!(!solutions.is_empty());
    }

    #[tokio::test]
    async fn test_fact_discovery() {
        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())),
            ],
        ));
        kb.add_fact(Predicate::new(
            "parent".to_string(),
            vec![
                Term::Const(Constant::String("Bob".to_string())),
                Term::Const(Constant::String("Charlie".to_string())),
            ],
        ));

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

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

        let response = provider
            .discover_facts(request)
            .await
            .expect("test: should succeed");
        assert_eq!(response.facts.len(), 2);
        assert_eq!(response.peers_queried, 1);
    }

    #[tokio::test]
    async fn test_incremental_loading() {
        let mut kb = KnowledgeBase::new();
        for i in 0..10 {
            kb.add_fact(Predicate::new(
                "number".to_string(),
                vec![Term::Const(Constant::Int(i))],
            ));
        }

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

        // Load first batch
        let request = IncrementalLoadRequest {
            predicate_name: "number".to_string(),
            batch_size: 3,
            offset: 0,
            filter: None,
        };

        let response = provider
            .load_incremental(request)
            .await
            .expect("test: should succeed");
        assert_eq!(response.batch.len(), 3);
        assert_eq!(response.total_count, 10);
        assert!(!response.is_last);
        assert_eq!(response.next_offset, Some(3));
    }

    #[tokio::test]
    async fn test_goal_resolution() {
        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: HashMap::new(),
            depth: 0,
            requester: "test".to_string(),
            request_id: "test_123".to_string(),
        };

        let response = provider
            .resolve_goal(request)
            .await
            .expect("test: should succeed");
        assert!(response.solved);
        assert!(!response.solutions.is_empty());
        assert!(response.proof.is_some());
    }
}