tenant.openagents/omega

1use anyhow::Result;
2use quick_xml::events::{BytesStart, Event};
3
4use super::{accent_class_name, add_to_event};
5
6pub(super) struct SequenceDiagramAccents {
7    accent_count: usize,
8    actor_bottom_counter: usize,
9    actor_top_counter: usize,
10    last_actor_accent: Option<usize>,
11    current_text_accent: Option<usize>,
12}
13
14impl SequenceDiagramAccents {
15    pub(super) fn new(accent_count: usize) -> Self {
16        Self {
17            accent_count,
18            actor_bottom_counter: 0,
19            actor_top_counter: 0,
20            last_actor_accent: None,
21            current_text_accent: None,
22        }
23    }
24
25    pub(super) fn process_event<'a>(&mut self, event: Event<'a>) -> Result<Event<'a>> {
26        if self.accent_count == 0 {
27            return Ok(event);
28        }
29
30        match &event {
31            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"rect" => {
32                if let Some(idx) = self.check_actor_rect(e)? {
33                    add_to_event(&event, e, &accent_class_name(idx))
34                } else {
35                    Ok(event)
36                }
37            }
38
39            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"text" => {
40                if let Some(idx) = self.check_actor_text(e)? {
41                    self.current_text_accent = Some(idx);
42                    add_to_event(&event, e, &accent_class_name(idx))
43                } else {
44                    Ok(event)
45                }
46            }
47
48            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"tspan" => {
49                if let Some(idx) = self.current_text_accent {
50                    add_to_event(&event, e, &accent_class_name(idx))
51                } else {
52                    Ok(event)
53                }
54            }
55
56            Event::End(e) if e.name().as_ref() == b"text" => {
57                self.current_text_accent = None;
58                Ok(event)
59            }
60
61            _ => Ok(event),
62        }
63    }
64
65    fn check_actor_rect(&mut self, e: &BytesStart<'_>) -> Result<Option<usize>> {
66        let class_attr = match e.try_get_attribute("class")? {
67            Some(a) => a,
68            None => return Ok(None),
69        };
70        let class_val = class_attr.unescape_value()?;
71        if class_val.contains("actor-bottom") {
72            let idx = self.actor_bottom_counter % self.accent_count;
73            self.actor_bottom_counter += 1;
74            self.last_actor_accent = Some(idx);
75            Ok(Some(idx))
76        } else if class_val.contains("actor-top") {
77            let idx = self.actor_top_counter % self.accent_count;
78            self.actor_top_counter += 1;
79            self.last_actor_accent = Some(idx);
80            Ok(Some(idx))
81        } else {
82            Ok(None)
83        }
84    }
85
86    fn check_actor_text(&mut self, e: &BytesStart<'_>) -> Result<Option<usize>> {
87        let class_attr = match e.try_get_attribute("class")? {
88            Some(a) => a,
89            None => return Ok(None),
90        };
91        let class_val = class_attr.unescape_value()?;
92        if class_val.contains("actor") && class_val.contains("actor-box") {
93            Ok(self.last_actor_accent.take())
94        } else {
95            Ok(None)
96        }
97    }
98}
99