1use std::{
2 collections::{BTreeMap, BTreeSet},
3 sync::{OnceLock, RwLock},
4};
5
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
12pub struct FieldInfo {
13 pub name: String,
15 #[serde(rename = "type")]
17 pub field_type: String,
18 pub optional: bool,
20}
21
22impl FieldInfo {
23 #[must_use]
25 pub fn new(name: impl Into<String>, field_type: impl Into<String>) -> Self {
26 Self {
27 name: name.into(),
28 field_type: field_type.into(),
29 optional: false,
30 }
31 }
32
33 #[must_use]
35 pub fn optional(mut self) -> Self {
36 self.optional = true;
37 self
38 }
39}
40
41#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
43pub struct SchemaInfo {
44 pub command: String,
46 #[serde(default, skip_serializing_if = "Vec::is_empty")]
48 pub fields: Vec<FieldInfo>,
49 #[serde(default, skip_serializing_if = "Option::is_none")]
51 pub schema: Option<Value>,
52}
53
54impl SchemaInfo {
55 #[must_use]
57 pub fn new(command: impl Into<String>) -> Self {
58 Self {
59 command: command.into(),
60 fields: Vec::new(),
61 schema: None,
62 }
63 }
64
65 #[must_use]
67 pub fn with_fields(mut self, fields: impl Into<Vec<FieldInfo>>) -> Self {
68 self.fields = fields.into();
69 self
70 }
71
72 #[must_use]
74 pub fn with_schema(mut self, schema: Value) -> Self {
75 self.schema = Some(schema);
76 self
77 }
78}
79
80pub(crate) fn no_schema_response(command_path: &str) -> Value {
89 serde_json::json!({
90 "command": command_path,
91 "fields": [],
92 "message": "No output schema is registered for this command.",
93 })
94}
95
96#[derive(Clone, Copy, Debug, Eq, PartialEq)]
98pub struct OutputField {
99 pub name: &'static str,
101 pub field_type: &'static str,
103 pub optional: bool,
105}
106
107impl OutputField {
108 #[must_use]
110 pub const fn new(name: &'static str, field_type: &'static str) -> Self {
111 Self {
112 name,
113 field_type,
114 optional: false,
115 }
116 }
117
118 #[must_use]
120 pub const fn optional(mut self) -> Self {
121 self.optional = true;
122 self
123 }
124
125 #[must_use]
127 pub const fn string(name: &'static str) -> Self {
128 Self::new(name, "string")
129 }
130
131 #[must_use]
133 pub const fn int(name: &'static str) -> Self {
134 Self::new(name, "int")
135 }
136
137 #[must_use]
139 pub const fn float(name: &'static str) -> Self {
140 Self::new(name, "float")
141 }
142
143 #[must_use]
145 pub const fn bool(name: &'static str) -> Self {
146 Self::new(name, "bool")
147 }
148
149 #[must_use]
151 pub const fn list(name: &'static str, field_type: &'static str) -> Self {
152 Self::new(name, field_type)
153 }
154
155 #[must_use]
157 pub const fn string_list(name: &'static str) -> Self {
158 Self::new(name, "[]string")
159 }
160
161 #[must_use]
163 pub const fn int_list(name: &'static str) -> Self {
164 Self::new(name, "[]int")
165 }
166
167 #[must_use]
169 pub const fn float_list(name: &'static str) -> Self {
170 Self::new(name, "[]float")
171 }
172
173 #[must_use]
175 pub const fn bool_list(name: &'static str) -> Self {
176 Self::new(name, "[]bool")
177 }
178
179 #[must_use]
181 pub const fn object_list(name: &'static str) -> Self {
182 Self::new(name, "[]object")
183 }
184
185 #[must_use]
187 pub const fn object(name: &'static str) -> Self {
188 Self::new(name, "object")
189 }
190
191 #[must_use]
193 pub const fn any(name: &'static str) -> Self {
194 Self::new(name, "any")
195 }
196}
197
198pub trait OutputSchema {
200 fn fields() -> &'static [OutputField];
202}
203
204#[derive(Clone, Debug, Default)]
206pub struct SchemaRegistry {
207 by_path: BTreeMap<String, SchemaInfo>,
208}
209
210impl SchemaRegistry {
211 #[must_use]
213 pub fn new() -> Self {
214 Self::default()
215 }
216
217 pub fn register<T: OutputSchema>(&mut self, command_path: impl Into<String>) {
219 self.register_fields(command_path, fields_for::<T>());
220 }
221
222 pub fn register_json_schema<T: JsonSchema>(&mut self, command_path: impl Into<String>) {
224 let command_path = command_path.into();
225 self.register_info(command_path.clone(), json_schema_info::<T>(command_path));
226 }
227
228 pub fn register_fields(
230 &mut self,
231 command_path: impl Into<String>,
232 fields: impl Into<Vec<FieldInfo>>,
233 ) {
234 let command_path = command_path.into();
235 self.register_info(
236 command_path.clone(),
237 SchemaInfo {
238 command: command_path,
239 fields: fields.into(),
240 schema: None,
241 },
242 );
243 }
244
245 pub fn register_info(&mut self, command_path: impl Into<String>, mut info: SchemaInfo) {
247 let command_path = command_path.into();
248 info.command = command_path.clone();
249 self.by_path.insert(command_path, info);
250 }
251
252 pub fn merge(&mut self, other: &Self) {
254 self.by_path.extend(other.by_path.clone());
255 }
256
257 #[must_use]
259 pub fn get_by_path(&self, command_path: &str) -> Option<SchemaInfo> {
260 if let Some(info) = self.by_path.get(command_path) {
261 return Some(schema_with_command(info, command_path));
262 }
263
264 let space_path = command_path.replace(':', " ");
265 self.by_path.iter().find_map(|(registered, info)| {
266 let matches = registered == &space_path
267 || registered
268 .split_once(' ')
269 .is_some_and(|(_, without_root)| without_root == space_path);
270 matches.then(|| schema_with_command(info, command_path))
271 })
272 }
273}
274
275fn schema_with_command(info: &SchemaInfo, command_path: &str) -> SchemaInfo {
276 SchemaInfo {
277 command: command_path.to_owned(),
278 fields: info.fields.clone(),
279 schema: info.schema.clone(),
280 }
281}
282
283#[must_use]
285pub fn fields_for<T: OutputSchema>() -> Vec<FieldInfo> {
286 T::fields()
287 .iter()
288 .map(|field| FieldInfo {
289 name: field.name.to_owned(),
290 field_type: field.field_type.to_owned(),
291 optional: field.optional,
292 })
293 .collect()
294}
295
296#[must_use]
298pub fn json_schema_for<T: JsonSchema>() -> Value {
299 serde_json::to_value(schemars::schema_for!(T)).unwrap_or(Value::Null)
300}
301
302#[must_use]
304pub fn json_schema_info<T: JsonSchema>(command_path: impl Into<String>) -> SchemaInfo {
305 let schema = json_schema_for::<T>();
306 SchemaInfo {
307 command: command_path.into(),
308 fields: fields_from_json_schema(&schema),
309 schema: Some(schema),
310 }
311}
312
313#[must_use]
315pub fn fields_from_json_schema(schema: &Value) -> Vec<FieldInfo> {
316 let Some(properties) = schema.get("properties").and_then(Value::as_object) else {
317 return Vec::new();
318 };
319 let required = schema
320 .get("required")
321 .and_then(Value::as_array)
322 .map(|items| {
323 items
324 .iter()
325 .filter_map(Value::as_str)
326 .collect::<BTreeSet<_>>()
327 })
328 .unwrap_or_default();
329
330 properties
331 .iter()
332 .map(|(name, property)| FieldInfo {
333 name: name.clone(),
334 field_type: json_schema_type_name(property, schema),
335 optional: !required.contains(name.as_str()) || schema_allows_null(property),
336 })
337 .collect()
338}
339
340fn json_schema_type_name(schema: &Value, root: &Value) -> String {
341 let schema = non_null_schema(schema);
342 let schema = resolve_local_ref(schema, root).unwrap_or(schema);
343 match primary_json_type(schema).as_deref() {
344 Some("string") => "string".to_owned(),
345 Some("integer") => "int".to_owned(),
346 Some("number") => "float".to_owned(),
347 Some("boolean") => "bool".to_owned(),
348 Some("array") => {
349 let item_type = schema
350 .get("items")
351 .map(|items| json_schema_type_name(items, root))
352 .unwrap_or_else(|| "any".to_owned());
353 format!("[]{item_type}")
354 }
355 Some("object") => "object".to_owned(),
356 Some(other) => other.to_owned(),
357 None => {
358 if schema.get("properties").is_some() {
359 "object".to_owned()
360 } else {
361 "any".to_owned()
362 }
363 }
364 }
365}
366
367fn non_null_schema(schema: &Value) -> &Value {
368 for key in ["anyOf", "oneOf"] {
369 if let Some(items) = schema.get(key).and_then(Value::as_array)
370 && let Some(non_null) = items
371 .iter()
372 .find(|item| item.get("type").and_then(Value::as_str) != Some("null"))
373 {
374 return non_null;
375 }
376 }
377 schema
378}
379
380fn resolve_local_ref<'schema>(
381 schema: &'schema Value,
382 root: &'schema Value,
383) -> Option<&'schema Value> {
384 let reference = schema.get("$ref").and_then(Value::as_str)?;
385 let pointer = reference.strip_prefix('#')?;
386 root.pointer(pointer)
387}
388
389fn primary_json_type(schema: &Value) -> Option<String> {
390 match schema.get("type") {
391 Some(Value::String(value)) => Some(value.clone()),
392 Some(Value::Array(values)) => values
393 .iter()
394 .filter_map(Value::as_str)
395 .find(|value| *value != "null")
396 .map(str::to_owned),
397 Some(Value::Null | Value::Bool(_) | Value::Number(_) | Value::Object(_)) | None => None,
398 }
399}
400
401fn schema_allows_null(schema: &Value) -> bool {
402 matches!(schema.get("type"), Some(Value::String(value)) if value == "null")
403 || schema
404 .get("type")
405 .and_then(Value::as_array)
406 .is_some_and(|items| items.iter().any(|item| item.as_str() == Some("null")))
407 || ["anyOf", "oneOf"].iter().any(|key| {
408 schema
409 .get(key)
410 .and_then(Value::as_array)
411 .is_some_and(|items| {
412 items
413 .iter()
414 .any(|item| item.get("type").and_then(Value::as_str) == Some("null"))
415 })
416 })
417}
418
419static GLOBAL_SCHEMA_REGISTRY: OnceLock<RwLock<SchemaRegistry>> = OnceLock::new();
420
421fn global_schema_registry() -> &'static RwLock<SchemaRegistry> {
422 GLOBAL_SCHEMA_REGISTRY.get_or_init(|| RwLock::new(SchemaRegistry::new()))
423}
424
425pub fn register_global_schema<T: OutputSchema>(command_path: impl Into<String>) {
427 let mut registry = global_schema_registry()
428 .write()
429 .unwrap_or_else(|poisoned| poisoned.into_inner());
430 registry.register::<T>(command_path);
431}
432
433pub fn register_global_json_schema<T: JsonSchema>(command_path: impl Into<String>) {
435 let mut registry = global_schema_registry()
436 .write()
437 .unwrap_or_else(|poisoned| poisoned.into_inner());
438 registry.register_json_schema::<T>(command_path);
439}
440
441pub fn register_global_schema_fields(
443 command_path: impl Into<String>,
444 fields: impl Into<Vec<FieldInfo>>,
445) {
446 let mut registry = global_schema_registry()
447 .write()
448 .unwrap_or_else(|poisoned| poisoned.into_inner());
449 registry.register_fields(command_path, fields);
450}
451
452pub fn register_global_schema_info(command_path: impl Into<String>, info: SchemaInfo) {
454 let mut registry = global_schema_registry()
455 .write()
456 .unwrap_or_else(|poisoned| poisoned.into_inner());
457 registry.register_info(command_path, info);
458}
459
460#[must_use]
462pub fn get_global_schema_by_path(command_path: &str) -> Option<SchemaInfo> {
463 global_schema_registry()
464 .read()
465 .unwrap_or_else(|poisoned| poisoned.into_inner())
466 .get_by_path(command_path)
467}
468
469#[must_use]
471pub fn global_schema_registry_snapshot() -> SchemaRegistry {
472 global_schema_registry()
473 .read()
474 .unwrap_or_else(|poisoned| poisoned.into_inner())
475 .clone()
476}
477
478#[must_use]
480pub fn format_help_section(fields: &[FieldInfo]) -> String {
481 if fields.is_empty() {
482 return String::new();
483 }
484 let max_name = fields
485 .iter()
486 .map(|field| field.name.len())
487 .max()
488 .unwrap_or_default();
489 let mut out = String::from("Output fields:\n");
490 for field in fields {
491 let optional = if field.optional { " (optional)" } else { "" };
492 out.push_str(&format!(
493 " {:<width$} {}{}\n",
494 field.name,
495 field.field_type,
496 optional,
497 width = max_name
498 ));
499 }
500
501 let first_string = fields
502 .iter()
503 .find(|field| field.field_type == "string")
504 .map(|field| field.name.as_str());
505 let first_bool = fields
506 .iter()
507 .find(|field| field.field_type == "bool")
508 .map(|field| field.name.as_str());
509 if first_string.is_some() || first_bool.is_some() {
510 out.push_str("\nFilter examples:\n");
511 if let Some(name) = first_string {
512 out.push_str(&format!(" --filter \"contains({name}, 'example')\"\n"));
513 }
514 if let Some(name) = first_bool {
515 out.push_str(&format!(" --filter '{name}'\n"));
516 }
517 }
518
519 out.push_str("\nExpr examples:\n");
520 out.push_str(" --expr 'length(@)'\n");
521 if let Some(name) = first_string {
522 out.push_str(&format!(" --expr '[].{name}'\n"));
523 }
524 out
525}