tenant.openagents/omega

1//! Fixes various issues in merman's SVG output.
2//!
3//! Replaces hardcoded white backgrounds with the theme background:
4//! ```xml
5//! <!-- before --> <svg style="background-color: white">
6//! <!-- after  --> <svg style="background-color: #1e1e2e">
7//! ```
8//!
9//! Removes `<rect>` elements with missing or invalid dimensions:
10//! ```xml
11//! <!-- before --> <rect width="NaN" height="10"/>
12//! <!-- after  --> (removed)
13//! ```
14//!
15//! Replaces hardcoded text colors with the theme text color:
16//! ```xml
17//! <!-- before --> <text fill="#333">Hello</text>
18//! <!-- after  --> <text fill="#cdd6f4">Hello</text>
19//! ```
20
21use std::borrow::Cow;
22use std::fmt::Write as _;
23
24use anyhow::Result;
25use quick_xml::events::{BytesStart, Event};
26
27use crate::MermaidTheme;
28
29struct ElementFixup<I> {
30    inner: I,
31    background_css: String,
32    text_color_css: String,
33    font_family_css: String,
34    svg_seen: bool,
35    skip_rect_depth: usize,
36}
37
38fn rewrite_attr<'a>(
39    e: &BytesStart<'_>,
40    attr_name: &[u8],
41    new_value: &str,
42) -> Result<BytesStart<'a>> {
43    let name = e.name();
44    let tag = std::str::from_utf8(name.as_ref())?;
45    let mut new_elem = BytesStart::new(tag.to_owned());
46    for attr in e.attributes() {
47        let attr = attr?;
48        if attr.key.local_name().as_ref() == attr_name {
49            let local_name = attr.key.local_name();
50            let key = std::str::from_utf8(local_name.as_ref())?;
51            new_elem.push_attribute((key, new_value));
52        } else {
53            new_elem.push_attribute(attr);
54        }
55    }
56    Ok(new_elem)
57}
58
59fn rewrap<'a>(event: &Event<'_>, elem: BytesStart<'a>) -> Event<'a> {
60    match event {
61        Event::Start(_) => Event::Start(elem),
62        _ => Event::Empty(elem),
63    }
64}
65
66fn is_bad_rect(e: &BytesStart) -> Result<bool> {
67    for attr_name in ["width", "height"] {
68        match e.try_get_attribute(attr_name)? {
69            None => return Ok(true),
70            Some(attr) => {
71                let val = attr.unescape_value()?;
72                let trimmed = val.trim();
73                if trimmed.is_empty() {
74                    return Ok(true);
75                }
76                if let Ok(n) = trimmed.parse::<f64>() {
77                    if !n.is_finite() || n <= 0.0 {
78                        return Ok(true);
79                    }
80                }
81            }
82        }
83    }
84    Ok(false)
85}
86
87fn is_hardcoded_text_fill(val: &str) -> bool {
88    matches!(
89        val,
90        "" | "#333" | "black" | "#000" | "#000000" | "white" | "#fff" | "#ffffff"
91    )
92}
93
94fn push_font_style(style: &mut String, font_family: &str) {
95    write!(style, "font-family: {font_family};").expect("write to String cannot fail");
96}
97
98fn font_style(font_family: &str) -> String {
99    let mut style = String::with_capacity(font_family.len() + "font-family: ;".len());
100    push_font_style(&mut style, font_family);
101    style
102}
103
104fn rewrite_background_style<'a>(style: &'a str, background_css: &str) -> Cow<'a, str> {
105    const PROPERTY: &str = "background-color:";
106
107    let Some(property_start) = style.find(PROPERTY) else {
108        return Cow::Borrowed(style);
109    };
110    let value_start = property_start + PROPERTY.len();
111    let value_end = style[value_start..]
112        .find(';')
113        .map_or(style.len(), |offset| value_start + offset);
114    let value = style[value_start..value_end].trim();
115
116    let is_white = value.eq_ignore_ascii_case("white")
117        || value.eq_ignore_ascii_case("#fff")
118        || value.eq_ignore_ascii_case("#ffffff");
119    if !is_white {
120        return Cow::Borrowed(style);
121    }
122
123    let value_len = value_end.saturating_sub(value_start);
124    let mut rewritten = String::with_capacity(
125        style
126            .len()
127            .saturating_sub(value_len)
128            .saturating_add(background_css.len()),
129    );
130    rewritten.push_str(&style[..value_start]);
131    rewritten.push_str(background_css);
132    rewritten.push_str(&style[value_end..]);
133    Cow::Owned(rewritten)
134}
135
136fn font_family_declaration_value(declaration: &str) -> Option<&str> {
137    let (property, value) = declaration.split_once(':')?;
138    property
139        .trim()
140        .eq_ignore_ascii_case("font-family")
141        .then(|| value.trim())
142}
143
144fn rewrite_font_style<'a>(style: &'a str, font_family: &str) -> Cow<'a, str> {
145    let mut font_family_declaration_count = 0;
146    let mut has_target_font_family = false;
147    for declaration in style
148        .split(';')
149        .map(str::trim)
150        .filter(|declaration| !declaration.is_empty())
151    {
152        if let Some(value) = font_family_declaration_value(declaration) {
153            font_family_declaration_count += 1;
154            has_target_font_family = value == font_family;
155        }
156    }
157
158    if font_family_declaration_count == 1 && has_target_font_family {
159        return Cow::Borrowed(style);
160    }
161
162    let mut rewritten =
163        String::with_capacity(style.len() + font_family.len() + " font-family: ;".len());
164    for declaration in style.split(';') {
165        let declaration = declaration.trim();
166        if declaration.is_empty() || font_family_declaration_value(declaration).is_some() {
167            continue;
168        }
169        if !rewritten.is_empty() {
170            rewritten.push(' ');
171        }
172        rewritten.push_str(declaration);
173        rewritten.push(';');
174    }
175    if !rewritten.is_empty() {
176        rewritten.push(' ');
177    }
178    push_font_style(&mut rewritten, font_family);
179    Cow::Owned(rewritten)
180}
181
182impl<'a, I: Iterator<Item = Result<Event<'a>>>> ElementFixup<I> {
183    fn rewrite_svg_style(&self, e: &BytesStart<'_>) -> Result<Option<BytesStart<'a>>> {
184        let Some(style) = e
185            .try_get_attribute("style")?
186            .map(|a| a.unescape_value())
187            .transpose()?
188        else {
189            return Ok(None);
190        };
191        let new_style = rewrite_background_style(&style, &self.background_css);
192        if matches!(new_style, Cow::Borrowed(_)) {
193            return Ok(None);
194        }
195
196        Ok(Some(rewrite_attr(e, b"style", &new_style)?))
197    }
198
199    fn rewrite_text_element(&self, e: &BytesStart<'_>, fix_fill: bool) -> Result<BytesStart<'a>> {
200        let name = e.name();
201        let tag = std::str::from_utf8(name.as_ref())?;
202        let mut new_elem = BytesStart::new(tag.to_owned());
203        let mut has_font_family = false;
204        let mut has_style = false;
205
206        for attr in e.attributes() {
207            let attr = attr?;
208            match attr.key.local_name().as_ref() {
209                b"fill" if fix_fill => {
210                    let val = attr.unescape_value()?;
211                    if is_hardcoded_text_fill(&val) {
212                        new_elem.push_attribute(("fill", self.text_color_css.as_str()));
213                    } else {
214                        new_elem.push_attribute(attr);
215                    }
216                }
217                b"font-family" => {
218                    has_font_family = true;
219                    new_elem.push_attribute(("font-family", self.font_family_css.as_str()));
220                }
221                b"style" => {
222                    has_style = true;
223                    let style = attr.unescape_value()?;
224                    let style = rewrite_font_style(&style, &self.font_family_css);
225                    new_elem.push_attribute(("style", style.as_ref()));
226                }
227                _ => new_elem.push_attribute(attr),
228            }
229        }
230
231        if !has_font_family {
232            new_elem.push_attribute(("font-family", self.font_family_css.as_str()));
233        }
234        if !has_style {
235            let style = font_style(&self.font_family_css);
236            new_elem.push_attribute(("style", style.as_str()));
237        }
238
239        Ok(new_elem)
240    }
241
242    fn process_event(&mut self, event: Event<'a>) -> Result<Option<Event<'a>>> {
243        match &event {
244            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"svg" && !self.svg_seen => {
245                self.svg_seen = true;
246                if let Some(new_elem) = self.rewrite_svg_style(e)? {
247                    Ok(Some(rewrap(&event, new_elem)))
248                } else {
249                    Ok(Some(event))
250                }
251            }
252
253            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"rect" => {
254                if is_bad_rect(e)? {
255                    if matches!(event, Event::Start(_)) {
256                        self.skip_rect_depth = 1;
257                    }
258                    Ok(None)
259                } else {
260                    Ok(Some(event))
261                }
262            }
263
264            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"text" => {
265                Ok(Some(rewrap(&event, self.rewrite_text_element(e, true)?)))
266            }
267
268            Event::Start(e) | Event::Empty(e) if e.name().as_ref() == b"tspan" => {
269                Ok(Some(rewrap(&event, self.rewrite_text_element(e, false)?)))
270            }
271
272            _ => Ok(Some(event)),
273        }
274    }
275}
276
277impl<'a, I: Iterator<Item = Result<Event<'a>>>> Iterator for ElementFixup<I> {
278    type Item = Result<Event<'a>>;
279
280    fn next(&mut self) -> Option<Self::Item> {
281        loop {
282            let event = match self.inner.next()? {
283                Ok(ev) => ev,
284                Err(e) => return Some(Err(e)),
285            };
286
287            if self.skip_rect_depth > 0 {
288                match &event {
289                    Event::Start(_) => self.skip_rect_depth += 1,
290                    Event::End(_) => self.skip_rect_depth -= 1,
291                    _ => {}
292                }
293                continue;
294            }
295
296            match self.process_event(event) {
297                Ok(Some(ev)) => return Some(Ok(ev)),
298                Ok(None) => continue,
299                Err(e) => return Some(Err(e)),
300            }
301        }
302    }
303}
304
305pub(super) fn process<'a>(
306    events: impl Iterator<Item = Result<Event<'a>>>,
307    theme: &MermaidTheme,
308) -> impl Iterator<Item = Result<Event<'a>>> {
309    ElementFixup {
310        inner: events,
311        background_css: crate::css_color(theme.background),
312        text_color_css: crate::css_color(theme.text_color),
313        font_family_css: theme.font_family.clone(),
314        svg_seen: false,
315        skip_rect_depth: 0,
316    }
317}
318