use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::command_message::CommandMessage;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ArgType {
String,
Integer,
Float,
Boolean,
Json,
}
impl std::fmt::Display for ArgType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ArgType::String => write!(f, "string"),
ArgType::Integer => write!(f, "integer"),
ArgType::Float => write!(f, "float"),
ArgType::Boolean => write!(f, "boolean"),
ArgType::Json => write!(f, "json"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArgDef {
pub name: std::string::String,
pub arg_type: ArgType,
pub required: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub default: Option<Value>,
pub description: std::string::String,
}
#[derive(Debug, Clone)]
pub struct CommandDef {
pub name: std::string::String,
pub description: std::string::String,
pub args: Vec<ArgDef>,
pub allow_extra_args: bool,
pub listed: bool,
}
impl CommandDef {
pub fn new(name: &str, description: &str) -> Self {
CommandDef {
name: name.to_string(),
description: description.to_string(),
args: Vec::new(),
allow_extra_args: false,
listed: true,
}
}
pub fn arg(mut self, name: &str, arg_type: ArgType, required: bool, description: &str) -> Self {
self.args.push(ArgDef {
name: name.to_string(),
arg_type,
required,
default: None,
description: description.to_string(),
});
self
}
pub fn arg_with_default(mut self, name: &str, arg_type: ArgType, default: Value, description: &str) -> Self {
self.args.push(ArgDef {
name: name.to_string(),
arg_type,
required: false,
default: Some(default),
description: description.to_string(),
});
self
}
pub fn extra_args(mut self) -> Self {
self.allow_extra_args = true;
self
}
pub fn unlisted(mut self) -> Self {
self.listed = false;
self
}
pub fn parse(&self, data: &Value) -> Result<ParsedArgs, std::string::String> {
let obj = match data.as_object() {
Some(o) => o,
None if data.is_null() && self.args.iter().all(|a| !a.required) => {
let mut result = serde_json::Map::new();
for arg in &self.args {
if let Some(ref def) = arg.default {
result.insert(arg.name.clone(), def.clone());
}
}
return Ok(ParsedArgs { inner: result });
}
None => return Err("Expected a JSON object for command arguments".to_string()),
};
let mut result = serde_json::Map::new();
let known_names: Vec<&str> = self.args.iter().map(|a| a.name.as_str()).collect();
for arg in &self.args {
match obj.get(&arg.name) {
Some(val) => {
validate_type(&arg.name, val, &arg.arg_type)?;
result.insert(arg.name.clone(), val.clone());
}
None if arg.required => {
return Err(format!("Missing required argument: --{}", arg.name));
}
None => {
if let Some(ref def) = arg.default {
result.insert(arg.name.clone(), def.clone());
}
}
}
}
const SERVER_INJECTED: &[&str] = &["instance_name", "project_file"];
for (key, val) in obj {
if known_names.contains(&key.as_str()) || SERVER_INJECTED.contains(&key.as_str()) {
continue;
}
if self.allow_extra_args {
result.insert(key.clone(), val.clone());
} else {
return Err(format!("Unknown argument: --{}", key));
}
}
Ok(ParsedArgs { inner: result })
}
pub fn help_text(&self, domain: &str) -> std::string::String {
let mut out = format!("{}.{} — {}\n", domain, self.name, self.description);
if self.args.is_empty() {
out.push_str("\n No arguments.\n");
return out;
}
out.push_str("\nArguments:\n");
for arg in &self.args {
let req = if arg.required { "required" } else { "optional" };
let default_str = match &arg.default {
Some(v) => format!(" [default: {}]", v),
None => std::string::String::new(),
};
out.push_str(&format!(
" --{:<24} ({}, {}){} {}\n",
arg.name, arg.arg_type, req, default_str, arg.description
));
}
if self.allow_extra_args {
out.push_str("\n Additional --key value flags are passed through as extra parameters.\n");
}
out
}
}
fn validate_type(name: &str, val: &Value, expected: &ArgType) -> Result<(), std::string::String> {
let ok = match expected {
ArgType::String => val.is_string(),
ArgType::Integer => val.is_i64() || val.is_u64() || val.is_f64(),
ArgType::Float => val.is_number(),
ArgType::Boolean => val.is_boolean(),
ArgType::Json => true, };
if ok {
Ok(())
} else {
Err(format!(
"Argument --{} expected type {}, got {}",
name,
expected,
value_type_name(val),
))
}
}
fn value_type_name(v: &Value) -> &'static str {
match v {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
#[derive(Debug, Clone)]
pub struct ParsedArgs {
inner: serde_json::Map<std::string::String, Value>,
}
impl ParsedArgs {
pub fn get_str(&self, name: &str) -> Option<&str> {
self.inner.get(name).and_then(|v| v.as_str())
}
pub fn get_f64(&self, name: &str) -> Option<f64> {
self.inner.get(name).and_then(|v| v.as_f64())
}
pub fn get_i64(&self, name: &str) -> Option<i64> {
self.inner.get(name).and_then(|v| v.as_i64())
}
pub fn get_bool(&self, name: &str) -> Option<bool> {
self.inner.get(name).and_then(|v| v.as_bool())
}
pub fn get_json(&self, name: &str) -> Option<&Value> {
self.inner.get(name)
}
pub fn require_str(&self, name: &str) -> Result<&str, std::string::String> {
self.get_str(name)
.ok_or_else(|| format!("Missing required argument: --{}", name))
}
pub fn require_f64(&self, name: &str) -> Result<f64, std::string::String> {
self.get_f64(name)
.ok_or_else(|| format!("Missing required argument: --{}", name))
}
pub fn require_i64(&self, name: &str) -> Result<i64, std::string::String> {
self.get_i64(name)
.ok_or_else(|| format!("Missing required argument: --{}", name))
}
pub fn require_bool(&self, name: &str) -> Result<bool, std::string::String> {
self.get_bool(name)
.ok_or_else(|| format!("Missing required argument: --{}", name))
}
pub fn iter(&self) -> serde_json::map::Iter<'_> {
self.inner.iter()
}
pub fn into_map(self) -> serde_json::Map<std::string::String, Value> {
self.inner
}
pub fn as_map(&self) -> &serde_json::Map<std::string::String, Value> {
&self.inner
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CatalogEntry {
pub fqdn: std::string::String,
pub description: std::string::String,
pub args: Vec<ArgDef>,
}
pub struct CommandRegistry {
domain: std::string::String,
commands: IndexMap<std::string::String, CommandDef>,
}
impl CommandRegistry {
pub fn new(domain: &str) -> Self {
CommandRegistry {
domain: domain.to_string(),
commands: IndexMap::new(),
}
}
pub fn register(&mut self, cmd: CommandDef) {
self.commands.insert(cmd.name.clone(), cmd);
}
pub fn catalog_fqdns(&self) -> Vec<std::string::String> {
self.commands
.values()
.filter(|c| c.listed)
.map(|c| format!("{}.{}", self.domain, c.name))
.collect()
}
pub fn catalog_entries(&self) -> Vec<CatalogEntry> {
self.commands
.values()
.filter(|c| c.listed)
.map(|c| CatalogEntry {
fqdn: format!("{}.{}", self.domain, c.name),
description: c.description.clone(),
args: c.args.clone(),
})
.collect()
}
pub fn get(&self, name: &str) -> Option<&CommandDef> {
self.commands.get(name)
}
pub fn parse_args(&self, subtopic: &str, data: &Value) -> Result<ParsedArgs, std::string::String> {
match self.commands.get(subtopic) {
Some(cmd) => cmd.parse(data),
None => {
let inner = data.as_object().cloned().unwrap_or_default();
Ok(ParsedArgs { inner })
}
}
}
pub fn handle_meta(&self, msg: &CommandMessage) -> Option<CommandMessage> {
let subtopic = msg.subtopic();
match subtopic.as_str() {
"help" => {
let command_name = msg.data.get("command").and_then(|v| v.as_str());
Some(self.handle_help(msg, command_name))
}
"get_catalog" => {
let detailed = msg.data.get("detailed")
.and_then(|v| v.as_bool())
.unwrap_or(false);
if detailed {
let entries = self.catalog_entries();
Some(msg.clone().into_response(
serde_json::to_value(&entries).unwrap_or(Value::Null),
))
} else {
let fqdns = self.catalog_fqdns();
Some(msg.clone().into_response(
serde_json::to_value(&fqdns).unwrap_or(Value::Null),
))
}
}
_ => None,
}
}
fn handle_help(&self, msg: &CommandMessage, command_name: Option<&str>) -> CommandMessage {
match command_name {
Some(name) => {
match self.commands.get(name) {
Some(cmd) => {
let text = cmd.help_text(&self.domain);
let entry = CatalogEntry {
fqdn: format!("{}.{}", self.domain, cmd.name),
description: cmd.description.clone(),
args: cmd.args.clone(),
};
msg.clone().into_response(serde_json::json!({
"text": text,
"command": entry,
}))
}
None => {
msg.clone().into_error_response(&format!(
"Unknown command '{}'. Use {}.help for a list of commands.",
name, self.domain,
))
}
}
}
None => {
let mut text = format!("{} commands:\n\n", self.domain);
for cmd in self.commands.values() {
if cmd.listed {
text.push_str(&format!(" {:<28} {}\n", cmd.name, cmd.description));
}
}
text.push_str(&format!(
"\nUse {}.help --command <name> for detailed argument info.\n",
self.domain
));
msg.clone().into_response(serde_json::json!({
"text": text,
"commands": self.catalog_entries(),
}))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn sample_command() -> CommandDef {
CommandDef::new("add_task", "Add a new task")
.arg("name", ArgType::String, true, "Task name")
.arg("sample_rate", ArgType::Float, true, "Sample rate in Hz")
.arg_with_default("timeout_ms", ArgType::Integer, json!(2500), "Read timeout")
.arg_with_default("clock_type", ArgType::String, json!("internal"), "Clock type")
}
#[test]
fn test_parse_required_args_present() {
let cmd = sample_command();
let data = json!({"name": "Encoder1", "sample_rate": 100000.0});
let args = cmd.parse(&data).unwrap();
assert_eq!(args.get_str("name"), Some("Encoder1"));
assert_eq!(args.get_f64("sample_rate"), Some(100000.0));
assert_eq!(args.get_i64("timeout_ms"), Some(2500));
assert_eq!(args.get_str("clock_type"), Some("internal"));
}
#[test]
fn test_parse_required_arg_missing() {
let cmd = sample_command();
let data = json!({"sample_rate": 1000.0});
let err = cmd.parse(&data).unwrap_err();
assert!(err.contains("--name"), "Error should mention missing arg: {err}");
}
#[test]
fn test_parse_type_mismatch() {
let cmd = sample_command();
let data = json!({"name": 123, "sample_rate": 1000.0});
let err = cmd.parse(&data).unwrap_err();
assert!(err.contains("--name"), "Error should mention the arg: {err}");
assert!(err.contains("string"), "Error should mention expected type: {err}");
}
#[test]
fn test_parse_extra_args_rejected() {
let cmd = sample_command();
let data = json!({"name": "T", "sample_rate": 1000.0, "bogus": 42});
let err = cmd.parse(&data).unwrap_err();
assert!(err.contains("--bogus"));
}
#[test]
fn test_parse_extra_args_allowed() {
let cmd = sample_command().extra_args();
let data = json!({"name": "T", "sample_rate": 1000.0, "dist_per_pulse": 0.001});
let args = cmd.parse(&data).unwrap();
assert_eq!(args.get_f64("dist_per_pulse"), Some(0.001));
}
#[test]
fn test_parse_null_data_no_required() {
let cmd = CommandDef::new("status", "Show status");
let args = cmd.parse(&Value::Null).unwrap();
assert!(args.as_map().is_empty());
}
#[test]
fn test_parse_null_data_with_defaults() {
let cmd = CommandDef::new("status", "Show status")
.arg_with_default("verbose", ArgType::Boolean, json!(false), "Verbose output");
let args = cmd.parse(&Value::Null).unwrap();
assert_eq!(args.get_bool("verbose"), Some(false));
}
#[test]
fn test_parse_null_data_with_required_fails() {
let cmd = CommandDef::new("start", "Start")
.arg("name", ArgType::String, true, "Name");
let err = cmd.parse(&Value::Null).unwrap_err();
assert!(err.contains("JSON object"));
}
#[test]
fn test_help_text() {
let cmd = sample_command();
let text = cmd.help_text("ni");
assert!(text.contains("ni.add_task"));
assert!(text.contains("Add a new task"));
assert!(text.contains("--name"));
assert!(text.contains("required"));
assert!(text.contains("--timeout_ms"));
assert!(text.contains("2500"));
}
#[test]
fn test_help_text_extra_args() {
let cmd = sample_command().extra_args();
let text = cmd.help_text("ni");
assert!(text.contains("passed through"));
}
#[test]
fn test_registry_catalog_fqdns() {
let mut reg = CommandRegistry::new("ni");
reg.register(CommandDef::new("start", "Start acquisition"));
reg.register(CommandDef::new("stop", "Stop acquisition"));
reg.register(CommandDef::new("internal", "Hidden").unlisted());
let fqdns = reg.catalog_fqdns();
assert_eq!(fqdns, vec!["ni.start", "ni.stop"]);
}
#[test]
fn test_registry_catalog_entries() {
let mut reg = CommandRegistry::new("ni");
reg.register(
CommandDef::new("add_task", "Add a task")
.arg("name", ArgType::String, true, "Task name"),
);
let entries = reg.catalog_entries();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].fqdn, "ni.add_task");
assert_eq!(entries[0].args.len(), 1);
}
#[test]
fn test_registry_parse_args_known_command() {
let mut reg = CommandRegistry::new("ni");
reg.register(sample_command());
let data = json!({"name": "T", "sample_rate": 1000.0});
let args = reg.parse_args("add_task", &data).unwrap();
assert_eq!(args.require_str("name").unwrap(), "T");
}
#[test]
fn test_registry_parse_args_unknown_command() {
let reg = CommandRegistry::new("ni");
let data = json!({"foo": "bar"});
let args = reg.parse_args("unknown", &data).unwrap();
assert_eq!(args.get_str("foo"), Some("bar"));
}
#[test]
fn test_handle_meta_help_summary() {
let mut reg = CommandRegistry::new("ni");
reg.register(CommandDef::new("start", "Start acquisition"));
reg.register(CommandDef::new("stop", "Stop acquisition"));
let msg = CommandMessage::request("ni.help", json!({}));
let resp = reg.handle_meta(&msg).unwrap();
assert!(resp.success);
let text = resp.data["text"].as_str().unwrap();
assert!(text.contains("start"));
assert!(text.contains("stop"));
}
#[test]
fn test_handle_meta_help_specific() {
let mut reg = CommandRegistry::new("ni");
reg.register(sample_command());
let msg = CommandMessage::request("ni.help", json!({"command": "add_task"}));
let resp = reg.handle_meta(&msg).unwrap();
assert!(resp.success);
let text = resp.data["text"].as_str().unwrap();
assert!(text.contains("--name"));
assert!(text.contains("--sample_rate"));
}
#[test]
fn test_handle_meta_help_unknown_command() {
let reg = CommandRegistry::new("ni");
let msg = CommandMessage::request("ni.help", json!({"command": "nonexistent"}));
let resp = reg.handle_meta(&msg).unwrap();
assert!(!resp.success);
}
#[test]
fn test_handle_meta_get_catalog() {
let mut reg = CommandRegistry::new("ni");
reg.register(CommandDef::new("start", "Start"));
reg.register(CommandDef::new("stop", "Stop"));
let msg = CommandMessage::request("ni.get_catalog", json!({}));
let resp = reg.handle_meta(&msg).unwrap();
assert!(resp.success);
let fqdns: Vec<String> = serde_json::from_value(resp.data).unwrap();
assert_eq!(fqdns, vec!["ni.start", "ni.stop"]);
}
#[test]
fn test_handle_meta_get_catalog_detailed() {
let mut reg = CommandRegistry::new("ni");
reg.register(
CommandDef::new("start", "Start acquisition")
.arg("mode", ArgType::String, false, "Start mode"),
);
let msg = CommandMessage::request("ni.get_catalog", json!({"detailed": true}));
let resp = reg.handle_meta(&msg).unwrap();
assert!(resp.success);
let entries: Vec<CatalogEntry> = serde_json::from_value(resp.data).unwrap();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].fqdn, "ni.start");
assert_eq!(entries[0].args.len(), 1);
}
#[test]
fn test_handle_meta_non_meta_returns_none() {
let reg = CommandRegistry::new("ni");
let msg = CommandMessage::request("ni.start", json!({}));
assert!(reg.handle_meta(&msg).is_none());
}
#[test]
fn test_parsed_args_require_missing() {
let args = ParsedArgs { inner: serde_json::Map::new() };
assert!(args.require_str("name").is_err());
assert!(args.require_f64("rate").is_err());
assert!(args.require_i64("count").is_err());
assert!(args.require_bool("flag").is_err());
}
#[test]
fn test_parsed_args_iter() {
let cmd = CommandDef::new("test", "Test").extra_args();
let data = json!({"a": 1, "b": "two"});
let args = cmd.parse(&data).unwrap();
let keys: Vec<&String> = args.iter().map(|(k, _)| k).collect();
assert!(keys.contains(&&"a".to_string()));
assert!(keys.contains(&&"b".to_string()));
}
#[test]
fn test_integer_allows_float_input() {
let cmd = CommandDef::new("test", "Test")
.arg("count", ArgType::Integer, true, "Count");
let data = json!({"count": 100000.0});
let args = cmd.parse(&data).unwrap();
assert!(args.get_f64("count").is_some());
}
#[test]
fn test_into_map() {
let cmd = CommandDef::new("test", "Test").extra_args();
let data = json!({"a": 1, "b": "two"});
let args = cmd.parse(&data).unwrap();
let map = args.into_map();
assert_eq!(map.len(), 2);
}
}