tokenix 0.67.2

//! Shareable index snapshots.
//!
//! Indexing is the one expensive thing tokenix does, and today every developer
//! on a team pays it in full on a repo where the answer is identical for all of
//! them. A snapshot is the index itself — compacted, stripped of the local
//! legacy embedding tables, gzipped — small enough to commit next to the code, so the
//! second person onward bootstraps in seconds and only indexes the diff.
//!
//! Two rules keep this honest:
//! * A snapshot never overwrites a *newer* local index without `--force`.
//! * Import stamps nothing: the snapshot's own `git_fingerprint` travels with
//!   it, so `index_staleness` immediately reports the gap between the snapshot's
//!   commit and the local checkout, and `tokenix index` closes it incrementally.

use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use rusqlite::Connection;

use crate::store;

/// Ceiling on the decompressed size of an imported snapshot (2 GiB). The default
/// snapshot path sits inside the repo, so a clone controls it: without a cap a
/// tiny gzip bomb fills the disk. Real indexes with vectors stay far below this.
pub const MAX_IMPORT_BYTES: u64 = 2 * 1024 * 1024 * 1024;

/// Decompress `reader` into `writer`, failing once output would exceed `cap`.
fn gunzip_capped<R: Read, W: Write>(reader: R, mut writer: W, cap: u64) -> std::io::Result<u64> {
    let mut limited = GzDecoder::new(reader).take(cap + 1);
    let n = std::io::copy(&mut limited, &mut writer)?;
    if n > cap {
        return Err(std::io::Error::other(format!(
            "decompressed snapshot exceeds the {cap}-byte limit"
        )));
    }
    Ok(n)
}

/// Where a team snapshot lives by default: inside the repo, next to the code it
/// describes, so it can be committed and reviewed like any other artifact.
pub const DEFAULT_SNAPSHOT_REL: &str = ".tokenix/index.db.gz";

pub struct ExportReport {
    pub output: PathBuf,
    pub bytes: u64,
    pub files: i64,
    pub chunks: i64,
    pub head: Option<String>,
    /// Chunks whose content had a known secret masked on the way into the
    /// snapshot. Non-zero means the *index* still holds the plaintext.
    pub redacted: usize,
}

#[derive(Debug)]
pub struct ImportReport {
    pub source: PathBuf,
    pub target: PathBuf,
    pub backup: Option<PathBuf>,
    pub files: i64,
    pub chunks: i64,
    pub created_by: Option<String>,
    pub replaced_newer: bool,
}

fn sql_literal(path: &Path) -> String {
    format!("'{}'", path.to_string_lossy().replace('\'', "''"))
}

fn counts(conn: &Connection) -> (i64, i64) {
    let one = |sql: &str| conn.query_row(sql, [], |r| r.get::<_, i64>(0)).unwrap_or(0);
    (
        one("SELECT COUNT(*) FROM files"),
        one("SELECT COUNT(*) FROM chunks"),
    )
}

fn indexed_at(conn: &Connection) -> f64 {
    store::meta_value(conn, "indexed_at")
        .and_then(|v| v.parse::<f64>().ok())
        .unwrap_or(0.0)
}

/// Write a compressed, self-contained copy of this repo's index.
pub fn export(repo_root: &Path, output: Option<&Path>) -> Result<ExportReport> {
    let source = store::db_path(repo_root);
    if !source.exists() {
        anyhow::bail!("no index for this repo yet. Run `tokenix index` first.");
    }
    let output = output
        .map(|p| p.to_path_buf())
        .unwrap_or_else(|| repo_root.join(DEFAULT_SNAPSHOT_REL));
    if let Some(parent) = output.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("cannot create {}", parent.display()))?;
    }

    // VACUUM INTO gives a compact, WAL-free copy without blocking on the live DB
    // being mid-write, and it refuses to overwrite, so clear any stale temp.
    let staging = output.with_extension(format!("staging-{}", std::process::id()));
    let _ = std::fs::remove_file(&staging);

    let redacted;
    let (files, chunks, head) = {
        let conn = Connection::open(&source)
            .with_context(|| format!("cannot open index at {}", source.display()))?;
        conn.execute_batch(&format!("VACUUM INTO {}", sql_literal(&staging)))
            .context("failed to snapshot the index (VACUUM INTO)")?;

        let staged = Connection::open(&staging)?;
        // Indexes built before embeddings were removed still carry vector tables; they
        // are dead weight, so they never ship.
        staged.execute_batch(
            "DROP TABLE IF EXISTS embeddings; DROP TABLE IF EXISTS embedding_cache;",
        )?;
        // `chunks.content` is verbatim source unless the repo opted into
        // `[index] redact_secrets`, which defaults to false. A snapshot is a
        // file this tool tells you to commit and share, so it gets scrubbed on
        // the way out regardless of that setting — a hardcoded key inside a
        // gzipped blob nobody diffs is exactly how one gets published.
        redacted = redact_staged_chunks(&staged)?;
        store::set_meta(&staged, "snapshot_version", crate::VERSION)?;
        store::set_meta(
            &staged,
            "snapshot_created_at",
            &chrono::Utc::now().to_rfc3339(),
        )?;
        // VACUUM again so the deleted cache pages actually leave the file.
        staged.execute_batch("VACUUM;")?;

        let (files, chunks) = counts(&staged);
        let head = store::meta_value(&staged, "git_fingerprint");
        (files, chunks, head)
    };

    let mut reader = BufReader::new(File::open(&staging)?);
    let out_file = File::create(&output)
        .with_context(|| format!("cannot write snapshot to {}", output.display()))?;
    let mut encoder = GzEncoder::new(BufWriter::new(out_file), Compression::default());
    std::io::copy(&mut reader, &mut encoder)?;
    encoder.finish()?.into_inner()?.sync_all()?;
    let _ = std::fs::remove_file(&staging);

    let bytes = std::fs::metadata(&output).map(|m| m.len()).unwrap_or(0);
    Ok(ExportReport {
        output,
        bytes,
        files,
        chunks,
        head,
        redacted,
    })
}

/// Mask known secrets in every chunk of the staged snapshot copy. Returns the
/// number of chunks changed. Operates on the staging database only — the live
/// index is untouched.
fn redact_staged_chunks(staged: &Connection) -> Result<usize> {
    let rows: Vec<(i64, String)> = {
        let mut stmt = staged.prepare("SELECT id, content FROM chunks")?;
        let mapped = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
        mapped.collect::<std::result::Result<Vec<_>, _>>()?
    };
    let mut changed = 0usize;
    let mut update = staged.prepare("UPDATE chunks SET content = ?1 WHERE id = ?2")?;
    for (id, content) in rows {
        let (clean, hit) = crate::secrets_scan::redact_known_secrets(&content);
        if hit {
            update.execute(rusqlite::params![clean, id])?;
            changed += 1;
        }
    }
    Ok(changed)
}

/// Restore a snapshot as this repo's index.
pub fn import(repo_root: &Path, input: Option<&Path>, force: bool) -> Result<ImportReport> {
    let source = input
        .map(|p| p.to_path_buf())
        .unwrap_or_else(|| repo_root.join(DEFAULT_SNAPSHOT_REL));
    if !source.exists() {
        anyhow::bail!("snapshot not found: {}", source.display());
    }

    let target = store::db_path(repo_root);
    if let Some(parent) = target.parent() {
        std::fs::create_dir_all(parent)?;
    }
    // Serialize with `tokenix index` so the swap never races a live writer.
    let _lock = store::acquire_index_lock(repo_root)?;
    let staging = target.with_extension(format!("import-{}", std::process::id()));
    let _ = std::fs::remove_file(&staging);

    {
        let mut out = BufWriter::new(File::create(&staging)?);
        let copied = gunzip_capped(
            BufReader::new(File::open(&source)?),
            &mut out,
            MAX_IMPORT_BYTES,
        )
        .and_then(|_| out.flush());
        if let Err(e) = copied {
            drop(out);
            let _ = std::fs::remove_file(&staging);
            return Err(anyhow::Error::from(e)
                .context(format!("{} is not a valid gzip snapshot", source.display())));
        }
    }

    let (files, chunks, created_by, snapshot_at) = {
        let staged = Connection::open(&staging)
            .with_context(|| format!("{} is not a readable index", source.display()))?;
        // A snapshot is clone-controlled input: verify page structure before
        // it can replace a working index.
        let integrity: String = staged
            .query_row("PRAGMA integrity_check", [], |r| r.get(0))
            .unwrap_or_else(|e| e.to_string());
        if integrity != "ok" {
            drop(staged);
            let _ = std::fs::remove_file(&staging);
            anyhow::bail!(
                "{} failed the SQLite integrity check: {integrity}",
                source.display()
            );
        }
        // Reject anything that is not actually a tokenix index before it can
        // replace a working one.
        let ok: i64 = staged
            .query_row(
                "SELECT COUNT(*) FROM sqlite_master WHERE type='table'
                   AND name IN ('files','chunks','meta','graph_nodes')",
                [],
                |r| r.get(0),
            )
            .unwrap_or(0);
        if ok < 4 {
            let _ = std::fs::remove_file(&staging);
            anyhow::bail!(
                "{} does not look like a tokenix index (missing core tables)",
                source.display()
            );
        }
        // Bring a snapshot from an older tokenix up to the current schema.
        store::init_schema(&staged, 768)?;
        let (files, chunks) = counts(&staged);
        if chunks == 0 {
            let _ = std::fs::remove_file(&staging);
            anyhow::bail!(
                "{} holds no chunks — refusing to import an empty index",
                source.display()
            );
        }
        (
            files,
            chunks,
            store::meta_value(&staged, "snapshot_version"),
            indexed_at(&staged),
        )
    };

    // Never silently discard a local index that is newer than the snapshot.
    let mut replaced_newer = false;
    if target.exists() {
        let local_at = Connection::open(&target)
            .ok()
            .map(|c| indexed_at(&c))
            .unwrap_or(0.0);
        if local_at > snapshot_at {
            if !force {
                let _ = std::fs::remove_file(&staging);
                anyhow::bail!(
                    "the local index is newer than {} — pass --force to replace it anyway",
                    source.display()
                );
            }
            replaced_newer = true;
        }
    }

    let backup = if target.exists() {
        let backup = target.with_extension("pre-import.bak");
        let _ = std::fs::remove_file(&backup);
        std::fs::rename(&target, &backup).with_context(|| {
            format!(
                "cannot move the existing index aside to {}",
                backup.display()
            )
        })?;
        Some(backup)
    } else {
        None
    };
    // A WAL/SHM pair left from the replaced database would be applied on top of
    // the imported file and corrupt it.
    for suffix in ["-wal", "-shm"] {
        let mut sidecar = target.clone().into_os_string();
        sidecar.push(suffix);
        let _ = std::fs::remove_file(PathBuf::from(sidecar));
    }
    std::fs::rename(&staging, &target).context("cannot move the imported index into place")?;

    Ok(ImportReport {
        source,
        target,
        backup,
        files,
        chunks,
        created_by,
        replaced_newer,
    })
}

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

    fn temp_repo(tag: &str) -> PathBuf {
        let dir = std::env::temp_dir().join(format!(
            "tokenix_snapshot_{tag}_{}_{:?}",
            std::process::id(),
            std::thread::current().id()
        ));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    /// Build a minimal but real index database at `path`.
    fn seed_index(path: &Path, chunk_body: &str, indexed_at_value: f64) {
        let conn = Connection::open(path).unwrap();
        store::init_schema(&conn, 768).unwrap();
        conn.execute(
            "INSERT INTO files(id,path,mtime,content_hash) VALUES(1,'src/a.rs',1.0,'h')",
            [],
        )
        .unwrap();
        conn.execute(
            "INSERT INTO chunks(id,file_id,path,start_line,end_line,symbol,kind,content,token_count)
             VALUES(1,1,'src/a.rs',1,5,'alpha','function',?1,10)",
            [chunk_body],
        )
        .unwrap();
        conn.execute_batch(
            "CREATE TABLE embedding_cache(content_hash TEXT PRIMARY KEY, embedding BLOB, updated_at REAL);
             INSERT INTO embedding_cache(content_hash,embedding,updated_at) VALUES('x',zeroblob(4096),1.0)",
        )
        .unwrap();
        store::set_meta(&conn, "indexed_at", &indexed_at_value.to_string()).unwrap();
    }

    #[test]
    fn export_then_import_round_trips_the_index() {
        let dir = temp_repo("roundtrip");
        let db = dir.join("index.db");
        seed_index(&db, "fn alpha() {}", 100.0);

        // Export straight from the seeded database (not via db_path, which is
        // keyed to the real home directory).
        let out = dir.join("snap.db.gz");
        let staging = dir.join("staging.db");
        {
            let conn = Connection::open(&db).unwrap();
            conn.execute_batch(&format!("VACUUM INTO {}", sql_literal(&staging)))
                .unwrap();
            let staged = Connection::open(&staging).unwrap();
            staged
                .execute_batch(
                    "DROP TABLE IF EXISTS embeddings; DROP TABLE IF EXISTS embedding_cache;",
                )
                .unwrap();
            store::set_meta(&staged, "snapshot_version", "test").unwrap();
        }
        let mut reader = BufReader::new(File::open(&staging).unwrap());
        let mut enc = GzEncoder::new(
            BufWriter::new(File::create(&out).unwrap()),
            Compression::default(),
        );
        std::io::copy(&mut reader, &mut enc).unwrap();
        enc.finish()
            .unwrap()
            .into_inner()
            .unwrap()
            .sync_all()
            .unwrap();

        // Decompress and verify the payload survived intact and cache-free.
        let restored = dir.join("restored.db");
        {
            let mut dec = GzDecoder::new(BufReader::new(File::open(&out).unwrap()));
            let mut w = BufWriter::new(File::create(&restored).unwrap());
            std::io::copy(&mut dec, &mut w).unwrap();
        }
        let conn = Connection::open(&restored).unwrap();
        let (files, chunks) = counts(&conn);
        assert_eq!((files, chunks), (1, 1));
        let legacy_tables: i64 = conn
            .query_row(
                "SELECT COUNT(*) FROM sqlite_master WHERE name IN ('embeddings','embedding_cache')",
                [],
                |r| r.get(0),
            )
            .unwrap();
        assert_eq!(legacy_tables, 0, "legacy embedding tables must not ship");
        let body: String = conn
            .query_row("SELECT content FROM chunks WHERE id=1", [], |r| r.get(0))
            .unwrap();
        assert_eq!(body, "fn alpha() {}");

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn import_rejects_a_file_that_is_not_an_index() {
        let dir = temp_repo("garbage");
        let bogus = dir.join("bogus.db.gz");
        {
            let mut enc = GzEncoder::new(
                BufWriter::new(File::create(&bogus).unwrap()),
                Compression::default(),
            );
            std::io::Write::write_all(&mut enc, b"definitely not sqlite").unwrap();
            enc.finish().unwrap();
        }
        let err = import(&dir, Some(&bogus), false).unwrap_err().to_string();
        assert!(
            err.contains("does not look like a tokenix index")
                || err.contains("not a readable index")
                || err.contains("integrity check"),
            "unexpected error: {err}"
        );
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn gunzip_stops_at_the_cap() {
        let mut gz = Vec::new();
        {
            let mut enc = GzEncoder::new(&mut gz, Compression::default());
            std::io::Write::write_all(&mut enc, &vec![0u8; 10_000]).unwrap();
            enc.finish().unwrap();
        }
        let mut sink = Vec::new();
        assert!(gunzip_capped(&gz[..], &mut sink, 9_999).is_err());
        let mut sink = Vec::new();
        assert_eq!(gunzip_capped(&gz[..], &mut sink, 10_000).unwrap(), 10_000);
    }

    #[test]
    fn import_rejects_a_corrupted_index_and_keeps_the_local_one() {
        let dir = temp_repo("corrupt");
        let db = dir.join("seed.db");
        seed_index(&db, &"fn alpha() {}\n".repeat(400), 100.0);
        let mut bytes = std::fs::read(&db).unwrap();
        // Scribble over the interior pages, past the header and schema page.
        let end = bytes.len();
        for b in &mut bytes[end / 2..end] {
            *b = 0xA5;
        }
        let bad = dir.join("bad.db.gz");
        {
            let mut enc = GzEncoder::new(File::create(&bad).unwrap(), Compression::default());
            std::io::Write::write_all(&mut enc, &bytes).unwrap();
            enc.finish().unwrap();
        }
        assert!(import(&dir, Some(&bad), false).is_err());
        assert!(!store::db_path(&dir).exists());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn import_rejects_non_gzip_input() {
        let dir = temp_repo("plain");
        let plain = dir.join("plain.db.gz");
        std::fs::write(&plain, b"not gzip at all").unwrap();
        assert!(import(&dir, Some(&plain), false).is_err());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn missing_snapshot_is_a_clear_error() {
        let dir = temp_repo("missing");
        let err = import(&dir, Some(&dir.join("nope.gz")), false)
            .unwrap_err()
            .to_string();
        assert!(err.contains("snapshot not found"), "{err}");
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn sql_literal_escapes_quotes() {
        assert_eq!(
            sql_literal(Path::new("/tmp/it's.db")),
            "'/tmp/it''s.db'".to_string()
        );
    }
}