ckg-storage 1.0.1

//! Storage facade over CozoDB. One DB per repo, plus a shared registry DB.
//!
//! DDL is run idempotently on first open: Cozo errors if a relation already
//! exists, and we detect that string and swallow it so re-opens are no-ops.
//!
//! Sub-modules own the implementation of each concern:
//! - `lifecycle`  — open / schema-version check / rebuild / path-swap
//! - `meta`       — boolean Meta sentinels (needs_reindex, index_in_progress)
//! - `insert`     — put_symbols / put_edges
//! - `resolve`    — resolve_cross_file_calls / detect_test_edges
//! - `embed`      — put_embeddings / hnsw_search / iter_embeddings_capped

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use ckg_core::{Error, Result};
use cozo::{DataValue, DbInstance, ScriptMutability};

use self::meta::{read_meta_bool, stamp_meta_bool, stamp_needs_reindex};

mod embed;
mod insert;
mod lifecycle;
mod meta;
mod registry;
mod resolve;

pub use embed::STORAGE_EMBED_DIM;
pub use registry::RegistryStorage;

// ---------------------------------------------------------------------------
// Shared helpers (used by multiple sub-modules via `super::map_err`)
// ---------------------------------------------------------------------------

pub(super) fn map_err(e: impl std::fmt::Display) -> Error {
    Error::Storage(e.to_string())
}

// ---------------------------------------------------------------------------
// Core structs
// ---------------------------------------------------------------------------

/// Per-repo Cozo DB handle.
pub struct Storage {
    pub(super) repo_id: ckg_core::RepoId,
    pub(super) db_path: PathBuf,
    pub(super) db: DbInstance,
}

impl Storage {
    // --- accessors ----------------------------------------------------------

    pub fn repo_id(&self) -> &ckg_core::RepoId {
        &self.repo_id
    }

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

    /// CR-storage-M-7: returns the currently-recorded `Meta.root_path`,
    /// which may differ from the caller's input shape if open_at auto-
    /// migrated a symlink-equivalent path (canonicalize-equal). Callers
    /// that also stamp the path elsewhere (e.g. registry's `Repo.root_path`
    /// via `RegistryStorage::put_repo`) should use this accessor so the
    /// two authoritative shapes stay consistent.
    ///
    /// CR-storage-H3: returns `Result<Option<String>>` so callers can
    /// distinguish "no row recorded" (Ok(None) — fresh DB) from "Meta
    /// read failed" (Err — disk error / corruption). Pre-fix the
    /// `.ok()?` collapse made both cases indistinguishable, mirroring
    /// the same bug pattern that `RootPathProbe::ReadFailed` was
    /// introduced to fix in `lifecycle.rs`.
    pub fn recorded_root_path(&self) -> Result<Option<String>> {
        let rows = self
            .db
            .run_script(
                "?[v] := *Meta{key: \"root_path\", value: v}",
                BTreeMap::new(),
                ScriptMutability::Immutable,
            )
            .map_err(map_err)?;
        Ok(rows.rows.first().and_then(|r| r.first()).and_then(|v| match v {
            DataValue::Str(s) => Some(s.to_string()),
            _ => None,
        }))
    }

    pub fn db(&self) -> &DbInstance {
        &self.db
    }

    // --- script runners -----------------------------------------------------

    /// Run a Cozo script with **mutable** access (`:put`, `:rm`, `:create`
    /// etc. are allowed). Safety relies on Cozo's `ScriptMutability::Mutable`
    /// runtime gate — there is no string-level prefilter.
    ///
    /// **STORAGE-H2 / danger:** The name is intentionally verbose. Callers
    /// must hold an explicit intent to mutate; prefer `Self::run_immutable`
    /// or `Self::run_with_immutable` for all read paths and any
    /// caller-supplied Datalog (e.g. MCP `query` tool). This keeps the
    /// footprint of mutable execution small and auditable.
    pub fn run_mutable_unchecked(&self, script: &str) -> Result<cozo::NamedRows> {
        self.db
            .run_script(script, BTreeMap::new(), ScriptMutability::Mutable)
            .map_err(map_err)
    }

    /// Read-only run — passes `ScriptMutability::Immutable` to Cozo which
    /// rejects scripts containing `:put`, `:rm`, `:create`, `:replace`,
    /// `:ensure_not`, etc. **at execution time** (not string-prefiltered).
    /// Use this for any caller-supplied Datalog (MCP `query` tool) so a
    /// malicious client can't drop or mutate relations.
    pub fn run_immutable(&self, script: &str) -> Result<cozo::NamedRows> {
        self.db
            .run_script(script, BTreeMap::new(), ScriptMutability::Immutable)
            .map_err(map_err)
    }

    /// Mutable run with parameters. Same safety model as `Self::run` — relies
    /// on `ScriptMutability::Mutable` runtime gate, no string prefilter.
    pub fn run_with(
        &self,
        script: &str,
        params: BTreeMap<String, DataValue>,
    ) -> Result<cozo::NamedRows> {
        self.db
            .run_script(script, params, ScriptMutability::Mutable)
            .map_err(map_err)
    }

    /// Read-only variant of `run_with` — caller-supplied params, but the
    /// script is rejected if it contains `:put` / `:rm` / `:create` /
    /// `:replace`. Use for any caller-controlled Datalog so a malicious /
    /// typo'd script can't mutate.
    pub fn run_with_immutable(
        &self,
        script: &str,
        params: BTreeMap<String, DataValue>,
    ) -> Result<cozo::NamedRows> {
        self.db
            .run_script(script, params, ScriptMutability::Immutable)
            .map_err(map_err)
    }

    // --- meta sentinels -----------------------------------------------------

    /// True if this repo was schema-rebuilt and hasn't been re-indexed since.
    /// Set when `Storage::open_*` triggers `rebuild_at_path`, cleared via
    /// `mark_indexed()` after a successful `ckg index` run.
    ///
    /// Returns `false` on any read failure — the sentinel is best-effort UX
    /// guidance, not a correctness gate.
    pub fn needs_reindex(&self) -> bool {
        read_meta_bool(&self.db, "needs_reindex")
    }

    /// CR-I-2: Atomicity sentinel. Stamp `index_in_progress=true` BEFORE the
    /// first `put_symbols` / `put_edges` of a fresh index run. On next
    /// `Storage::open`, if this flag is still set, `needs_reindex` is promoted.
    /// Cleared by `mark_indexed`.
    pub fn mark_index_in_progress(&self) -> Result<()> {
        stamp_meta_bool(&self.db, "index_in_progress", true)
    }

    /// CR-I-2: True if a previous index run started but didn't reach
    /// `mark_indexed`. Mostly internal — `Storage::open_at` checks this on
    /// every open.
    pub fn is_index_in_progress(&self) -> bool {
        read_meta_bool(&self.db, "index_in_progress")
    }

    /// Clear the `needs_reindex` AND `index_in_progress` sentinels after a
    /// successful `ckg index` run. Safe to call when the sentinels are already
    /// absent.
    pub fn mark_indexed(&self) -> Result<()> {
        stamp_needs_reindex(&self.db, false)?;
        stamp_meta_bool(&self.db, "index_in_progress", false)?;
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use ckg_core::{EdgeKind, Kind, RepoId};
    use cozo::DataValue;
    use tempfile::tempdir;

    use crate::store::lifecycle::sweep_rebuild_debris;

    fn make_sym(id: &str, qname: &str) -> ckg_core::Symbol {
        ckg_core::Symbol {
            id: id.into(),
            qname: qname.into(),
            name: qname.split("::").last().unwrap_or(qname).into(),
            kind: Kind::Function,
            file: "x.rs".into(),
            line: 1,
            col: 0,
            is_public: true,
            doc: String::new(),
            hash: "h".into(),
        }
    }

    #[test]
    fn open_and_idempotent() {
        // RocksDB takes an exclusive process-wide lock on the path, so the
        // first instance is dropped before re-opening.
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("aaaaaaaaaaaaaaaaaaaaaaaa").unwrap();
        {
            let _s1 = Storage::open_unverified(id.clone(), dir.path()).unwrap();
        }
        let _s2 = Storage::open_unverified(id, dir.path()).unwrap();
    }

    #[test]
    fn put_symbols_and_edges() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("bbbbbbbbbbbbbbbbbbbbbbbb").unwrap();
        let st = Storage::open_unverified(id, dir.path()).unwrap();
        let syms = vec![make_sym("a", "Foo::bar"), make_sym("b", "Baz::qux")];
        st.put_symbols(&syms).unwrap();
        let edges = vec![ckg_core::Edge {
            kind: EdgeKind::Calls,
            src: "a".into(),
            dst: "b".into(),
            confidence: 1.0,
        }];
        st.put_edges(&edges).unwrap();

        let rows = st.run_mutable_unchecked("?[id] := *Symbol{id}").unwrap();
        assert_eq!(rows.rows.len(), 2);
        let edges = st.run_mutable_unchecked("?[s, d] := *Calls{src: s, dst: d}").unwrap();
        assert_eq!(edges.rows.len(), 1);
    }

    /// CR-I-2: a previous indexer that crashed mid-batch leaves the
    /// `index_in_progress` sentinel set. The next `Storage::open_at`
    /// must promote `needs_reindex=true` so the user is warned, and
    /// clear the in-progress flag so a fresh run can claim it.
    #[test]
    fn crashed_run_promotes_needs_reindex() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("cccccccccccccccccccccccc").unwrap();
        // First open: simulate an indexer that started but crashed
        // before mark_indexed.
        {
            let st = Storage::open_unverified(id.clone(), dir.path()).unwrap();
            st.mark_index_in_progress().unwrap();
            assert!(st.is_index_in_progress());
            // Drop without mark_indexed — simulates crash.
        }
        // Second open: recovery path must fire.
        let st = Storage::open_unverified(id, dir.path()).unwrap();
        assert!(
            st.needs_reindex(),
            "open after crashed run must set needs_reindex=true"
        );
        assert!(
            !st.is_index_in_progress(),
            "open must clear the in-progress flag (one-shot recovery)"
        );
    }

    /// CR-I-2: the happy path — `mark_indexed` clears BOTH sentinels.
    #[test]
    fn mark_indexed_clears_both_sentinels() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("dddddddddddddddddddddddd").unwrap();
        let st = Storage::open_unverified(id, dir.path()).unwrap();
        st.mark_index_in_progress().unwrap();
        st.mark_indexed().unwrap();
        assert!(!st.needs_reindex());
        assert!(!st.is_index_in_progress());
    }

    /// CR-M-7: opening with a symlink-equivalent path (different
    /// byte shape, same canonical resolution) auto-migrates the
    /// recorded `root_path` instead of failing closed. Simulates the
    /// macOS `/var ↔ /private/var` case via a real symlink.
    #[test]
    fn open_at_canonicalize_equal_path_auto_migrates() {
        // Set up a real-repo dir + a symlink to it.
        let outer = tempdir().unwrap();
        let real = outer.path().join("real_repo");
        std::fs::create_dir_all(&real).unwrap();
        let alias = outer.path().join("alias_repo");
        #[cfg(unix)]
        std::os::unix::fs::symlink(&real, &alias).unwrap();
        #[cfg(not(unix))]
        {
            // Skip on platforms where we can't make a symlink without admin.
            return;
        }

        let base = tempdir().unwrap();
        let id = RepoId::try_new("eeeeeeeeeeeeeeeeeeeeeeee").unwrap();

        // First open: stamp the alias path.
        {
            let _ = Storage::open_at(id.clone(), &alias, base.path()).unwrap();
        }
        // Second open: pass the canonical path (real). Pre-M-7 this
        // failed with "RepoId collision detected"; post-M-7 it auto-
        // migrates and succeeds.
        let _ = Storage::open_at(id, &real, base.path())
            .expect("symlink-equivalent path must auto-migrate, not collide");
    }

    /// CR-review-#5: negative path — stored root_path no longer exists
    /// on disk (canonicalize fails), caller's path differs byte-wise.
    /// Auto-migrate must NOT silently accept; the original collision-
    /// detection error must fire.
    #[test]
    fn open_at_canonicalize_fails_falls_through_to_collision() {
        let outer = tempdir().unwrap();
        let real = outer.path().join("real_repo");
        std::fs::create_dir_all(&real).unwrap();

        let base = tempdir().unwrap();
        let id = RepoId::try_new("ffffffffffffffffffffffff").unwrap();

        // First open: stamp `real_repo`.
        {
            let _ = Storage::open_at(id.clone(), &real, base.path()).unwrap();
        }
        // Delete the real path so canonicalize on the stored value now fails.
        std::fs::remove_dir_all(&real).unwrap();

        // Re-open with a different path that DOES exist — but canonicalize
        // of the stored value fails, so paths_canonicalize_equal returns
        // false, and the collision-detection error must surface.
        let other = outer.path().join("other_repo");
        std::fs::create_dir_all(&other).unwrap();
        match Storage::open_at(id, &other, base.path()) {
            Ok(_) => panic!("must fail closed when stored path can't canonicalize"),
            Err(e) => {
                let msg = e.to_string();
                assert!(
                    msg.contains("RepoId collision"),
                    "expected collision error, got: {msg}"
                );
            }
        }
    }

    /// CR-review-#5: negative path — two distinct repos that happen to
    /// hash to the same RepoId must continue to fail closed even with
    /// auto-migrate enabled. Both paths exist but canonicalize to
    /// different absolutes.
    #[test]
    fn open_at_distinct_repos_still_collide() {
        let outer = tempdir().unwrap();
        let repo_a = outer.path().join("repo_a");
        let repo_b = outer.path().join("repo_b");
        std::fs::create_dir_all(&repo_a).unwrap();
        std::fs::create_dir_all(&repo_b).unwrap();

        let base = tempdir().unwrap();
        let id = RepoId::try_new("0123456789abcdef01234567").unwrap();

        {
            let _ = Storage::open_at(id.clone(), &repo_a, base.path()).unwrap();
        }
        match Storage::open_at(id, &repo_b, base.path()) {
            Ok(_) => panic!("two distinct repos must surface collision error"),
            Err(e) => {
                let msg = e.to_string();
                assert!(
                    msg.contains("RepoId collision"),
                    "expected collision error, got: {msg}"
                );
            }
        }
    }

    #[test]
    fn registry_open() {
        let dir = tempdir().unwrap();
        let _r = RegistryStorage::open(dir.path()).unwrap();
    }

    /// I2: opening the same RepoId with a DIFFERENT canonical root_path
    /// must refuse — that's exactly the collision case where two repos
    /// hashed to the same 64-bit RepoId would otherwise silently merge.
    #[test]
    fn open_at_refuses_root_path_mismatch() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("cccccccccccccccccccccccc").unwrap();
        let root_a = std::path::Path::new("/tmp/ckg-test-repo-A");
        let root_b = std::path::Path::new("/tmp/ckg-test-repo-B");
        // First open stamps root_a.
        {
            let _s = Storage::open_at(id.clone(), root_a, dir.path()).unwrap();
        }
        // Re-opening with root_a succeeds.
        {
            let _s = Storage::open_at(id.clone(), root_a, dir.path()).unwrap();
        }
        // Re-opening with root_b refuses (collision detected).
        let res = Storage::open_at(id.clone(), root_b, dir.path());
        match res {
            Ok(_) => panic!("expected collision error, got Ok"),
            Err(e) => assert!(
                e.to_string().contains("RepoId collision detected"),
                "expected collision error; got {e}"
            ),
        }
    }

    /// I2: backward-compat — `Storage::open_unverified` (no root) does NOT
    /// verify, so existing tests / ad-hoc CLI usage stays unbroken.
    #[test]
    fn open_without_root_skips_verification() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("dddddddddddddddddddddddd").unwrap();
        // Stamp via open_at:
        {
            let _s = Storage::open_at(
                id.clone(),
                std::path::Path::new("/tmp/some/root"),
                dir.path(),
            )
            .unwrap();
        }
        // Open without root — must succeed even though Meta has a
        // mismatched-style root recorded; backward-compat path doesn't
        // perform the check.
        let _s = Storage::open_unverified(id, dir.path()).unwrap();
    }

    /// R1 (re-fix): `sweep_rebuild_debris` must NOT delete `<path>.old`
    /// when `<path>` contains partial RocksDB debris (e.g. a stray LOG
    /// file with no `CURRENT`). The old gate (`read_dir empty`) was
    /// satisfied by any single file, destroying the recovery dir from
    /// a prior C3 hard-error path.
    #[test]
    fn sweep_preserves_old_when_path_lacks_current_file() {
        let dir = tempdir().unwrap();
        let workspace = dir
            .path()
            .join("workspace_folders")
            .join("recovery_test_id");
        std::fs::create_dir_all(&workspace).unwrap();
        // Simulate post-crash state: <path> has a stray LOG file (RocksDB
        // debris) but NO CURRENT — the marker of a healthy DB.
        std::fs::write(workspace.join("LOG"), b"stale rocksdb log").unwrap();
        // Recovery dir from a prior failed swap.
        let old_path = workspace.with_extension("old");
        std::fs::create_dir_all(&old_path).unwrap();
        // The healthy old DB has CURRENT.
        std::fs::write(old_path.join("CURRENT"), b"MANIFEST-000001\n").unwrap();
        std::fs::write(old_path.join("MANIFEST-000001"), b"").unwrap();

        sweep_rebuild_debris(&workspace);

        assert!(
            old_path.exists(),
            "sweep destroyed the recovery dir at {} despite <path> being unhealthy",
            old_path.display()
        );
        assert!(
            old_path.join("CURRENT").is_file(),
            "recovery dir contents lost"
        );
    }

    /// R1 inverse: when `<path>` IS a healthy RocksDB (has CURRENT),
    /// `<path>.old` is correctly swept (old behavior preserved).
    #[test]
    fn sweep_removes_old_when_path_has_current() {
        let dir = tempdir().unwrap();
        let workspace = dir.path().join("workspace_folders").join("healthy_path");
        std::fs::create_dir_all(&workspace).unwrap();
        std::fs::write(workspace.join("CURRENT"), b"MANIFEST-000001\n").unwrap();
        let old_path = workspace.with_extension("old");
        std::fs::create_dir_all(&old_path).unwrap();
        std::fs::write(old_path.join("LOG"), b"stale").unwrap();

        sweep_rebuild_debris(&workspace);

        assert!(
            !old_path.exists(),
            "healthy <path> should have triggered .old sweep"
        );
    }

    /// NI1: opening with `open_at` after a Meta read-failure must NOT
    /// silently stamp the caller's path. We verify the stamp-then-verify
    /// cycle across the NoRow / Recorded transitions.
    #[test]
    fn open_at_stamp_then_verify_round_trip() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("01101011001011001011001a").unwrap();
        let root = std::path::Path::new("/tmp/ckg-ni1-test");
        // Open #1: NoRow → stamp.
        {
            let _ = Storage::open_at(id.clone(), root, dir.path()).unwrap();
        }
        // Open #2: Recorded(matches) → ok.
        {
            let _ = Storage::open_at(id.clone(), root, dir.path()).unwrap();
        }
        // Open #3: Recorded(mismatch) → refuse.
        let mismatch_root = std::path::Path::new("/tmp/ckg-ni1-other");
        let res = Storage::open_at(id, mismatch_root, dir.path());
        match res {
            Ok(_) => panic!("should have refused mismatched root"),
            Err(e) => {
                let m = e.to_string();
                assert!(
                    m.contains("RepoId collision detected"),
                    "expected collision error, got: {m}"
                );
            }
        }
    }

    /// I8: `resolve_cross_file_calls` must not regress correctness. A
    /// bare-name dst should still rewrite to the unique Symbol id even
    /// after the narrowed-load refactor.
    #[test]
    fn resolve_cross_file_calls_rewrites_unique_bare_name() {
        let dir = tempdir().unwrap();
        let id = RepoId::try_new("eeeeeeeeeeeeeeeeeeeeeeee").unwrap();
        let st = Storage::open_unverified(id, dir.path()).unwrap();
        // Two symbols with distinct names; one Calls edge with a bare
        // dst that exactly matches the second symbol's name.
        st.put_symbols(&[
            ckg_core::Symbol {
                id: "src_id".into(),
                qname: "Caller::go".into(),
                name: "go".into(),
                kind: Kind::Function,
                file: "a.rs".into(),
                line: 1,
                col: 0,
                is_public: true,
                doc: String::new(),
                hash: "h1".into(),
            },
            ckg_core::Symbol {
                id: "tgt_id".into(),
                qname: "Target::run".into(),
                name: "uniq_target".into(),
                kind: Kind::Function,
                file: "b.rs".into(),
                line: 1,
                col: 0,
                is_public: true,
                doc: String::new(),
                hash: "h2".into(),
            },
        ])
        .unwrap();
        st.put_edges(&[ckg_core::Edge {
            kind: EdgeKind::Calls,
            src: "src_id".into(),
            dst: "uniq_target".into(), // bare name, not yet a Symbol.id
            confidence: 0.5,
        }])
        .unwrap();
        let rewritten = st.resolve_cross_file_calls().unwrap();
        assert_eq!(rewritten, 1, "expected exactly one rewrite");
        let rows = st.run_mutable_unchecked("?[s, d] := *Calls{src: s, dst: d}").unwrap();
        let dst = rows
            .rows
            .first()
            .and_then(|r| r.get(1))
            .and_then(|v| match v {
                DataValue::Str(s) => Some(s.to_string()),
                _ => None,
            })
            .unwrap();
        assert_eq!(dst, "tgt_id", "dst should rewrite to target id");
    }
}