kiromi-ai-cli 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! `kiromi-ai-memory snapshot {create,list,delete,restore}` (Plan 12 phase J).

use kiromi_ai_memory::{RestoreOpts, SnapshotId, SnapshotOpts, SnapshotRef};

use crate::cli::{
    GlobalArgs, SnapshotCmd, SnapshotCreateArgs, SnapshotDeleteArgs, SnapshotRestoreArgs,
};
use crate::error::CliError;
use crate::runtime::Runtime;

pub(crate) async fn run(cmd: SnapshotCmd, globals: &GlobalArgs) -> Result<(), CliError> {
    let rt = Runtime::open(globals).await?;
    match cmd {
        SnapshotCmd::Create(a) => create(rt, a).await,
        SnapshotCmd::List => list(rt).await,
        SnapshotCmd::Delete(a) => delete(rt, a).await,
        SnapshotCmd::Restore(a) => restore(rt, a).await,
    }
}

async fn create(rt: Runtime, a: SnapshotCreateArgs) -> Result<(), CliError> {
    let mut opts = SnapshotOpts::default();
    if let Some(t) = a.tag {
        opts = opts.with_tag(t);
    }
    if let Some(r) = a.reason {
        opts = opts.with_reason(r);
    }
    let s = rt.mem.snapshot(opts).await?;
    if rt.json {
        println!(
            "{}",
            serde_json::to_string_pretty(&snapshot_json(&s)).unwrap_or_default()
        );
    } else {
        println!("{}\t{}\t{}", s.id, s.seq, s.tag.as_deref().unwrap_or(""));
    }
    rt.mem.close().await?;
    Ok(())
}

async fn list(rt: Runtime) -> Result<(), CliError> {
    let snaps = rt.mem.list_snapshots().await?;
    if rt.json {
        let arr: Vec<_> = snaps.iter().map(snapshot_json).collect();
        println!("{}", serde_json::to_string_pretty(&arr).unwrap_or_default());
    } else {
        for s in &snaps {
            println!(
                "{}\t{}\t{}\t{}",
                s.id,
                s.seq,
                s.created_at_ms,
                s.tag.as_deref().unwrap_or("")
            );
        }
    }
    rt.mem.close().await?;
    Ok(())
}

async fn delete(rt: Runtime, a: SnapshotDeleteArgs) -> Result<(), CliError> {
    let id: SnapshotId = a.id.parse::<SnapshotId>().map_err(|e| CliError {
        kind: crate::error::ExitCode::Config,
        source: anyhow::anyhow!("snapshot id: {e}"),
    })?;
    // Locate the row to recover seq/created_at.
    let snaps = rt.mem.list_snapshots().await?;
    let Some(sref) = snaps.into_iter().find(|s| s.id == id) else {
        return Err(CliError {
            kind: crate::error::ExitCode::NotFound,
            source: anyhow::anyhow!("snapshot not found: {}", a.id),
        });
    };
    rt.mem.delete_snapshot(&sref).await?;
    if !rt.json {
        println!("deleted {}", sref.id);
    }
    rt.mem.close().await?;
    Ok(())
}

async fn restore(rt: Runtime, a: SnapshotRestoreArgs) -> Result<(), CliError> {
    let id: SnapshotId = a.id.parse::<SnapshotId>().map_err(|e| CliError {
        kind: crate::error::ExitCode::Config,
        source: anyhow::anyhow!("snapshot id: {e}"),
    })?;
    let snaps = rt.mem.list_snapshots().await?;
    let Some(sref) = snaps.into_iter().find(|s| s.id == id) else {
        return Err(CliError {
            kind: crate::error::ExitCode::NotFound,
            source: anyhow::anyhow!("snapshot not found: {}", a.id),
        });
    };
    let opts = RestoreOpts::default().with_also_restore_attributes(!a.no_attributes);
    let report = rt.mem.restore(&sref, opts).await?;
    if rt.json {
        let body = serde_json::json!({
            "snapshot_id": sref.id.to_string(),
            "memories_re_tombstoned": report.memories_re_tombstoned,
            "memories_un_tombstoned": report.memories_un_tombstoned,
            "summaries_re_tombstoned": report.summaries_re_tombstoned,
            "summaries_un_tombstoned": report.summaries_un_tombstoned,
            "links_added": report.links_added,
            "links_removed": report.links_removed,
            "attributes_set": report.attributes_set,
            "attributes_cleared": report.attributes_cleared,
        });
        println!(
            "{}",
            serde_json::to_string_pretty(&body).unwrap_or_default()
        );
    } else {
        println!(
            "restored to {}: {} mems re-tombstoned, {} un-tombstoned",
            sref.id, report.memories_re_tombstoned, report.memories_un_tombstoned
        );
    }
    rt.mem.close().await?;
    Ok(())
}

fn snapshot_json(s: &SnapshotRef) -> serde_json::Value {
    serde_json::json!({
        "id": s.id.to_string(),
        "seq": s.seq,
        "created_at_ms": s.created_at_ms,
        "tag": s.tag,
        "reason": s.reason,
    })
}