ipfrs-semantic 0.2.0

//! Provenance Tracking for Embeddings
//!
//! This module tracks the provenance and lineage of embeddings:
//! - Source tracking for embedding generation
//! - Version control for embeddings
//! - Immutable audit trails
//! - Explanation generation for search results

use ipfrs_core::{Cid, Error, Result};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

/// Source of an embedding
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum EmbeddingSource {
    /// Generated by a machine learning model
    Model {
        /// Model name
        name: String,
        /// Model version
        version: String,
        /// Model parameters/config
        config: HashMap<String, String>,
    },
    /// Manually created
    Manual {
        /// Creator identifier
        creator: String,
        /// Description
        description: String,
    },
    /// Derived from another embedding
    Derived {
        /// Source embedding CID (as string)
        #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
        source_cid: Cid,
        /// Transformation applied
        transformation: String,
    },
    /// Aggregated from multiple embeddings
    Aggregated {
        /// Source embedding CIDs (as strings)
        #[serde(
            serialize_with = "serialize_cid_vec",
            deserialize_with = "deserialize_cid_vec"
        )]
        source_cids: Vec<Cid>,
        /// Aggregation method
        method: String,
    },
}

/// Helper function to serialize a CID
fn serialize_cid<S>(cid: &Cid, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    serializer.serialize_str(&cid.to_string())
}

/// Helper function to deserialize a CID
fn deserialize_cid<'de, D>(deserializer: D) -> std::result::Result<Cid, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let s = String::deserialize(deserializer)?;
    s.parse().map_err(serde::de::Error::custom)
}

/// Helper function to serialize a Vec<CID>
fn serialize_cid_vec<S>(cids: &[Cid], serializer: S) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    let strings: Vec<String> = cids.iter().map(|c| c.to_string()).collect();
    strings.serialize(serializer)
}

/// Helper function to deserialize a Vec<CID>
fn deserialize_cid_vec<'de, D>(deserializer: D) -> std::result::Result<Vec<Cid>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let strings: Vec<String> = Vec::deserialize(deserializer)?;
    strings
        .into_iter()
        .map(|s| s.parse().map_err(serde::de::Error::custom))
        .collect()
}

/// Embedding metadata with provenance
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingMetadata {
    /// Content identifier
    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
    pub cid: Cid,
    /// Embedding version
    pub version: u32,
    /// Source of embedding
    pub source: EmbeddingSource,
    /// Creation timestamp (Unix epoch ms)
    pub created_at: u64,
    /// Input data reference
    pub input_reference: Option<String>,
    /// Embedding dimension
    pub dimension: usize,
    /// Additional metadata
    pub extra: HashMap<String, String>,
}

impl EmbeddingMetadata {
    /// Create new embedding metadata
    pub fn new(cid: Cid, dimension: usize, source: EmbeddingSource) -> Self {
        Self {
            cid,
            version: 1,
            source,
            created_at: current_timestamp_ms(),
            input_reference: None,
            dimension,
            extra: HashMap::new(),
        }
    }

    /// Set input reference
    pub fn with_input_reference(mut self, reference: impl Into<String>) -> Self {
        self.input_reference = Some(reference.into());
        self
    }

    /// Add extra metadata
    pub fn with_extra(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.extra.insert(key.into(), value.into());
        self
    }
}

/// Audit log entry for embedding operations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuditLogEntry {
    /// Entry ID
    pub id: u64,
    /// Timestamp (Unix epoch ms)
    pub timestamp: u64,
    /// Operation type
    pub operation: AuditOperation,
    /// CID affected
    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
    pub cid: Cid,
    /// User/system identifier
    pub actor: String,
    /// Additional context
    pub context: HashMap<String, String>,
}

/// Type of audit operation
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum AuditOperation {
    /// Embedding created
    Create,
    /// Embedding updated
    Update,
    /// Embedding deleted
    Delete,
    /// Embedding queried
    Query,
    /// Embedding accessed
    Access,
}

/// Version history for an embedding
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VersionHistory {
    /// Embedding CID
    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
    pub cid: Cid,
    /// All versions
    pub versions: Vec<EmbeddingVersion>,
}

/// A single version of an embedding
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingVersion {
    /// Version number
    pub version: u32,
    /// Timestamp
    pub timestamp: u64,
    /// Change description
    pub change_log: String,
    /// Previous version CID (if updated)
    #[serde(
        skip_serializing_if = "Option::is_none",
        serialize_with = "serialize_cid_option",
        deserialize_with = "deserialize_cid_option"
    )]
    pub previous_cid: Option<Cid>,
    /// Metadata for this version
    pub metadata: EmbeddingMetadata,
}

/// Helper function to serialize an Option<CID>
fn serialize_cid_option<S>(cid: &Option<Cid>, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    match cid {
        Some(c) => serializer.serialize_some(&c.to_string()),
        None => serializer.serialize_none(),
    }
}

/// Helper function to deserialize an Option<CID>
fn deserialize_cid_option<'de, D>(deserializer: D) -> std::result::Result<Option<Cid>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let opt: Option<String> = Option::deserialize(deserializer)?;
    opt.map(|s| s.parse().map_err(serde::de::Error::custom))
        .transpose()
}

/// Explanation for a search result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchExplanation {
    /// Query CID
    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
    pub query_cid: Cid,
    /// Result CID
    #[serde(serialize_with = "serialize_cid", deserialize_with = "deserialize_cid")]
    pub result_cid: Cid,
    /// Similarity score
    pub score: f32,
    /// Feature attributions (which features contributed to similarity)
    pub attributions: Vec<FeatureAttribution>,
    /// Explanation text
    pub explanation: String,
}

/// Attribution of a feature to similarity
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureAttribution {
    /// Feature index
    pub feature_idx: usize,
    /// Contribution to score (-1.0 to 1.0)
    pub contribution: f32,
    /// Feature description
    pub description: Option<String>,
}

/// Provenance tracker for embeddings
pub struct ProvenanceTracker {
    /// Embedding metadata storage
    metadata: Arc<RwLock<HashMap<Cid, EmbeddingMetadata>>>,
    /// Version history
    versions: Arc<RwLock<HashMap<Cid, VersionHistory>>>,
    /// Audit log (immutable append-only)
    audit_log: Arc<RwLock<Vec<AuditLogEntry>>>,
    /// Next audit log ID
    next_audit_id: Arc<RwLock<u64>>,
}

impl ProvenanceTracker {
    /// Create a new provenance tracker
    pub fn new() -> Self {
        Self {
            metadata: Arc::new(RwLock::new(HashMap::new())),
            versions: Arc::new(RwLock::new(HashMap::new())),
            audit_log: Arc::new(RwLock::new(Vec::new())),
            next_audit_id: Arc::new(RwLock::new(0)),
        }
    }

    /// Track a new embedding
    pub fn track_embedding(&self, metadata: EmbeddingMetadata) -> Result<()> {
        let cid = metadata.cid;

        // Store metadata
        self.metadata.write().insert(cid, metadata.clone());

        // Initialize version history
        let version = EmbeddingVersion {
            version: 1,
            timestamp: current_timestamp_ms(),
            change_log: "Initial version".to_string(),
            previous_cid: None,
            metadata: metadata.clone(),
        };

        let history = VersionHistory {
            cid,
            versions: vec![version],
        };

        self.versions.write().insert(cid, history);

        // Add audit log entry
        self.add_audit_entry(
            AuditOperation::Create,
            cid,
            "system".to_string(),
            HashMap::new(),
        )?;

        Ok(())
    }

    /// Update embedding metadata (creates new version)
    pub fn update_embedding(
        &self,
        cid: Cid,
        new_cid: Cid,
        change_log: impl Into<String>,
    ) -> Result<()> {
        let metadata = self
            .metadata
            .read()
            .get(&cid)
            .cloned()
            .ok_or_else(|| Error::InvalidInput(format!("Embedding not found: {}", cid)))?;

        // Create new version
        let mut new_metadata = metadata.clone();
        new_metadata.cid = new_cid;
        new_metadata.version += 1;
        new_metadata.created_at = current_timestamp_ms();

        // Store new metadata
        self.metadata.write().insert(new_cid, new_metadata.clone());

        // Update version history
        let mut versions = self.versions.write();
        let history = versions.entry(cid).or_insert_with(|| VersionHistory {
            cid,
            versions: Vec::new(),
        });

        history.versions.push(EmbeddingVersion {
            version: new_metadata.version,
            timestamp: new_metadata.created_at,
            change_log: change_log.into(),
            previous_cid: Some(cid),
            metadata: new_metadata,
        });

        // Add audit log entry
        drop(versions);
        self.add_audit_entry(
            AuditOperation::Update,
            new_cid,
            "system".to_string(),
            HashMap::from([("previous_cid".to_string(), cid.to_string())]),
        )?;

        Ok(())
    }

    /// Get embedding metadata
    pub fn get_metadata(&self, cid: &Cid) -> Option<EmbeddingMetadata> {
        self.metadata.read().get(cid).cloned()
    }

    /// Get version history
    pub fn get_version_history(&self, cid: &Cid) -> Option<VersionHistory> {
        self.versions.read().get(cid).cloned()
    }

    /// Get audit log entries for a CID
    pub fn get_audit_log(&self, cid: &Cid) -> Vec<AuditLogEntry> {
        self.audit_log
            .read()
            .iter()
            .filter(|e| &e.cid == cid)
            .cloned()
            .collect()
    }

    /// Get all audit log entries (for compliance/export)
    pub fn get_full_audit_log(&self) -> Vec<AuditLogEntry> {
        self.audit_log.read().clone()
    }

    /// Add an audit log entry
    fn add_audit_entry(
        &self,
        operation: AuditOperation,
        cid: Cid,
        actor: String,
        context: HashMap<String, String>,
    ) -> Result<()> {
        let id = {
            let mut next_id = self.next_audit_id.write();
            let id = *next_id;
            *next_id += 1;
            id
        };

        let entry = AuditLogEntry {
            id,
            timestamp: current_timestamp_ms(),
            operation,
            cid,
            actor,
            context,
        };

        self.audit_log.write().push(entry);
        Ok(())
    }

    /// Log a query operation
    pub fn log_query(&self, query_cid: Cid, result_cids: &[Cid]) -> Result<()> {
        let context = HashMap::from([("result_count".to_string(), result_cids.len().to_string())]);

        self.add_audit_entry(
            AuditOperation::Query,
            query_cid,
            "system".to_string(),
            context,
        )?;

        // Log access to each result
        for result_cid in result_cids {
            self.add_audit_entry(
                AuditOperation::Access,
                *result_cid,
                "query".to_string(),
                HashMap::from([("query_cid".to_string(), query_cid.to_string())]),
            )?;
        }

        Ok(())
    }

    /// Generate explanation for a search result
    pub fn explain_result(
        &self,
        query_embedding: &[f32],
        result_cid: &Cid,
        result_embedding: &[f32],
        score: f32,
    ) -> SearchExplanation {
        // Calculate feature attributions
        let mut attributions = Vec::new();

        for (idx, (q, r)) in query_embedding
            .iter()
            .zip(result_embedding.iter())
            .enumerate()
        {
            let contribution = q * r; // Simplified: dot product contribution
            if contribution.abs() > 0.1 {
                // Only include significant contributions
                attributions.push(FeatureAttribution {
                    feature_idx: idx,
                    contribution,
                    description: Some(format!("Dimension {}", idx)),
                });
            }
        }

        // Sort by absolute contribution
        attributions.sort_by(|a, b| {
            b.contribution
                .abs()
                .partial_cmp(&a.contribution.abs())
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Keep top 10 features
        attributions.truncate(10);

        // Generate explanation text
        let explanation = if attributions.is_empty() {
            format!("Result matched with similarity score {:.3}", score)
        } else {
            let top_features: Vec<String> = attributions
                .iter()
                .take(3)
                .map(|a| format!("dim {}: {:.3}", a.feature_idx, a.contribution))
                .collect();

            format!(
                "Result matched with similarity score {:.3}. Top contributing features: {}",
                score,
                top_features.join(", ")
            )
        };

        SearchExplanation {
            query_cid: Cid::default(), // Placeholder
            result_cid: *result_cid,
            score,
            attributions,
            explanation,
        }
    }

    /// Get statistics about tracked embeddings
    pub fn stats(&self) -> ProvenanceStats {
        let metadata = self.metadata.read();
        let versions = self.versions.read();
        let audit_log = self.audit_log.read();

        ProvenanceStats {
            total_embeddings: metadata.len(),
            total_versions: versions.values().map(|h| h.versions.len()).sum(),
            total_audit_entries: audit_log.len(),
            oldest_timestamp: audit_log.first().map(|e| e.timestamp),
            newest_timestamp: audit_log.last().map(|e| e.timestamp),
        }
    }
}

impl Default for ProvenanceTracker {
    fn default() -> Self {
        Self::new()
    }
}

/// Statistics about provenance tracking
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProvenanceStats {
    /// Total number of tracked embeddings
    pub total_embeddings: usize,
    /// Total number of versions across all embeddings
    pub total_versions: usize,
    /// Total audit log entries
    pub total_audit_entries: usize,
    /// Oldest audit entry timestamp
    pub oldest_timestamp: Option<u64>,
    /// Newest audit entry timestamp
    pub newest_timestamp: Option<u64>,
}

/// Get current timestamp in milliseconds
fn current_timestamp_ms() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .expect("system time is after UNIX epoch")
        .as_millis() as u64
}

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

    fn test_cid() -> Cid {
        "bafybeigdyrzt5sfp7udm7hu76uh7y26nf3efuylqabf3oclgtqy55fbzdi"
            .parse()
            .expect("test: hardcoded CID string is valid")
    }

    fn test_cid2() -> Cid {
        "bafybeiczsscdsbs7ffqz55asqdf3smv6klcw3gofszvwlyarci47bgf354"
            .parse()
            .expect("test: hardcoded CID string is valid")
    }

    #[test]
    fn test_track_embedding() {
        let tracker = ProvenanceTracker::new();

        let metadata = EmbeddingMetadata::new(
            test_cid(),
            768,
            EmbeddingSource::Model {
                name: "bert-base".to_string(),
                version: "1.0".to_string(),
                config: HashMap::new(),
            },
        );

        assert!(tracker.track_embedding(metadata).is_ok());

        let retrieved = tracker.get_metadata(&test_cid());
        assert!(retrieved.is_some());
        assert_eq!(
            retrieved
                .expect("test: metadata was just inserted so it must exist")
                .dimension,
            768
        );
    }

    #[test]
    fn test_version_history() {
        let tracker = ProvenanceTracker::new();

        let metadata = EmbeddingMetadata::new(
            test_cid(),
            768,
            EmbeddingSource::Manual {
                creator: "test".to_string(),
                description: "test embedding".to_string(),
            },
        );

        tracker
            .track_embedding(metadata)
            .expect("test: track_embedding should succeed");

        // Update embedding
        tracker
            .update_embedding(test_cid(), test_cid2(), "Updated embedding")
            .expect("test: update_embedding should succeed");

        let history = tracker.get_version_history(&test_cid());
        assert!(history.is_some());
        assert_eq!(
            history
                .expect("test: version history was just created")
                .versions
                .len(),
            2
        );
    }

    #[test]
    fn test_audit_log() {
        let tracker = ProvenanceTracker::new();

        let metadata = EmbeddingMetadata::new(
            test_cid(),
            768,
            EmbeddingSource::Derived {
                source_cid: test_cid2(),
                transformation: "normalize".to_string(),
            },
        );

        tracker
            .track_embedding(metadata)
            .expect("test: track_embedding should succeed");

        let audit_entries = tracker.get_audit_log(&test_cid());
        assert!(!audit_entries.is_empty());
        assert_eq!(audit_entries[0].operation, AuditOperation::Create);
    }

    #[test]
    fn test_log_query() {
        let tracker = ProvenanceTracker::new();

        let result_cids = vec![test_cid(), test_cid2()];
        let query_cid = test_cid();

        assert!(tracker.log_query(query_cid, &result_cids).is_ok());

        let audit_log = tracker.get_full_audit_log();
        assert_eq!(audit_log.len(), 3); // 1 query + 2 access
    }

    #[test]
    fn test_explain_result() {
        let tracker = ProvenanceTracker::new();

        let query_emb = vec![1.0, 0.5, 0.3];
        let result_emb = vec![0.9, 0.6, 0.2];

        let explanation = tracker.explain_result(&query_emb, &test_cid(), &result_emb, 0.95);

        assert_eq!(explanation.result_cid, test_cid());
        assert_eq!(explanation.score, 0.95);
        assert!(!explanation.attributions.is_empty());
        assert!(!explanation.explanation.is_empty());
    }

    #[test]
    fn test_provenance_stats() {
        let tracker = ProvenanceTracker::new();

        let metadata1 = EmbeddingMetadata::new(
            test_cid(),
            768,
            EmbeddingSource::Manual {
                creator: "test".to_string(),
                description: "test".to_string(),
            },
        );

        let metadata2 = EmbeddingMetadata::new(
            test_cid2(),
            512,
            EmbeddingSource::Manual {
                creator: "test".to_string(),
                description: "test2".to_string(),
            },
        );

        tracker
            .track_embedding(metadata1)
            .expect("test: first track_embedding should succeed");
        tracker
            .track_embedding(metadata2)
            .expect("test: second track_embedding should succeed");

        let stats = tracker.stats();
        assert_eq!(stats.total_embeddings, 2);
        assert_eq!(stats.total_versions, 2);
        assert_eq!(stats.total_audit_entries, 2);
    }
}