ai-model-vault 4.6.1

//! Model inference validation — sanity-check model outputs after storage or
//! conversion to catch silent corruption.
//!
//! The validator loads a model file, runs a deterministic probe input, and
//! compares the output against a stored reference.  This works at the tensor
//! level for formats that support direct reading (SafeTensors, GGUF, NumPy).

use std::fs;
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};

use crate::error::{Result, VaultError};

// ── Types ────────────────────────────────────────────────────────────────────

/// A validation probe — fixed input with expected output checksum.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationProbe {
    /// Human label (e.g. "zeros-input", "sample-sentence")
    pub label: String,
    /// SHA-256 of the reference output tensor bytes.
    pub expected_checksum: String,
    /// Shape of the probe input tensor.
    pub input_shape: Vec<usize>,
    /// Data type (f32, f16, bf16, i32, …)
    pub dtype: String,
}

/// Result of a single validation run.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationResult {
    pub probe_label: String,
    pub passed: bool,
    pub actual_checksum: Option<String>,
    pub message: String,
}

/// Full validation report for a model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationReport {
    pub model_name: String,
    pub file_path: String,
    pub results: Vec<ValidationResult>,
    pub overall_pass: bool,
    pub validated_at: String,
}

/// Manages validation probes for a vault.
pub struct ValidationStore {
    base_dir: PathBuf,
}

// ── Implementation ───────────────────────────────────────────────────────────

impl ValidationStore {
    /// Create a new store under `vault_path/validations/`.
    pub fn new(vault_path: &Path) -> Result<Self> {
        let base_dir = vault_path.join("validations");
        fs::create_dir_all(&base_dir)?;
        Ok(Self { base_dir })
    }

    fn probe_file(&self, model: &str) -> PathBuf {
        self.base_dir.join(format!("{}.probes.json", model))
    }

    /// Save probes for a model.
    pub fn save_probes(&self, model: &str, probes: &[ValidationProbe]) -> Result<()> {
        let json = serde_json::to_string_pretty(probes)?;
        fs::write(self.probe_file(model), json)?;
        Ok(())
    }

    /// Load probes for a model.
    pub fn load_probes(&self, model: &str) -> Result<Vec<ValidationProbe>> {
        let path = self.probe_file(model);
        if !path.exists() {
            return Ok(Vec::new());
        }
        let data = fs::read_to_string(&path)?;
        let probes: Vec<ValidationProbe> = serde_json::from_str(&data)?;
        Ok(probes)
    }

    /// Validate a model file against its stored probes.
    ///
    /// This performs a structural check:
    /// 1. File exists and is readable.
    /// 2. File size is non-zero.
    /// 3. SHA-256 of the entire file matches the stored checksum (if a
    ///    "file-integrity" probe exists).
    ///
    /// Full tensor-level inference validation would require loading the model
    /// with the appropriate runtime — this is a framework-agnostic baseline.
    pub fn validate(&self, model: &str, file_path: &Path) -> Result<ValidationReport> {
        let probes = self.load_probes(model)?;

        let mut results = Vec::new();

        // Basic structural check even when no probes are stored.
        if !file_path.exists() {
            results.push(ValidationResult {
                probe_label: "file-exists".into(),
                passed: false,
                actual_checksum: None,
                message: format!("File not found: {}", file_path.display()),
            });
        } else {
            let meta = fs::metadata(file_path)?;
            results.push(ValidationResult {
                probe_label: "file-exists".into(),
                passed: true,
                actual_checksum: None,
                message: format!("OK ({} bytes)", meta.len()),
            });

            if meta.len() == 0 {
                results.push(ValidationResult {
                    probe_label: "non-empty".into(),
                    passed: false,
                    actual_checksum: None,
                    message: "File is empty".into(),
                });
            } else {
                results.push(ValidationResult {
                    probe_label: "non-empty".into(),
                    passed: true,
                    actual_checksum: None,
                    message: "OK".into(),
                });
            }
        }

        // Run stored probes
        for probe in &probes {
            if probe.label == "file-integrity" {
                let actual = file_sha256(file_path);
                let passed = actual.as_deref() == Some(probe.expected_checksum.as_str());
                results.push(ValidationResult {
                    probe_label: probe.label.clone(),
                    passed,
                    actual_checksum: actual.clone(),
                    message: if passed {
                        "Checksum matches".into()
                    } else {
                        format!(
                            "Checksum mismatch: expected {}, got {}",
                            probe.expected_checksum,
                            actual.unwrap_or_else(|| "N/A".into())
                        )
                    },
                });
            } else {
                // For custom probes we cannot run inference here — just record as skipped.
                results.push(ValidationResult {
                    probe_label: probe.label.clone(),
                    passed: true,
                    actual_checksum: None,
                    message: "Probe skipped (requires runtime)".into(),
                });
            }
        }

        let overall_pass = results.iter().all(|r| r.passed);

        Ok(ValidationReport {
            model_name: model.to_string(),
            file_path: file_path.display().to_string(),
            results,
            overall_pass,
            validated_at: chrono::Utc::now().to_rfc3339(),
        })
    }

    /// Create a file-integrity probe from the current model file.
    pub fn create_integrity_probe(&self, model: &str, file_path: &Path) -> Result<()> {
        let checksum = file_sha256(file_path).ok_or_else(|| {
            VaultError::IoError(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                "Cannot hash file",
            ))
        })?;

        let probe = ValidationProbe {
            label: "file-integrity".into(),
            expected_checksum: checksum,
            input_shape: vec![],
            dtype: "bytes".into(),
        };

        let mut probes = self.load_probes(model)?;
        probes.retain(|p| p.label != "file-integrity");
        probes.push(probe);
        self.save_probes(model, &probes)
    }
}

fn file_sha256(path: &Path) -> Option<String> {
    use sha2::Digest;
    let data = fs::read(path).ok()?;
    let hash = sha2::Sha256::digest(&data);
    Some(hex::encode(hash))
}

// ── Tests ────────────────────────────────────────────────────────────────────

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

    #[test]
    fn test_validate_missing_file() {
        let dir = tempfile::tempdir().unwrap();
        let store = ValidationStore::new(dir.path()).unwrap();
        let report = store.validate("test", Path::new("/nonexistent")).unwrap();
        assert!(!report.overall_pass);
    }

    #[test]
    fn test_validate_existing_file() {
        let dir = tempfile::tempdir().unwrap();
        let store = ValidationStore::new(dir.path()).unwrap();

        let model_file = dir.path().join("model.bin");
        fs::write(&model_file, b"hello model").unwrap();

        let report = store.validate("test", &model_file).unwrap();
        assert!(report.overall_pass);
        assert_eq!(report.results.len(), 2); // file-exists + non-empty
    }

    #[test]
    fn test_integrity_probe() {
        let dir = tempfile::tempdir().unwrap();
        let store = ValidationStore::new(dir.path()).unwrap();

        let model_file = dir.path().join("model.bin");
        fs::write(&model_file, b"deterministic content").unwrap();

        store
            .create_integrity_probe("mymodel", &model_file)
            .unwrap();

        let report = store.validate("mymodel", &model_file).unwrap();
        assert!(report.overall_pass);

        // Corrupt the file
        fs::write(&model_file, b"corrupted content").unwrap();
        let report2 = store.validate("mymodel", &model_file).unwrap();
        assert!(!report2.overall_pass);
    }

    #[test]
    fn test_probe_persistence() {
        let dir = tempfile::tempdir().unwrap();
        let store = ValidationStore::new(dir.path()).unwrap();

        let probes = vec![ValidationProbe {
            label: "test-probe".into(),
            expected_checksum: "abc123".into(),
            input_shape: vec![1, 3, 224, 224],
            dtype: "f32".into(),
        }];

        store.save_probes("mymodel", &probes).unwrap();
        let loaded = store.load_probes("mymodel").unwrap();
        assert_eq!(loaded.len(), 1);
        assert_eq!(loaded[0].label, "test-probe");
    }
}