tenant.openagents/omega

1//! Injects CSS classes to set accent colors. Mermaid broadly speaking does not
2//! provide a mechanism to color individual nodes. A few diagram types support
3//! `:::my-css-class` on nodes, but most don't.
4//!
5//! [`inject_css`](super::inject_css) then injects CSS rules for the classes
6//! that this module injects.
7
8mod class_diagram;
9mod mindmap;
10mod sequence_diagram;
11
12use anyhow::Result;
13use quick_xml::events::{BytesStart, Event};
14
15use crate::MermaidTheme;
16
17pub(crate) struct NodeRect {
18    pub cx: f64,
19    pub cy: f64,
20    pub half_height: f64,
21    pub accent_idx: usize,
22}
23
24#[derive(Default)]
25pub(crate) struct NodeTracker {
26    rects: Vec<NodeRect>,
27    building: Option<NodeRect>,
28}
29
30impl NodeTracker {
31    pub fn start_node(&mut self, cx: f64, cy: f64, half_height: f64, accent_idx: usize) {
32        self.building = Some(NodeRect {
33            cx,
34            cy,
35            half_height,
36            accent_idx,
37        });
38    }
39
40    pub fn maybe_finish_node(&mut self) {
41        if let Some(rect) = self.building.take() {
42            self.rects.push(rect);
43        }
44    }
45
46    pub fn update_half_height(&mut self, e: &BytesStart<'_>) {
47        if let Some(node) = &mut self.building
48            && let Some(hh) = parse_path_half_height(e)
49            && hh > node.half_height
50        {
51            node.half_height = hh;
52        }
53    }
54
55    pub fn lookup_accent(&self, e: &BytesStart<'_>) -> Option<usize> {
56        lookup_position_accent(&self.rects, e)
57    }
58}
59
60pub(crate) fn parse_translate(e: &BytesStart<'_>) -> Option<(f64, f64)> {
61    let attr = e.try_get_attribute("transform").ok()??;
62    let val = attr.unescape_value().ok()?;
63    let inner = val.strip_prefix("translate(")?.strip_suffix(')')?;
64    let (x_str, y_str) = inner.split_once(',')?;
65    Some((x_str.trim().parse().ok()?, y_str.trim().parse().ok()?))
66}
67
68pub(crate) fn parse_path_half_height(e: &BytesStart<'_>) -> Option<f64> {
69    let attr = e.try_get_attribute("d").ok()??;
70    let d = attr.unescape_value().ok()?;
71    let rest = d.strip_prefix('M')?.trim_start();
72    // The path data starts with `M x,y ...`; the y coordinate is the second
73    // whitespace/comma-separated token.
74    let y: f64 = rest
75        .split([' ', ','])
76        .filter(|token| !token.is_empty())
77        .nth(1)?
78        .parse()
79        .ok()?;
80    Some(y.abs())
81}
82
83// These arrays are basically just optimized versions of `format!("zed-accent-{i}")`
84const ACCENT_CLASSES: [&str; 8] = [
85    "zed-accent-0",
86    "zed-accent-1",
87    "zed-accent-2",
88    "zed-accent-3",
89    "zed-accent-4",
90    "zed-accent-5",
91    "zed-accent-6",
92    "zed-accent-7",
93];
94
95const CHART_COLOR_CLASSES: [&str; 8] = [
96    "zed-chart-0",
97    "zed-chart-1",
98    "zed-chart-2",
99    "zed-chart-3",
100    "zed-chart-4",
101    "zed-chart-5",
102    "zed-chart-6",
103    "zed-chart-7",
104];
105
106pub(crate) fn accent_class_name(index: usize) -> &'static str {
107    ACCENT_CLASSES[index % ACCENT_CLASSES.len()]
108}
109
110fn chart_color_class_name(index: usize) -> &'static str {
111    CHART_COLOR_CLASSES[index % CHART_COLOR_CLASSES.len()]
112}
113
114/// Wraps [`add_class`] and preserves the `Start`/`Empty` variant of the original event.
115pub(crate) fn add_to_event<'a>(ev: &Event<'_>, e: &BytesStart<'_>, cl: &str) -> Result<Event<'a>> {
116    let new_elem = add_class(e, cl)?;
117    Ok(match ev {
118        Event::Start(_) => Event::Start(new_elem),
119        _ => Event::Empty(new_elem),
120    })
121}
122
123/// Adds a CSS class to an element, preserving any existing classes.
124pub(crate) fn add_class<'a>(e: &BytesStart<'_>, class_to_add: &str) -> Result<BytesStart<'a>> {
125    let name = e.name();
126    let tag = std::str::from_utf8(name.as_ref())?;
127    let mut new_elem = BytesStart::new(tag.to_owned());
128    let mut class_found = false;
129    for attr in e.attributes() {
130        let attr = attr?;
131        if attr.key.local_name().as_ref() == b"class" {
132            let existing = attr.unescape_value()?;
133            let new_class = format!("{existing} {class_to_add}");
134            new_elem.push_attribute(("class", new_class.as_str()));
135            class_found = true;
136        } else {
137            new_elem.push_attribute(attr);
138        }
139    }
140    if !class_found {
141        new_elem.push_attribute(("class", class_to_add));
142    }
143    Ok(new_elem)
144}
145
146#[derive(Debug, Clone, Copy)]
147pub(crate) struct AccentStackEntry {
148    /// The accent inherited by nested text, even when the group has no layout geometry.
149    accent_idx: Option<usize>,
150    /// Whether this group called `NodeTracker::start_node` and must be finished later.
151    tracks_node: bool,
152}
153
154impl AccentStackEntry {
155    pub fn none() -> Self {
156        Self {
157            accent_idx: None,
158            tracks_node: false,
159        }
160    }
161
162    pub fn accent(accent_idx: usize, tracks_node: bool) -> Self {
163        Self {
164            accent_idx: Some(accent_idx),
165            tracks_node,
166        }
167    }
168
169    pub fn tracks_node(self) -> bool {
170        self.tracks_node
171    }
172
173    fn accent_idx(self) -> Option<usize> {
174        self.accent_idx
175    }
176}
177
178pub(crate) fn current_stack_accent(stack: &[AccentStackEntry]) -> Option<usize> {
179    stack.iter().rev().find_map(|entry| entry.accent_idx())
180}
181
182// merman's fallback overlay groups intentionally preserve source classes such
183// as `node` and `section-*` so host CSS can style fallback text. They are not
184// layout nodes though, so accent tracking must not count them as new nodes.
185pub(crate) fn is_foreign_object_fallback_group(e: &BytesStart<'_>) -> Result<bool> {
186    Ok(e.try_get_attribute("data-merman-foreignobject")?
187        .is_some_and(|attr| attr.value.as_ref() == b"fallback"))
188}
189
190pub(crate) fn lookup_position_accent(node_rects: &[NodeRect], e: &BytesStart<'_>) -> Option<usize> {
191    let parse_attr = |name| -> Option<f64> {
192        e.try_get_attribute(name)
193            .ok()??
194            .unescape_value()
195            .ok()?
196            .parse()
197            .ok()
198    };
199    let x: f64 = parse_attr("x")?;
200    let y: f64 = parse_attr("y")?;
201
202    node_rects.iter().find_map(|rect| {
203        let in_y = (y - rect.cy).abs() <= rect.half_height + 5.0;
204        let in_x = (x - rect.cx).abs() <= rect.half_height * 2.0;
205        (in_x && in_y).then_some(rect.accent_idx)
206    })
207}
208
209#[derive(Debug, Clone, Copy, PartialEq, Eq)]
210enum DiagramType {
211    Flowchart,
212    Mindmap,
213    ClassDiagram,
214    StateDiagram,
215    SequenceDiagram,
216    Unhandled,
217}
218
219fn detect_diagram_type(e: &BytesStart<'_>) -> DiagramType {
220    let class = match e
221        .try_get_attribute("class")
222        .ok()
223        .flatten()
224        .and_then(|a| a.unescape_value().ok())
225    {
226        Some(c) => c,
227        None => return DiagramType::SequenceDiagram,
228    };
229
230    for token in class.split_whitespace() {
231        match token {
232            "flowchart" => return DiagramType::Flowchart,
233            "mindmap" => return DiagramType::Mindmap,
234            "classDiagram" => return DiagramType::ClassDiagram,
235            "statediagram" => return DiagramType::StateDiagram,
236            "journey" => return DiagramType::Unhandled,
237            _ => {}
238        }
239    }
240
241    DiagramType::SequenceDiagram
242}
243
244/// Different diagrams require different state when computing accent colors.
245enum Handler {
246    /// Before we have identified the diagram type
247    Pending,
248    /// Diagram type doesn't require injecting classes.
249    Passthrough,
250    Flowchart(class_diagram::ClassDiagramAccents),
251    Mindmap(mindmap::MindmapAccents),
252    ClassDiagram(class_diagram::ClassDiagramAccents),
253    StateDiagram(class_diagram::ClassDiagramAccents),
254    Sequence(sequence_diagram::SequenceDiagramAccents),
255}
256
257struct AccentColors<'theme, I> {
258    inner: I,
259    theme: &'theme MermaidTheme,
260    handler: Handler,
261    in_legend: bool,
262    legend_color_idx: usize,
263    in_plot: bool,
264    plot_depth: usize,
265    plot_path_done: bool,
266    pie_color_idx: usize,
267    quadrant_point_idx: usize,
268}
269
270impl<'a, 'theme, I: Iterator<Item = Result<Event<'a>>>> AccentColors<'theme, I> {
271    fn process_chart_colors(&mut self, event: Event<'a>) -> Result<Event<'a>> {
272        match &event {
273            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"g" => {
274                let is_start = matches!(event, Event::Start(_));
275                // Only a real opening tag increases nesting depth. Self-closing `<g/>`
276                // elements have no matching `</g>`, so counting them would leave
277                // `plot_depth` permanently inflated and `in_plot` stuck on.
278                if self.in_plot && is_start {
279                    self.plot_depth += 1;
280                }
281                if let Some(class_attr) = e.try_get_attribute("class")? {
282                    let class = class_attr.unescape_value()?;
283                    if class.as_ref() == "plot" {
284                        if is_start && !self.in_plot {
285                            self.in_plot = true;
286                            self.plot_depth = 1;
287                            self.plot_path_done = false;
288                        }
289                    } else if class.as_ref() == "legend" {
290                        self.in_legend = true;
291                    } else if class.as_ref() == "data-point" {
292                        let accent_count = self.theme.accent_colors.len();
293                        if accent_count > 0 {
294                            let idx = self.quadrant_point_idx % accent_count;
295                            self.quadrant_point_idx += 1;
296                            return add_to_event(&event, e, &accent_class_name(idx));
297                        }
298                    }
299                }
300                Ok(event)
301            }
302
303            Event::End(e) if e.name().as_ref() == b"g" => {
304                if self.in_plot {
305                    self.plot_depth = self.plot_depth.saturating_sub(1);
306                    if self.plot_depth == 0 {
307                        self.in_plot = false;
308                    }
309                }
310                Ok(event)
311            }
312
313            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"rect" => {
314                if self.in_legend && self.legend_color_idx < 8 {
315                    let class = chart_color_class_name(self.legend_color_idx);
316                    self.legend_color_idx += 1;
317                    self.in_legend = false;
318                    add_to_event(&event, e, &class)
319                } else if self.in_plot {
320                    add_to_event(&event, e, &chart_color_class_name(0))
321                } else {
322                    Ok(event)
323                }
324            }
325
326            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"path" => {
327                let class_val = e
328                    .try_get_attribute("class")?
329                    .map(|a| a.unescape_value())
330                    .transpose()?;
331
332                if class_val.as_deref() == Some("pieCircle") {
333                    let class = chart_color_class_name(self.pie_color_idx % 8);
334                    self.pie_color_idx += 1;
335                    add_to_event(&event, e, &class)
336                } else if self.in_plot
337                    && !self.plot_path_done
338                    && e.try_get_attribute("stroke")?.is_some()
339                {
340                    self.plot_path_done = true;
341                    add_to_event(&event, e, &chart_color_class_name(1))
342                } else {
343                    Ok(event)
344                }
345            }
346
347            _ => Ok(event),
348        }
349    }
350}
351
352impl<'a, 'theme, I: Iterator<Item = Result<Event<'a>>>> Iterator for AccentColors<'theme, I> {
353    type Item = Result<Event<'a>>;
354
355    fn next(&mut self) -> Option<Self::Item> {
356        let event = match self.inner.next()? {
357            Ok(ev) => ev,
358            Err(e) => return Some(Err(e)),
359        };
360
361        if matches!(self.handler, Handler::Pending) {
362            if let Event::Start(e) | Event::Empty(e) = &event {
363                if e.name().as_ref() == b"svg" {
364                    let diagram_type = detect_diagram_type(e);
365                    let count = self.theme.accent_colors.len();
366                    self.handler = match diagram_type {
367                        DiagramType::Flowchart => {
368                            Handler::Flowchart(class_diagram::ClassDiagramAccents::new(count))
369                        }
370                        DiagramType::Mindmap => Handler::Mindmap(mindmap::MindmapAccents::new()),
371                        DiagramType::ClassDiagram => {
372                            Handler::ClassDiagram(class_diagram::ClassDiagramAccents::new(count))
373                        }
374                        DiagramType::StateDiagram => {
375                            Handler::StateDiagram(class_diagram::ClassDiagramAccents::new(count))
376                        }
377                        DiagramType::SequenceDiagram => {
378                            Handler::Sequence(sequence_diagram::SequenceDiagramAccents::new(count))
379                        }
380                        DiagramType::Unhandled => Handler::Passthrough,
381                    };
382                }
383            }
384        }
385
386        let event = match &mut self.handler {
387            Handler::Flowchart(h) | Handler::ClassDiagram(h) | Handler::StateDiagram(h) => {
388                h.process_event(event)
389            }
390            Handler::Mindmap(h) => h.process_event(event),
391            Handler::Sequence(h) => h.process_event(event),
392            Handler::Passthrough | Handler::Pending => Ok(event),
393        };
394
395        Some(match event {
396            Ok(event) => self.process_chart_colors(event),
397            err => err,
398        })
399    }
400}
401
402pub(super) fn process<'a, 'theme>(
403    events: impl Iterator<Item = Result<Event<'a>>>,
404    theme: &'theme MermaidTheme,
405) -> impl Iterator<Item = Result<Event<'a>>> {
406    AccentColors {
407        inner: events,
408        theme,
409        handler: Handler::Pending,
410        in_legend: false,
411        legend_color_idx: 0,
412        in_plot: false,
413        plot_depth: 0,
414        plot_path_done: false,
415        pie_color_idx: 0,
416        quadrant_point_idx: 0,
417    }
418}
419