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
80#[derive(Clone, Copy, Debug, Eq, PartialEq)]
82pub struct OutputField {
83 pub name: &'static str,
85 pub field_type: &'static str,
87 pub optional: bool,
89}
90
91impl OutputField {
92 #[must_use]
94 pub const fn new(name: &'static str, field_type: &'static str) -> Self {
95 Self {
96 name,
97 field_type,
98 optional: false,
99 }
100 }
101
102 #[must_use]
104 pub const fn optional(mut self) -> Self {
105 self.optional = true;
106 self
107 }
108
109 #[must_use]
111 pub const fn string(name: &'static str) -> Self {
112 Self::new(name, "string")
113 }
114
115 #[must_use]
117 pub const fn int(name: &'static str) -> Self {
118 Self::new(name, "int")
119 }
120
121 #[must_use]
123 pub const fn float(name: &'static str) -> Self {
124 Self::new(name, "float")
125 }
126
127 #[must_use]
129 pub const fn bool(name: &'static str) -> Self {
130 Self::new(name, "bool")
131 }
132
133 #[must_use]
135 pub const fn list(name: &'static str, field_type: &'static str) -> Self {
136 Self::new(name, field_type)
137 }
138
139 #[must_use]
141 pub const fn string_list(name: &'static str) -> Self {
142 Self::new(name, "[]string")
143 }
144
145 #[must_use]
147 pub const fn int_list(name: &'static str) -> Self {
148 Self::new(name, "[]int")
149 }
150
151 #[must_use]
153 pub const fn float_list(name: &'static str) -> Self {
154 Self::new(name, "[]float")
155 }
156
157 #[must_use]
159 pub const fn bool_list(name: &'static str) -> Self {
160 Self::new(name, "[]bool")
161 }
162
163 #[must_use]
165 pub const fn object_list(name: &'static str) -> Self {
166 Self::new(name, "[]object")
167 }
168
169 #[must_use]
171 pub const fn object(name: &'static str) -> Self {
172 Self::new(name, "object")
173 }
174
175 #[must_use]
177 pub const fn any(name: &'static str) -> Self {
178 Self::new(name, "any")
179 }
180}
181
182pub trait OutputSchema {
184 fn fields() -> &'static [OutputField];
186}
187
188#[derive(Clone, Debug, Default)]
190pub struct SchemaRegistry {
191 by_path: BTreeMap<String, SchemaInfo>,
192}
193
194impl SchemaRegistry {
195 #[must_use]
197 pub fn new() -> Self {
198 Self::default()
199 }
200
201 pub fn register<T: OutputSchema>(&mut self, command_path: impl Into<String>) {
203 self.register_fields(command_path, fields_for::<T>());
204 }
205
206 pub fn register_json_schema<T: JsonSchema>(&mut self, command_path: impl Into<String>) {
208 let command_path = command_path.into();
209 self.register_info(command_path.clone(), json_schema_info::<T>(command_path));
210 }
211
212 pub fn register_fields(
214 &mut self,
215 command_path: impl Into<String>,
216 fields: impl Into<Vec<FieldInfo>>,
217 ) {
218 let command_path = command_path.into();
219 self.register_info(
220 command_path.clone(),
221 SchemaInfo {
222 command: command_path,
223 fields: fields.into(),
224 schema: None,
225 },
226 );
227 }
228
229 pub fn register_info(&mut self, command_path: impl Into<String>, mut info: SchemaInfo) {
231 let command_path = command_path.into();
232 info.command = command_path.clone();
233 self.by_path.insert(command_path, info);
234 }
235
236 pub fn merge(&mut self, other: &Self) {
238 self.by_path.extend(other.by_path.clone());
239 }
240
241 #[must_use]
243 pub fn get_by_path(&self, command_path: &str) -> Option<SchemaInfo> {
244 if let Some(info) = self.by_path.get(command_path) {
245 return Some(schema_with_command(info, command_path));
246 }
247
248 let space_path = command_path.replace(':', " ");
249 self.by_path.iter().find_map(|(registered, info)| {
250 let matches = registered == &space_path
251 || registered
252 .split_once(' ')
253 .is_some_and(|(_, without_root)| without_root == space_path);
254 matches.then(|| schema_with_command(info, command_path))
255 })
256 }
257}
258
259fn schema_with_command(info: &SchemaInfo, command_path: &str) -> SchemaInfo {
260 SchemaInfo {
261 command: command_path.to_owned(),
262 fields: info.fields.clone(),
263 schema: info.schema.clone(),
264 }
265}
266
267#[must_use]
269pub fn fields_for<T: OutputSchema>() -> Vec<FieldInfo> {
270 T::fields()
271 .iter()
272 .map(|field| FieldInfo {
273 name: field.name.to_owned(),
274 field_type: field.field_type.to_owned(),
275 optional: field.optional,
276 })
277 .collect()
278}
279
280#[must_use]
282pub fn json_schema_for<T: JsonSchema>() -> Value {
283 serde_json::to_value(schemars::schema_for!(T)).unwrap_or(Value::Null)
284}
285
286#[must_use]
288pub fn json_schema_info<T: JsonSchema>(command_path: impl Into<String>) -> SchemaInfo {
289 let schema = json_schema_for::<T>();
290 SchemaInfo {
291 command: command_path.into(),
292 fields: fields_from_json_schema(&schema),
293 schema: Some(schema),
294 }
295}
296
297#[must_use]
299pub fn fields_from_json_schema(schema: &Value) -> Vec<FieldInfo> {
300 let Some(properties) = schema.get("properties").and_then(Value::as_object) else {
301 return Vec::new();
302 };
303 let required = schema
304 .get("required")
305 .and_then(Value::as_array)
306 .map(|items| {
307 items
308 .iter()
309 .filter_map(Value::as_str)
310 .collect::<BTreeSet<_>>()
311 })
312 .unwrap_or_default();
313
314 properties
315 .iter()
316 .map(|(name, property)| FieldInfo {
317 name: name.clone(),
318 field_type: json_schema_type_name(property, schema),
319 optional: !required.contains(name.as_str()) || schema_allows_null(property),
320 })
321 .collect()
322}
323
324fn json_schema_type_name(schema: &Value, root: &Value) -> String {
325 let schema = non_null_schema(schema);
326 let schema = resolve_local_ref(schema, root).unwrap_or(schema);
327 match primary_json_type(schema).as_deref() {
328 Some("string") => "string".to_owned(),
329 Some("integer") => "int".to_owned(),
330 Some("number") => "float".to_owned(),
331 Some("boolean") => "bool".to_owned(),
332 Some("array") => {
333 let item_type = schema
334 .get("items")
335 .map(|items| json_schema_type_name(items, root))
336 .unwrap_or_else(|| "any".to_owned());
337 format!("[]{item_type}")
338 }
339 Some("object") => "object".to_owned(),
340 Some(other) => other.to_owned(),
341 None => {
342 if schema.get("properties").is_some() {
343 "object".to_owned()
344 } else {
345 "any".to_owned()
346 }
347 }
348 }
349}
350
351fn non_null_schema(schema: &Value) -> &Value {
352 for key in ["anyOf", "oneOf"] {
353 if let Some(items) = schema.get(key).and_then(Value::as_array)
354 && let Some(non_null) = items
355 .iter()
356 .find(|item| item.get("type").and_then(Value::as_str) != Some("null"))
357 {
358 return non_null;
359 }
360 }
361 schema
362}
363
364fn resolve_local_ref<'schema>(
365 schema: &'schema Value,
366 root: &'schema Value,
367) -> Option<&'schema Value> {
368 let reference = schema.get("$ref").and_then(Value::as_str)?;
369 let pointer = reference.strip_prefix('#')?;
370 root.pointer(pointer)
371}
372
373fn primary_json_type(schema: &Value) -> Option<String> {
374 match schema.get("type") {
375 Some(Value::String(value)) => Some(value.clone()),
376 Some(Value::Array(values)) => values
377 .iter()
378 .filter_map(Value::as_str)
379 .find(|value| *value != "null")
380 .map(str::to_owned),
381 Some(Value::Null | Value::Bool(_) | Value::Number(_) | Value::Object(_)) | None => None,
382 }
383}
384
385fn schema_allows_null(schema: &Value) -> bool {
386 matches!(schema.get("type"), Some(Value::String(value)) if value == "null")
387 || schema
388 .get("type")
389 .and_then(Value::as_array)
390 .is_some_and(|items| items.iter().any(|item| item.as_str() == Some("null")))
391 || ["anyOf", "oneOf"].iter().any(|key| {
392 schema
393 .get(key)
394 .and_then(Value::as_array)
395 .is_some_and(|items| {
396 items
397 .iter()
398 .any(|item| item.get("type").and_then(Value::as_str) == Some("null"))
399 })
400 })
401}
402
403static GLOBAL_SCHEMA_REGISTRY: OnceLock<RwLock<SchemaRegistry>> = OnceLock::new();
404
405fn global_schema_registry() -> &'static RwLock<SchemaRegistry> {
406 GLOBAL_SCHEMA_REGISTRY.get_or_init(|| RwLock::new(SchemaRegistry::new()))
407}
408
409pub fn register_global_schema<T: OutputSchema>(command_path: impl Into<String>) {
411 let mut registry = global_schema_registry()
412 .write()
413 .unwrap_or_else(|poisoned| poisoned.into_inner());
414 registry.register::<T>(command_path);
415}
416
417pub fn register_global_json_schema<T: JsonSchema>(command_path: impl Into<String>) {
419 let mut registry = global_schema_registry()
420 .write()
421 .unwrap_or_else(|poisoned| poisoned.into_inner());
422 registry.register_json_schema::<T>(command_path);
423}
424
425pub fn register_global_schema_fields(
427 command_path: impl Into<String>,
428 fields: impl Into<Vec<FieldInfo>>,
429) {
430 let mut registry = global_schema_registry()
431 .write()
432 .unwrap_or_else(|poisoned| poisoned.into_inner());
433 registry.register_fields(command_path, fields);
434}
435
436pub fn register_global_schema_info(command_path: impl Into<String>, info: SchemaInfo) {
438 let mut registry = global_schema_registry()
439 .write()
440 .unwrap_or_else(|poisoned| poisoned.into_inner());
441 registry.register_info(command_path, info);
442}
443
444#[must_use]
446pub fn get_global_schema_by_path(command_path: &str) -> Option<SchemaInfo> {
447 global_schema_registry()
448 .read()
449 .unwrap_or_else(|poisoned| poisoned.into_inner())
450 .get_by_path(command_path)
451}
452
453#[must_use]
455pub fn global_schema_registry_snapshot() -> SchemaRegistry {
456 global_schema_registry()
457 .read()
458 .unwrap_or_else(|poisoned| poisoned.into_inner())
459 .clone()
460}
461
462#[must_use]
464pub fn format_help_section(fields: &[FieldInfo]) -> String {
465 if fields.is_empty() {
466 return String::new();
467 }
468 let max_name = fields
469 .iter()
470 .map(|field| field.name.len())
471 .max()
472 .unwrap_or_default();
473 let mut out = String::from("Output fields:\n");
474 for field in fields {
475 let optional = if field.optional { " (optional)" } else { "" };
476 out.push_str(&format!(
477 " {:<width$} {}{}\n",
478 field.name,
479 field.field_type,
480 optional,
481 width = max_name
482 ));
483 }
484
485 let first_string = fields
486 .iter()
487 .find(|field| field.field_type == "string")
488 .map(|field| field.name.as_str());
489 let first_bool = fields
490 .iter()
491 .find(|field| field.field_type == "bool")
492 .map(|field| field.name.as_str());
493 if first_string.is_some() || first_bool.is_some() {
494 out.push_str("\nFilter examples:\n");
495 if let Some(name) = first_string {
496 out.push_str(&format!(" --filter \"contains({name}, 'example')\"\n"));
497 }
498 if let Some(name) = first_bool {
499 out.push_str(&format!(" --filter '{name}'\n"));
500 }
501 }
502
503 out.push_str("\nExpr examples:\n");
504 out.push_str(" --expr 'length(@)'\n");
505 if let Some(name) = first_string {
506 out.push_str(&format!(" --expr '[].{name}'\n"));
507 }
508 out
509}