lean_ctx/core/context_package/

registry.rs

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};

use super::content::PackageContent;
use super::manifest::PackageManifest;

const INDEX_FILE: &str = "package-index.json";
const PACKAGES_DIR: &str = "packages";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PackageIndex {
    pub schema_version: u32,
    pub updated_at: DateTime<Utc>,
    pub entries: Vec<PackageEntry>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PackageEntry {
    pub name: String,
    pub version: String,
    pub description: String,
    pub installed_at: DateTime<Utc>,
    pub layers: Vec<String>,
    pub sha256: String,
    pub byte_size: u64,
    #[serde(default)]
    pub tags: Vec<String>,
    #[serde(default)]
    pub auto_load: bool,
}

impl PackageIndex {
    fn new() -> Self {
        Self {
            schema_version: crate::core::contracts::CONTEXT_PACKAGE_V1_SCHEMA_VERSION,
            updated_at: Utc::now(),
            entries: Vec::new(),
        }
    }
}

pub struct LocalRegistry {
    root: PathBuf,
}

impl LocalRegistry {
    pub fn open() -> Result<Self, String> {
        let data_dir = crate::core::data_dir::lean_ctx_data_dir()?;
        let root = data_dir.join(PACKAGES_DIR);
        std::fs::create_dir_all(&root).map_err(|e| format!("create packages dir: {e}"))?;
        Ok(Self { root })
    }

    pub fn open_at(root: &Path) -> Result<Self, String> {
        std::fs::create_dir_all(root).map_err(|e| format!("create packages dir: {e}"))?;
        Ok(Self {
            root: root.to_path_buf(),
        })
    }

    pub fn root(&self) -> &Path {
        &self.root
    }

    pub fn install(
        &self,
        manifest: &PackageManifest,
        content: &PackageContent,
    ) -> Result<PathBuf, String> {
        let pkg_dir = self.package_dir(&manifest.name, &manifest.version);
        std::fs::create_dir_all(&pkg_dir).map_err(|e| format!("create package dir: {e}"))?;

        let manifest_json = serde_json::to_string_pretty(manifest).map_err(|e| e.to_string())?;
        atomic_write(&pkg_dir.join("manifest.json"), manifest_json.as_bytes())?;

        let content_json = serde_json::to_string_pretty(content).map_err(|e| e.to_string())?;
        atomic_write(&pkg_dir.join("content.json"), content_json.as_bytes())?;

        let mut index = self.load_index()?;
        index
            .entries
            .retain(|e| !(e.name == manifest.name && e.version == manifest.version));
        index.entries.push(PackageEntry {
            name: manifest.name.clone(),
            version: manifest.version.clone(),
            description: manifest.description.clone(),
            installed_at: Utc::now(),
            layers: manifest
                .layers
                .iter()
                .map(|l| l.as_str().to_string())
                .collect(),
            sha256: manifest.integrity.sha256.clone(),
            byte_size: manifest.integrity.byte_size,
            tags: manifest.tags.clone(),
            auto_load: false,
        });
        index.updated_at = Utc::now();
        self.save_index(&index)?;

        Ok(pkg_dir)
    }

    pub fn remove(&self, name: &str, version: Option<&str>) -> Result<u32, String> {
        let mut index = self.load_index()?;
        let before = index.entries.len();

        let to_remove: Vec<(String, String)> = index
            .entries
            .iter()
            .filter(|e| e.name == name && version.is_none_or(|v| e.version == v))
            .map(|e| (e.name.clone(), e.version.clone()))
            .collect();

        for (n, v) in &to_remove {
            let dir = self.package_dir(n, v);
            if dir.exists() {
                let _ = std::fs::remove_dir_all(&dir);
            }
        }

        index.entries.retain(|e| {
            !to_remove
                .iter()
                .any(|(n, v)| e.name == *n && e.version == *v)
        });

        let removed = (before - index.entries.len()) as u32;
        if removed > 0 {
            index.updated_at = Utc::now();
            self.save_index(&index)?;
        }

        Ok(removed)
    }

    pub fn list(&self) -> Result<Vec<PackageEntry>, String> {
        let index = self.load_index()?;
        Ok(index.entries)
    }

    pub fn get(&self, name: &str, version: Option<&str>) -> Result<Option<PackageEntry>, String> {
        let index = self.load_index()?;
        Ok(index
            .entries
            .into_iter()
            .find(|e| e.name == name && version.is_none_or(|v| e.version == v)))
    }

    pub fn load_package(
        &self,
        name: &str,
        version: &str,
    ) -> Result<(PackageManifest, PackageContent), String> {
        let pkg_dir = self.package_dir(name, version);
        if !pkg_dir.exists() {
            return Err(format!("package {name}@{version} not found"));
        }

        let manifest_json = std::fs::read_to_string(pkg_dir.join("manifest.json"))
            .map_err(|e| format!("read manifest: {e}"))?;
        let content_json = std::fs::read_to_string(pkg_dir.join("content.json"))
            .map_err(|e| format!("read content: {e}"))?;

        let manifest: PackageManifest =
            serde_json::from_str(&manifest_json).map_err(|e| format!("parse manifest: {e}"))?;
        let content: PackageContent =
            serde_json::from_str(&content_json).map_err(|e| format!("parse content: {e}"))?;

        verify_integrity(&manifest, &content)?;

        Ok((manifest, content))
    }

    pub fn set_auto_load(&self, name: &str, version: &str, auto_load: bool) -> Result<(), String> {
        let mut index = self.load_index()?;
        if let Some(entry) = index
            .entries
            .iter_mut()
            .find(|e| e.name == name && e.version == version)
        {
            entry.auto_load = auto_load;
            index.updated_at = Utc::now();
            self.save_index(&index)?;
        } else {
            return Err(format!("package {name}@{version} not found in index"));
        }
        Ok(())
    }

    pub fn auto_load_packages(&self) -> Result<Vec<PackageEntry>, String> {
        let index = self.load_index()?;
        Ok(index.entries.into_iter().filter(|e| e.auto_load).collect())
    }

    pub fn export_to_file(&self, name: &str, version: &str, output: &Path) -> Result<u64, String> {
        let (manifest, content) = self.load_package(name, version)?;

        let bundle = ExportBundle { manifest, content };
        let json = serde_json::to_string_pretty(&bundle).map_err(|e| e.to_string())?;
        let bytes = json.as_bytes();

        atomic_write(output, bytes)?;
        Ok(bytes.len() as u64)
    }

    pub fn import_from_file(&self, path: &Path) -> Result<PackageManifest, String> {
        let json = std::fs::read_to_string(path).map_err(|e| format!("read package file: {e}"))?;
        let bundle: ExportBundle =
            serde_json::from_str(&json).map_err(|e| format!("parse package: {e}"))?;

        bundle.manifest.validate().map_err(|errs| errs.join("; "))?;

        verify_integrity(&bundle.manifest, &bundle.content)?;

        self.install(&bundle.manifest, &bundle.content)?;
        Ok(bundle.manifest)
    }

    fn package_dir(&self, name: &str, version: &str) -> PathBuf {
        self.root.join(format!("{name}-{version}"))
    }

    fn load_index(&self) -> Result<PackageIndex, String> {
        let path = self.root.join(INDEX_FILE);
        if !path.exists() {
            return Ok(PackageIndex::new());
        }
        let json = std::fs::read_to_string(&path).map_err(|e| format!("read index: {e}"))?;
        serde_json::from_str(&json).map_err(|e| format!("parse index: {e}"))
    }

    fn save_index(&self, index: &PackageIndex) -> Result<(), String> {
        let json = serde_json::to_string_pretty(index).map_err(|e| e.to_string())?;
        atomic_write(&self.root.join(INDEX_FILE), json.as_bytes())
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct ExportBundle {
    manifest: PackageManifest,
    content: PackageContent,
}

fn verify_integrity(manifest: &PackageManifest, content: &PackageContent) -> Result<(), String> {
    let canonical = serde_json::to_string(content).map_err(|e| e.to_string())?;
    let content_bytes = canonical.as_bytes();

    let mut h1 = Sha256::new();
    h1.update(content_bytes);
    let actual_content_hash = format!("{:x}", h1.finalize());

    if actual_content_hash != manifest.integrity.content_hash {
        return Err(format!(
            "integrity check failed: content_hash mismatch (expected {}, got {actual_content_hash})",
            manifest.integrity.content_hash
        ));
    }

    let expected_sha256 = {
        let composite = format!(
            "{}:{}:{actual_content_hash}",
            manifest.name, manifest.version
        );
        let mut h2 = Sha256::new();
        h2.update(composite.as_bytes());
        format!("{:x}", h2.finalize())
    };

    if manifest.integrity.sha256 != expected_sha256 {
        return Err(format!(
            "integrity check failed: sha256 mismatch (expected {expected_sha256}, got {})",
            manifest.integrity.sha256
        ));
    }

    if manifest.integrity.byte_size != content_bytes.len() as u64 {
        return Err(format!(
            "integrity check failed: byte_size mismatch (expected {}, got {})",
            manifest.integrity.byte_size,
            content_bytes.len()
        ));
    }

    Ok(())
}

fn atomic_write(path: &Path, data: &[u8]) -> Result<(), String> {
    if path.exists()
        && path
            .symlink_metadata()
            .is_ok_and(|m| m.file_type().is_symlink())
    {
        return Err(format!(
            "refusing to write through symlink: {}",
            path.display()
        ));
    }
    let parent = path.parent().ok_or_else(|| "invalid path".to_string())?;
    let tmp = parent.join(format!(
        ".{}.tmp",
        path.file_name().and_then(|s| s.to_str()).unwrap_or("pkg")
    ));
    std::fs::write(&tmp, data).map_err(|e| format!("write tmp: {e}"))?;
    std::fs::rename(&tmp, path).map_err(|e| format!("rename: {e}"))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::context_package::manifest::{CompatibilitySpec, PackageStats};

    #[test]
    fn registry_round_trip() {
        let dir = tempfile::tempdir().unwrap();
        let reg = LocalRegistry::open_at(dir.path()).unwrap();

        assert!(reg.list().unwrap().is_empty());

        let manifest = PackageManifest {
            schema_version: crate::core::contracts::CONTEXT_PACKAGE_V1_SCHEMA_VERSION,
            name: "test-pkg".into(),
            version: "1.0.0".into(),
            description: "test".into(),
            author: None,
            created_at: Utc::now(),
            updated_at: None,
            layers: vec![super::super::manifest::PackageLayer::Knowledge],
            dependencies: vec![],
            tags: vec!["rust".into()],
            integrity: {
                let c = PackageContent::default();
                let j = serde_json::to_string(&c).unwrap();
                let mut h = Sha256::new();
                h.update(j.as_bytes());
                let ch = format!("{:x}", h.finalize());
                let composite = format!("test-pkg:1.0.0:{ch}");
                let mut h2 = Sha256::new();
                h2.update(composite.as_bytes());
                let sha = format!("{:x}", h2.finalize());
                super::super::manifest::PackageIntegrity {
                    sha256: sha,
                    content_hash: ch,
                    byte_size: j.len() as u64,
                }
            },
            provenance: super::super::manifest::PackageProvenance {
                tool: "lean-ctx".into(),
                tool_version: "0.0.0".into(),
                project_hash: None,
                source_session_id: None,
            },
            compatibility: CompatibilitySpec::default(),
            stats: PackageStats::default(),
        };

        let content = PackageContent::default();

        reg.install(&manifest, &content).unwrap();
        let list = reg.list().unwrap();
        assert_eq!(list.len(), 1);
        assert_eq!(list[0].name, "test-pkg");

        let (loaded_m, _loaded_c) = reg.load_package("test-pkg", "1.0.0").unwrap();
        assert_eq!(loaded_m.name, "test-pkg");

        let removed = reg.remove("test-pkg", None).unwrap();
        assert_eq!(removed, 1);
        assert!(reg.list().unwrap().is_empty());
    }

    #[test]
    fn export_import_round_trip() {
        let dir = tempfile::tempdir().unwrap();
        let reg = LocalRegistry::open_at(dir.path()).unwrap();

        let content = PackageContent::default();
        let content_json = serde_json::to_string(&content).unwrap();
        let mut h = Sha256::new();
        h.update(content_json.as_bytes());
        let content_hash = format!("{:x}", h.finalize());

        let manifest = PackageManifest {
            schema_version: crate::core::contracts::CONTEXT_PACKAGE_V1_SCHEMA_VERSION,
            name: "export-test".into(),
            version: "2.0.0".into(),
            description: "round trip test".into(),
            author: Some("test".into()),
            created_at: Utc::now(),
            updated_at: None,
            layers: vec![super::super::manifest::PackageLayer::Knowledge],
            dependencies: vec![],
            tags: vec![],
            integrity: {
                let composite = format!("export-test:2.0.0:{content_hash}");
                let mut h2 = Sha256::new();
                h2.update(composite.as_bytes());
                super::super::manifest::PackageIntegrity {
                    sha256: format!("{:x}", h2.finalize()),
                    content_hash,
                    byte_size: content_json.len() as u64,
                }
            },
            provenance: super::super::manifest::PackageProvenance {
                tool: "lean-ctx".into(),
                tool_version: "0.0.0".into(),
                project_hash: None,
                source_session_id: None,
            },
            compatibility: CompatibilitySpec::default(),
            stats: PackageStats::default(),
        };

        reg.install(&manifest, &content).unwrap();

        let export_path = dir.path().join("test.lctxpkg");
        let bytes = reg
            .export_to_file("export-test", "2.0.0", &export_path)
            .unwrap();
        assert!(bytes > 0);

        let reg2 = LocalRegistry::open_at(&dir.path().join("other")).unwrap();
        let imported = reg2.import_from_file(&export_path).unwrap();
        assert_eq!(imported.name, "export-test");
        assert_eq!(reg2.list().unwrap().len(), 1);
    }
}