kiromi-ai-cli 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! `kiromi-ai-memory link {add|remove|list}`.

use std::str::FromStr;

use kiromi_ai_memory::{MemoryId, MemoryRef};

use crate::cli::{GlobalArgs, LinkCmd};
use crate::cmd::get::probe_partition;
use crate::error::{CliError, ExitCode};
use crate::output;
use crate::runtime::Runtime;

pub(crate) async fn run(cmd: LinkCmd, globals: &GlobalArgs) -> Result<(), CliError> {
    let rt = Runtime::open(globals).await?;
    let result = match cmd {
        LinkCmd::Add { src, dst } => {
            let s = make_ref(&src, &rt).await?;
            let d = make_ref(&dst, &rt).await?;
            rt.mem.add_link(&s, &d).await?;
            if globals.json {
                println!("{}", output::to_json(&serde_json::json!({"added": true})));
            } else {
                println!("ok");
            }
            Ok(())
        }
        LinkCmd::Remove { src, dst } => {
            let s = make_ref(&src, &rt).await?;
            let d = make_ref(&dst, &rt).await?;
            rt.mem.remove_link(&s, &d).await?;
            if globals.json {
                println!("{}", output::to_json(&serde_json::json!({"removed": true})));
            } else {
                println!("ok");
            }
            Ok(())
        }
        LinkCmd::List { id } => {
            let r = make_ref(&id, &rt).await?;
            let links = rt.mem.links_of(&r).await?;
            // `links_of` returns BOTH directions for a memory; restrict the
            // CLI surface to outbound links (src == requested id) so the table
            // matches the user's mental model of "links from this memory".
            let outbound: Vec<_> = links.iter().filter(|l| l.src == r.id).collect();
            if globals.json {
                let wire: Vec<_> = outbound
                    .iter()
                    .map(|l| {
                        serde_json::json!({
                            "src": l.src.to_string(),
                            "dst": l.dst.to_string(),
                            "kind": format!("{:?}", l.kind),
                            "created_at_ms": l.created_at_ms,
                        })
                    })
                    .collect();
                println!("{}", output::to_json(&wire));
            } else {
                let mut t = comfy_table::Table::new();
                t.set_content_arrangement(comfy_table::ContentArrangement::Dynamic);
                t.set_header(vec!["src", "dst", "kind"]);
                for l in &outbound {
                    t.add_row(vec![
                        l.src.to_string(),
                        l.dst.to_string(),
                        format!("{:?}", l.kind),
                    ]);
                }
                println!("{t}");
            }
            Ok(())
        }
    };
    rt.mem.close().await?;
    result
}

/// Resolve `<ulid>` → `MemoryRef` by fetching the row to learn its partition.
async fn make_ref(s: &str, rt: &Runtime) -> Result<MemoryRef, CliError> {
    let id = MemoryId::from_str(s).map_err(|e| CliError {
        kind: ExitCode::Config,
        source: anyhow::anyhow!("bad id {s:?}: {e}"),
    })?;
    let probe = MemoryRef {
        id,
        partition: probe_partition()?,
    };
    let rec = rt.mem.get(&probe).await?;
    Ok(rec.r#ref)
}