1use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event};
7use quick_xml::Writer;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::io::Cursor;
11
12use super::event::CotError;
13use super::types::OperationalStatus;
14
15pub const PEAT_EXTENSION_VERSION: &str = "1.0";
17
18#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
22pub struct PeatExtension {
23 pub source: Option<PeatSource>,
25 pub confidence: Option<PeatConfidence>,
27 pub hierarchy: Option<PeatHierarchy>,
29 pub attributes: HashMap<String, PeatAttribute>,
31 pub status: Option<PeatStatus>,
33 pub capabilities: Vec<PeatCapability>,
35 pub handoff: Option<PeatHandoff>,
37 pub classification: Option<PeatClassification>,
39}
40
41impl PeatExtension {
42 pub fn new() -> Self {
44 Self::default()
45 }
46
47 pub fn with_source(mut self, source: PeatSource) -> Self {
49 self.source = Some(source);
50 self
51 }
52
53 pub fn with_confidence(mut self, value: f64, threshold: Option<f64>) -> Self {
55 self.confidence = Some(PeatConfidence { value, threshold });
56 self
57 }
58
59 pub fn with_hierarchy(mut self, hierarchy: PeatHierarchy) -> Self {
61 self.hierarchy = Some(hierarchy);
62 self
63 }
64
65 pub fn with_attribute(mut self, key: &str, value: &str, attr_type: &str) -> Self {
67 self.attributes.insert(
68 key.to_string(),
69 PeatAttribute {
70 value: value.to_string(),
71 attr_type: attr_type.to_string(),
72 },
73 );
74 self
75 }
76
77 pub fn with_status(mut self, status: PeatStatus) -> Self {
79 self.status = Some(status);
80 self
81 }
82
83 pub fn with_capability(mut self, capability: PeatCapability) -> Self {
85 self.capabilities.push(capability);
86 self
87 }
88
89 pub fn with_handoff(mut self, handoff: PeatHandoff) -> Self {
91 self.handoff = Some(handoff);
92 self
93 }
94
95 pub fn with_classification(mut self, level: &str, caveat: Option<&str>) -> Self {
97 self.classification = Some(PeatClassification {
98 level: level.to_string(),
99 caveat: caveat.map(|s| s.to_string()),
100 });
101 self
102 }
103
104 pub fn write_xml(&self, writer: &mut Writer<Cursor<Vec<u8>>>) -> Result<(), CotError> {
106 let mut peat_elem = BytesStart::new("_peat_");
107 peat_elem.push_attribute(("version", PEAT_EXTENSION_VERSION));
108
109 writer
110 .write_event(Event::Start(peat_elem))
111 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
112
113 if let Some(ref source) = self.source {
115 let mut src_elem = BytesStart::new("source");
116 src_elem.push_attribute(("platform", source.platform.as_str()));
117 src_elem.push_attribute(("model", source.model.as_str()));
118 src_elem.push_attribute(("model_version", source.model_version.as_str()));
119
120 writer
121 .write_event(Event::Empty(src_elem))
122 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
123 }
124
125 if let Some(ref conf) = self.confidence {
127 let mut conf_elem = BytesStart::new("confidence");
128 conf_elem.push_attribute(("value", conf.value.to_string().as_str()));
129 if let Some(threshold) = conf.threshold {
130 conf_elem.push_attribute(("threshold", threshold.to_string().as_str()));
131 }
132
133 writer
134 .write_event(Event::Empty(conf_elem))
135 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
136 }
137
138 if let Some(ref hier) = self.hierarchy {
140 writer
141 .write_event(Event::Start(BytesStart::new("hierarchy")))
142 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
143
144 if let Some(ref cell) = hier.cell {
145 let mut cell_elem = BytesStart::new("cell");
146 cell_elem.push_attribute(("id", cell.id.as_str()));
147 if let Some(ref role) = cell.role {
148 cell_elem.push_attribute(("role", role.as_str()));
149 }
150 writer
151 .write_event(Event::Empty(cell_elem))
152 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
153 }
154
155 if let Some(ref formation) = hier.formation {
156 let mut form_elem = BytesStart::new("formation");
157 form_elem.push_attribute(("id", formation.as_str()));
158 writer
159 .write_event(Event::Empty(form_elem))
160 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
161 }
162
163 if let Some(ref zone) = hier.zone {
164 let mut zone_elem = BytesStart::new("zone");
165 zone_elem.push_attribute(("id", zone.as_str()));
166 writer
167 .write_event(Event::Empty(zone_elem))
168 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
169 }
170
171 writer
172 .write_event(Event::End(BytesEnd::new("hierarchy")))
173 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
174 }
175
176 if !self.attributes.is_empty() {
178 writer
179 .write_event(Event::Start(BytesStart::new("attributes")))
180 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
181
182 for (key, attr) in &self.attributes {
183 let mut attr_elem = BytesStart::new("attr");
184 attr_elem.push_attribute(("key", key.as_str()));
185 attr_elem.push_attribute(("type", attr.attr_type.as_str()));
186
187 writer
188 .write_event(Event::Start(attr_elem))
189 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
190 writer
191 .write_event(Event::Text(BytesText::new(&attr.value)))
192 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
193 writer
194 .write_event(Event::End(BytesEnd::new("attr")))
195 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
196 }
197
198 writer
199 .write_event(Event::End(BytesEnd::new("attributes")))
200 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
201 }
202
203 if let Some(ref status) = self.status {
205 let mut status_elem = BytesStart::new("status");
206 status_elem.push_attribute(("operational", status.operational.as_str()));
207 status_elem.push_attribute(("readiness", status.readiness.to_string().as_str()));
208
209 writer
210 .write_event(Event::Empty(status_elem))
211 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
212 }
213
214 for cap in &self.capabilities {
216 let mut cap_elem = BytesStart::new("capability");
217 cap_elem.push_attribute(("type", cap.capability_type.as_str()));
218 cap_elem.push_attribute(("model_version", cap.model_version.as_str()));
219 cap_elem.push_attribute(("precision", cap.precision.to_string().as_str()));
220 cap_elem.push_attribute(("status", cap.status.as_str()));
221
222 writer
223 .write_event(Event::Empty(cap_elem))
224 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
225 }
226
227 if let Some(ref handoff) = self.handoff {
229 let mut handoff_elem = BytesStart::new("handoff");
230 handoff_elem.push_attribute(("source_cell", handoff.source_cell.as_str()));
231 handoff_elem.push_attribute(("target_cell", handoff.target_cell.as_str()));
232 handoff_elem.push_attribute(("state", handoff.state.as_str()));
233 handoff_elem.push_attribute(("reason", handoff.reason.as_str()));
234
235 writer
236 .write_event(Event::Empty(handoff_elem))
237 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
238 }
239
240 if let Some(ref class) = self.classification {
242 let mut class_elem = BytesStart::new("classification");
243 class_elem.push_attribute(("level", class.level.as_str()));
244 if let Some(ref caveat) = class.caveat {
245 class_elem.push_attribute(("caveat", caveat.as_str()));
246 }
247
248 writer
249 .write_event(Event::Empty(class_elem))
250 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
251 }
252
253 writer
254 .write_event(Event::End(BytesEnd::new("_peat_")))
255 .map_err(|e| CotError::XmlWrite(e.to_string()))?;
256
257 Ok(())
258 }
259}
260
261#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
263pub struct PeatSource {
264 pub platform: String,
266 pub model: String,
268 pub model_version: String,
270}
271
272impl PeatSource {
273 pub fn new(platform: &str, model: &str, model_version: &str) -> Self {
275 Self {
276 platform: platform.to_string(),
277 model: model.to_string(),
278 model_version: model_version.to_string(),
279 }
280 }
281}
282
283#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
285pub struct PeatConfidence {
286 pub value: f64,
288 pub threshold: Option<f64>,
290}
291
292#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
294pub struct PeatHierarchy {
295 pub cell: Option<PeatCellMembership>,
297 pub formation: Option<String>,
299 pub zone: Option<String>,
301}
302
303impl PeatHierarchy {
304 pub fn new() -> Self {
306 Self::default()
307 }
308
309 pub fn with_cell(mut self, id: &str, role: Option<&str>) -> Self {
311 self.cell = Some(PeatCellMembership {
312 id: id.to_string(),
313 role: role.map(|s| s.to_string()),
314 });
315 self
316 }
317
318 pub fn with_formation(mut self, id: &str) -> Self {
320 self.formation = Some(id.to_string());
321 self
322 }
323
324 pub fn with_zone(mut self, id: &str) -> Self {
326 self.zone = Some(id.to_string());
327 self
328 }
329}
330
331#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
333pub struct PeatCellMembership {
334 pub id: String,
336 pub role: Option<String>,
338}
339
340#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
342pub struct PeatAttribute {
343 pub value: String,
345 pub attr_type: String,
347}
348
349#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
351pub struct PeatStatus {
352 pub operational: OperationalStatus,
354 pub readiness: f64,
356}
357
358impl PeatStatus {
359 pub fn new(operational: OperationalStatus, readiness: f64) -> Self {
361 Self {
362 operational,
363 readiness: readiness.clamp(0.0, 1.0),
364 }
365 }
366}
367
368#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
370pub struct PeatCapability {
371 pub capability_type: String,
373 pub model_version: String,
375 pub precision: f64,
377 pub status: OperationalStatus,
379}
380
381impl PeatCapability {
382 pub fn new(
384 capability_type: &str,
385 model_version: &str,
386 precision: f64,
387 status: OperationalStatus,
388 ) -> Self {
389 Self {
390 capability_type: capability_type.to_string(),
391 model_version: model_version.to_string(),
392 precision: precision.clamp(0.0, 1.0),
393 status,
394 }
395 }
396}
397
398#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
400pub struct PeatHandoff {
401 pub source_cell: String,
403 pub target_cell: String,
405 pub state: String,
407 pub reason: String,
409}
410
411impl PeatHandoff {
412 pub fn new(source_cell: &str, target_cell: &str, state: &str, reason: &str) -> Self {
414 Self {
415 source_cell: source_cell.to_string(),
416 target_cell: target_cell.to_string(),
417 state: state.to_string(),
418 reason: reason.to_string(),
419 }
420 }
421}
422
423#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
425pub struct PeatClassification {
426 pub level: String,
428 pub caveat: Option<String>,
430}
431
432#[cfg(test)]
433mod tests {
434 use super::*;
435
436 #[test]
437 fn test_peat_extension_creation() {
438 let ext = PeatExtension::new()
439 .with_source(PeatSource::new("Alpha-2", "object_tracker", "1.3.0"))
440 .with_confidence(0.89, Some(0.70))
441 .with_hierarchy(
442 PeatHierarchy::new()
443 .with_cell("Alpha-Team", Some("tracker"))
444 .with_formation("Formation-1"),
445 )
446 .with_attribute("jacket_color", "blue", "string")
447 .with_status(PeatStatus::new(OperationalStatus::Active, 0.91));
448
449 assert!(ext.source.is_some());
450 assert_eq!(ext.source.as_ref().unwrap().platform, "Alpha-2");
451 assert!(ext.confidence.is_some());
452 assert_eq!(ext.confidence.as_ref().unwrap().value, 0.89);
453 }
454
455 #[test]
456 fn test_peat_extension_to_xml() {
457 let ext = PeatExtension::new()
458 .with_source(PeatSource::new("Platform-1", "sensor", "1.0.0"))
459 .with_confidence(0.85, None);
460
461 let mut writer = Writer::new(Cursor::new(Vec::new()));
462 ext.write_xml(&mut writer).unwrap();
463
464 let xml = String::from_utf8(writer.into_inner().into_inner()).unwrap();
465 assert!(xml.contains("<_peat_"));
466 assert!(xml.contains("version=\"1.0\""));
467 assert!(xml.contains("platform=\"Platform-1\""));
468 assert!(xml.contains("value=\"0.85\""));
469 }
470
471 #[test]
472 fn test_peat_hierarchy() {
473 let hier = PeatHierarchy::new()
474 .with_cell("Alpha-Team", Some("leader"))
475 .with_formation("Formation-1")
476 .with_zone("Zone-A");
477
478 assert_eq!(hier.cell.as_ref().unwrap().id, "Alpha-Team");
479 assert_eq!(hier.cell.as_ref().unwrap().role, Some("leader".to_string()));
480 assert_eq!(hier.formation, Some("Formation-1".to_string()));
481 assert_eq!(hier.zone, Some("Zone-A".to_string()));
482 }
483
484 #[test]
485 fn test_peat_capability() {
486 let cap = PeatCapability::new("OBJECT_TRACKING", "1.3.0", 0.94, OperationalStatus::Active);
487
488 assert_eq!(cap.capability_type, "OBJECT_TRACKING");
489 assert_eq!(cap.precision, 0.94);
490 }
491
492 #[test]
493 fn test_peat_status_readiness_clamped() {
494 let status = PeatStatus::new(OperationalStatus::Active, 1.5);
495 assert_eq!(status.readiness, 1.0);
496
497 let status2 = PeatStatus::new(OperationalStatus::Degraded, -0.5);
498 assert_eq!(status2.readiness, 0.0);
499 }
500
501 #[test]
502 fn test_peat_extension_with_attributes() {
503 let ext = PeatExtension::new()
504 .with_attribute("color", "red", "string")
505 .with_attribute("count", "5", "number")
506 .with_attribute("active", "true", "boolean");
507
508 assert_eq!(ext.attributes.len(), 3);
509 assert_eq!(ext.attributes["color"].value, "red");
510 assert_eq!(ext.attributes["color"].attr_type, "string");
511 }
512
513 #[test]
514 fn test_peat_handoff() {
515 let handoff =
516 PeatHandoff::new("Alpha-Team", "Bravo-Team", "INITIATED", "boundary_crossing");
517
518 assert_eq!(handoff.source_cell, "Alpha-Team");
519 assert_eq!(handoff.target_cell, "Bravo-Team");
520 assert_eq!(handoff.state, "INITIATED");
521 }
522
523 #[test]
524 fn test_peat_classification() {
525 let ext = PeatExtension::new().with_classification("UNCLASSIFIED", Some("FOUO"));
526
527 assert!(ext.classification.is_some());
528 let class = ext.classification.unwrap();
529 assert_eq!(class.level, "UNCLASSIFIED");
530 assert_eq!(class.caveat, Some("FOUO".to_string()));
531 }
532}