clawhdf5/crates/clawhdf5-cli/src/main.rs

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

use clap::{Parser, Subcommand};
use clawhdf5_agent::signing::{self, SigningKey, VerifyingKey};
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};

/// ClawhDF5 — HDF5-backed cognitive memory for AI agents
#[derive(Parser)]
#[command(name = "clawhdf5", version, about)]
struct Cli {
    /// Path to the .h5 memory file (not needed for `keygen`)
    #[arg(short, long, env = "CLAWHDF5_PATH")]
    path: Option<PathBuf>,

    /// File holding an Ed25519 signing key (64 hex characters, from
    /// `keygen`). Every checkpoint this command makes is then signed; a
    /// signed store refuses to checkpoint without it.
    #[arg(long, env = "CLAWHDF5_SIGNING_KEY", global = true)]
    signing_key: Option<PathBuf>,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Create a new memory file
    Create {
        /// Agent identifier
        #[arg(long, default_value = "default")]
        agent_id: String,
        /// Embedding dimension
        #[arg(long, default_value_t = 384)]
        dim: usize,
        /// Enable write-ahead log
        #[arg(long)]
        wal: bool,
        /// Hold the vector index's copy of the embeddings as f32 instead of
        /// the default int8 (which uses a quarter of the memory and is faster
        /// at equal recall)
        #[arg(long)]
        f32_index: bool,
        /// Accepted for compatibility; int8 is now the default
        #[arg(long, hide = true, conflicts_with = "f32_index")]
        quantized_index: bool,
        /// Store embeddings as full-precision f32 instead of the default
        /// half precision (float16: half the bytes, about three significant
        /// digits, values within ±65504)
        #[arg(long)]
        f32: bool,
        /// Accepted for compatibility; float16 is now the default
        #[arg(long, hide = true, conflicts_with = "f32")]
        float16: bool,
    },
    /// Save a memory entry (reads JSON from stdin or --json)
    Save {
        /// JSON: {"chunk":"...","embedding":[...],"source_channel":"...","timestamp":0.0,"session_id":"...","tags":""}
        #[arg(long)]
        json: Option<String>,
    },
    /// Search memory by embedding vector
    Search {
        /// Query embedding as JSON array of f32
        #[arg(long)]
        embedding: String,
        /// Optional text query for hybrid BM25+vector search
        #[arg(long, default_value = "")]
        query: String,
        /// Number of results
        #[arg(short = 'k', long, default_value_t = 5)]
        top_k: usize,
        /// Vector similarity weight (0.0-1.0)
        #[arg(long, default_value_t = 0.7)]
        vector_weight: f32,
        /// BM25 keyword weight (0.0-1.0)
        #[arg(long, default_value_t = 0.3)]
        keyword_weight: f32,
    },
    /// Get a specific memory chunk by index
    Recall {
        /// Chunk index
        index: usize,
    },
    /// Show memory stats (count, active, config)
    Stats,
    /// Flush WAL to main HDF5 file
    FlushWal,
    /// Generate AGENTS.md from memory contents
    AgentsMd {
        /// Write to file instead of stdout
        #[arg(long)]
        output: Option<PathBuf>,
    },
    /// Export all entries as JSON lines to stdout
    Export,
    /// Snapshot (copy) memory file to destination
    Snapshot {
        /// Destination path
        dest: PathBuf,
    },
    /// Generate an Ed25519 signing key for signed checkpoints
    Keygen {
        /// Where to write the secret key (created new, owner-only on Unix)
        #[arg(long)]
        out: PathBuf,
    },
    /// Verify a signed store against a public key; exit status 2 if not valid
    Verify {
        /// The trusted public key: 64 hex characters, or a file holding them
        #[arg(long)]
        public_key: String,
    },
}

fn read_signing_key(path: &Path) -> Result<SigningKey, Box<dyn std::error::Error>> {
    let text = std::fs::read_to_string(path)
        .map_err(|e| format!("cannot read signing key {}: {e}", path.display()))?;
    let bytes = signing::from_hex::<32>(&text)
        .ok_or_else(|| format!("{} is not a 64-hex-character key", path.display()))?;
    Ok(SigningKey::from_bytes(&bytes))
}

/// Open for writing, with the signing key applied if one was given.
fn open_writable(
    path: &Path,
    key: &Option<SigningKey>,
) -> Result<HDF5Memory, Box<dyn std::error::Error>> {
    let mut mem = HDF5Memory::open(path)?;
    if let Some(k) = key {
        mem.set_signing_key(k.clone());
    }
    Ok(mem)
}

fn main() {
    let cli = Cli::parse();

    if let Err(e) = run(cli) {
        eprintln!("error: {e}");
        std::process::exit(1);
    }
}

fn run(cli: Cli) -> Result<(), Box<dyn std::error::Error>> {
    if let Commands::Keygen { out } = &cli.command {
        let key = signing::generate_key();
        let mut opts = std::fs::OpenOptions::new();
        opts.write(true).create_new(true);
        #[cfg(unix)]
        {
            use std::os::unix::fs::OpenOptionsExt;
            opts.mode(0o600);
        }
        use std::io::Write;
        let mut f = opts
            .open(out)
            .map_err(|e| format!("cannot create {}: {e}", out.display()))?;
        writeln!(f, "{}", signing::to_hex(&key.to_bytes()))?;
        let j = serde_json::json!({
            "status": "generated",
            "secret_key_file": out.display().to_string(),
            "public_key": signing::to_hex(&key.verifying_key().to_bytes()),
        });
        println!("{}", serde_json::to_string_pretty(&j)?);
        return Ok(());
    }
    let path = cli
        .path
        .clone()
        .ok_or("--path (or CLAWHDF5_PATH) is required")?;
    let key = cli
        .signing_key
        .as_deref()
        .map(read_signing_key)
        .transpose()?;
    match cli.command {
        Commands::Create {
            agent_id,
            dim,
            wal,
            f32_index,
            quantized_index: _,
            f32,
            float16: _,
        } => {
            let mut config = MemoryConfig::new(path.clone(), &agent_id, dim);
            config.wal_enabled = wal;
            // As with --f32-index: only ever switch the library default off.
            if f32 {
                config.float16 = false;
            }
            let config_float16 = config.float16;
            // Only ever switch *off* the library default: assigning the flag
            // outright would force every CLI-created store back to f32 unless
            // the caller knew to ask for int8.
            if f32_index {
                config.quantized_index = false;
            }
            let config_quantized = config.quantized_index;
            let mut mem = HDF5Memory::create(config)?;
            // Sign straight away, so the store is never on disk unsigned.
            if let Some(k) = &key {
                mem.set_signing_key(k.clone());
                mem.flush_wal()?;
            }
            let j = serde_json::json!({
                "status": "created",
                "path": path.display().to_string(),
                "agent_id": agent_id,
                "embedding_dim": dim,
                "wal_enabled": wal,
                "quantized_index": config_quantized,
                "float16": config_float16,
                "signed": mem.is_signed(),
                "count": mem.count(),
            });
            println!("{}", serde_json::to_string_pretty(&j)?);
        }

        Commands::Save { json } => {
            let input = match json {
                Some(s) => s,
                None => {
                    use std::io::Read;
                    let mut buf = String::new();
                    std::io::stdin().read_to_string(&mut buf)?;
                    buf
                }
            };
            let entry: MemoryEntry = serde_json::from_str(&input)?;
            let mut mem = open_writable(&path, &key)?;
            let idx = mem.save(entry)?;
            let j = serde_json::json!({ "status": "saved", "index": idx, "count": mem.count() });
            println!("{}", serde_json::to_string(&j)?);
        }

        Commands::Search {
            embedding,
            query,
            top_k,
            vector_weight,
            keyword_weight,
        } => {
            let emb: Vec<f32> = serde_json::from_str(&embedding)?;
            let mut mem = open_writable(&path, &key)?;
            let results = mem.hybrid_search(&emb, &query, vector_weight, keyword_weight, top_k);
            let j: Vec<serde_json::Value> = results
                .iter()
                .map(|r| {
                    serde_json::json!({
                        "index": r.index,
                        "score": r.score,
                        "chunk": &r.chunk,
                        "timestamp": r.timestamp,
                        "source_channel": &r.source_channel,
                    })
                })
                .collect();
            println!("{}", serde_json::to_string_pretty(&j)?);
        }

        Commands::Recall { index } => {
            let mem = HDF5Memory::open_read_only(&path)?;
            match mem.get_chunk(index) {
                Some(content) => {
                    let j = serde_json::json!({ "index": index, "chunk": content });
                    println!("{}", serde_json::to_string_pretty(&j)?);
                }
                None => {
                    eprintln!("no entry at index {index}");
                    std::process::exit(1);
                }
            }
        }

        Commands::Stats => {
            let mem = HDF5Memory::open_read_only(&path)?;
            let cfg = mem.config();
            let j = serde_json::json!({
                "path": path.display().to_string(),
                "agent_id": cfg.agent_id,
                "embedding_dim": cfg.embedding_dim,
                "count": mem.count(),
                "active": mem.count_active(),
                "wal_enabled": cfg.wal_enabled,
                "wal_pending": mem.wal_pending_count(),
                "signed": mem.is_signed(),
            });
            println!("{}", serde_json::to_string_pretty(&j)?);
        }

        Commands::FlushWal => {
            let mut mem = open_writable(&path, &key)?;
            let before = mem.wal_pending_count();
            mem.flush_wal()?;
            let j = serde_json::json!({
                "status": "flushed",
                "entries_flushed": before,
                "wal_pending": mem.wal_pending_count(),
            });
            println!("{}", serde_json::to_string(&j)?);
        }

        Commands::AgentsMd { output } => {
            let mem = HDF5Memory::open_read_only(&path)?;
            let md = mem.generate_agents_md();
            match output {
                Some(p) => {
                    std::fs::write(&p, &md)?;
                    eprintln!("wrote {}", p.display());
                }
                None => print!("{md}"),
            }
        }

        Commands::Export => {
            let mem = HDF5Memory::open_read_only(&path)?;
            for i in 0..mem.count() {
                if let Some(chunk) = mem.get_chunk(i) {
                    let j = serde_json::json!({ "index": i, "chunk": chunk });
                    println!("{}", serde_json::to_string(&j)?);
                }
            }
        }

        Commands::Keygen { .. } => unreachable!("handled before opening a store"),

        Commands::Verify { public_key } => {
            let text = if Path::new(&public_key).is_file() {
                std::fs::read_to_string(&public_key)?
            } else {
                public_key
            };
            let bytes = signing::from_hex::<32>(&text)
                .ok_or("--public-key must be 64 hex characters or a file holding them")?;
            let trusted = VerifyingKey::from_bytes(&bytes)?;
            let r = HDF5Memory::verify(&path, &trusted)?;
            let j = serde_json::json!({
                "valid": r.is_valid(),
                "signed": r.signed,
                "key_matches": r.key_matches,
                "signature_valid": r.signature_valid,
                "records_match": r.records_match,
                "settings_match": r.settings_match,
                "sessions_match": r.sessions_match,
                "graph_match": r.graph_match,
                "changed_records": r.changed_records,
                "record_count": r.record_count,
                "signed_record_count": r.signed_record_count,
                "signed_by": r.public_key.map(|k| signing::to_hex(&k)),
                "wal_entries_unsigned": r.wal_entries_unsigned,
            });
            println!("{}", serde_json::to_string_pretty(&j)?);
            if !r.is_valid() {
                std::process::exit(2);
            }
        }

        Commands::Snapshot { dest } => {
            let _result = clawhdf5_agent::storage::snapshot_file(&path, &dest)?;
            let j = serde_json::json!({
                "status": "snapshot_created",
                "source": path.display().to_string(),
                "dest": dest.display().to_string(),
            });
            println!("{}", serde_json::to_string(&j)?);
        }
    }

    Ok(())
}