ai-app/app/src/client/ansi.rs

use std::ops::Range;

/// An RGB colour, the same shape wherever this crate names one -- no alpha,
/// because the one place that needs partial transparency (dimming) says so
/// with a separate flag rather than baking it into the colour.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Rgb {
    pub r: u8,
    pub g: u8,
    pub b: u8,
}

impl Rgb {
    pub const fn new(r: u8, g: u8, b: u8) -> Self {
        Self { r, g, b }
    }
}

#[derive(Debug, Clone)]
pub struct AnsiPalette {
    pub colours: [Rgb; 16],
    pub foreground: Rgb,
    pub background: Rgb,
}

/// One span's worth of styling. `None` means unspecified.
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub struct Style {
    pub color: Option<Rgb>,
    /// How much of `color`'s alpha survives, 0.0-1.0; `None` is opaque.
    pub alpha: Option<f32>,
    pub background: Option<Rgb>,
    pub bold: bool,
    pub italic: bool,
    pub underline: bool,
    pub strikethrough: bool,
}

#[derive(Debug, Clone, PartialEq, Default)]
pub struct StyledText {
    pub text: String,
    pub spans: Vec<(Range<usize>, Style)>,
}

impl StyledText {
    fn plain(text: String) -> Self {
        Self {
            text,
            spans: Vec::new(),
        }
    }
}

const ESC: char = '\u{1B}';
const BELL: char = '\u{7}';

pub fn ansi_styled(text: &str, palette: &AnsiPalette) -> StyledText {
    if !text.contains(ESC) && !text.contains('\r') {
        return StyledText::plain(text.to_string());
    }

    let chars: Vec<char> = text.chars().collect();
    let mut runs: Vec<(String, Option<Style>)> = Vec::new();
    let mut sgr = Sgr::PLAIN;
    let mut at = 0usize;
    let mut plain = String::new();

    let flush = |plain: &mut String, sgr: Sgr, runs: &mut Vec<(String, Option<Style>)>| {
        if !plain.is_empty() {
            runs.push((std::mem::take(plain), sgr.span(palette)));
        }
    };

    while at < chars.len() {
        let c = chars[at];
        if c == ESC {
            flush(&mut plain, sgr, &mut runs);
            at = skip_escape(&chars, at, |params, final_byte| {
                if final_byte == 'm' {
                    sgr = sgr.apply(params, palette);
                }
            });
        } else if c == '\r' && chars.get(at + 1) != Some(&'\n') {
            flush(&mut plain, sgr, &mut runs);
            drop_line(&mut runs);
            at += 1;
        } else if c == '\r' {
            at += 1;
        } else if c >= ' ' || c == '\n' || c == '\t' {
            plain.push(c);
            at += 1;
        } else {
            at += 1;
        }
    }
    flush(&mut plain, sgr, &mut runs);

    let mut out = String::new();
    let mut spans = Vec::new();
    for (run_text, style) in runs {
        let start = out.len();
        out.push_str(&run_text);
        if let Some(style) = style {
            spans.push((start..out.len(), style));
        }
    }
    StyledText { text: out, spans }
}

fn drop_line(runs: &mut Vec<(String, Option<Style>)>) {
    while let Some((text, style)) = runs.pop() {
        if let Some(break_at) = text.rfind('\n') {
            runs.push((text[..=break_at].to_string(), style));
            return;
        }
    }
}

fn is_csi_final(c: char) -> bool {
    ('@'..='~').contains(&c)
}

/// Steps over the escape sequence starting at `at`, reporting a CSI's
/// parameters and final byte. One reader for every kind, because the point
/// is to *leave* them all behind: a sequence this did not recognise would
/// otherwise have its body printed as ordinary text. Three shapes -- the CSI
/// (`ESC [ ... letter`), the string escapes which run to a terminator, and
/// the two-character ones.
fn skip_escape(chars: &[char], at: usize, mut on_csi: impl FnMut(&str, char)) -> usize {
    let Some(&next) = chars.get(at + 1) else {
        return at + 1;
    };
    match next {
        '[' => {
            let mut end = at + 2;
            while end < chars.len() && !is_csi_final(chars[end]) {
                end += 1;
            }
            if end >= chars.len() {
                chars.len()
            } else {
                let params: String = chars[at + 2..end].iter().collect();
                on_csi(&params, chars[end]);
                end + 1
            }
        }
        ']' | 'P' | 'X' | '^' | '_' => {
            let mut end = at + 2;
            while end < chars.len() {
                if chars[end] == BELL {
                    return end + 1;
                }
                if chars[end] == ESC && chars.get(end + 1) == Some(&'\\') {
                    return end + 2;
                }
                end += 1;
            }
            chars.len()
        }
        _ => at + 2,
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct Sgr {
    fg: Option<Rgb>,
    bg: Option<Rgb>,
    bold: bool,
    dim: bool,
    italic: bool,
    underline: bool,
    strike: bool,
    reverse: bool,
}

const DIM_ALPHA: f32 = 0.65;

impl Sgr {
    const PLAIN: Sgr = Sgr {
        fg: None,
        bg: None,
        bold: false,
        dim: false,
        italic: false,
        underline: false,
        strike: false,
        reverse: false,
    };

    fn span(&self, palette: &AnsiPalette) -> Option<Style> {
        if *self == Sgr::PLAIN {
            return None;
        }
        let front = if self.reverse {
            Some(self.bg.unwrap_or(palette.background))
        } else {
            self.fg
        };
        let back = if self.reverse {
            Some(self.fg.unwrap_or(palette.foreground))
        } else {
            self.bg
        };
        // Dim has to have a colour to dim, so where none was named it dims
        // the ordinary one.
        let stated = front.or(if self.dim {
            Some(palette.foreground)
        } else {
            None
        });
        Some(Style {
            color: stated,
            alpha: if self.dim { Some(DIM_ALPHA) } else { None },
            background: back,
            bold: self.bold,
            italic: self.italic,
            underline: self.underline,
            strikethrough: self.strike,
        })
    }

    fn apply(&self, params: &str, palette: &AnsiPalette) -> Sgr {
        let codes: Vec<i64> = params
            .split(';')
            .map(|p| p.trim().parse::<i64>().unwrap_or(0))
            .collect();
        let mut state = *self;
        let mut at = 0usize;
        while at < codes.len() {
            let code = codes[at];
            state = match code {
                0 => Sgr::PLAIN,
                1 => Sgr {
                    bold: true,
                    ..state
                },
                2 => Sgr { dim: true, ..state },
                3 => Sgr {
                    italic: true,
                    ..state
                },
                4 => Sgr {
                    underline: true,
                    ..state
                },
                7 => Sgr {
                    reverse: true,
                    ..state
                },
                9 => Sgr {
                    strike: true,
                    ..state
                },
                21 | 22 => Sgr {
                    bold: false,
                    dim: false,
                    ..state
                },
                23 => Sgr {
                    italic: false,
                    ..state
                },
                24 => Sgr {
                    underline: false,
                    ..state
                },
                27 => Sgr {
                    reverse: false,
                    ..state
                },
                29 => Sgr {
                    strike: false,
                    ..state
                },
                30..=37 => Sgr {
                    fg: Some(palette.colours[(code - 30) as usize]),
                    ..state
                },
                90..=97 => Sgr {
                    fg: Some(palette.colours[(code - 90 + 8) as usize]),
                    ..state
                },
                40..=47 => Sgr {
                    bg: Some(palette.colours[(code - 40) as usize]),
                    ..state
                },
                100..=107 => Sgr {
                    bg: Some(palette.colours[(code - 100 + 8) as usize]),
                    ..state
                },
                39 => Sgr { fg: None, ..state },
                49 => Sgr { bg: None, ..state },
                38 | 48 => {
                    let (colour, last) = extended_colour(&codes, at, palette);
                    at = last;
                    if code == 38 {
                        Sgr {
                            fg: colour,
                            ..state
                        }
                    } else {
                        Sgr {
                            bg: colour,
                            ..state
                        }
                    }
                }
                _ => state,
            };
            at += 1;
        }
        state
    }
}

fn extended_colour(codes: &[i64], at: usize, palette: &AnsiPalette) -> (Option<Rgb>, usize) {
    match codes.get(at + 1) {
        Some(&5) => match codes.get(at + 2) {
            None => (None, at + 1),
            Some(&n) => (Some(indexed_colour(n, palette)), at + 2),
        },
        Some(&2) => {
            let r = codes.get(at + 2);
            let g = codes.get(at + 3);
            let b = codes.get(at + 4);
            match (r, g, b) {
                (Some(&r), Some(&g), Some(&b)) => (
                    Some(Rgb::new(
                        r.clamp(0, 255) as u8,
                        g.clamp(0, 255) as u8,
                        b.clamp(0, 255) as u8,
                    )),
                    at + 4,
                ),
                _ => (None, at + 1),
            }
        }
        _ => (None, at + 1),
    }
}

const CUBE: [u8; 6] = [0, 95, 135, 175, 215, 255];

fn indexed_colour(n: i64, palette: &AnsiPalette) -> Rgb {
    if n < 0 {
        palette.foreground
    } else if n < 16 {
        palette.colours[n as usize]
    } else if n < 232 {
        let i = (n - 16) as usize;
        Rgb::new(CUBE[i / 36], CUBE[i / 6 % 6], CUBE[i % 6])
    } else if n < 256 {
        let grey = (8 + (n - 232) * 10) as u8;
        Rgb::new(grey, grey, grey)
    } else {
        palette.foreground
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn palette() -> AnsiPalette {
        let mut colours = [Rgb::new(0, 0, 0); 16];
        for (i, c) in colours.iter_mut().enumerate() {
            *c = Rgb::new(i as u8, 0, 0);
        }
        AnsiPalette {
            colours,
            foreground: Rgb::new(255, 255, 255),
            background: Rgb::new(0, 0, 0),
        }
    }

    fn styled(text: &str) -> StyledText {
        ansi_styled(text, &palette())
    }

    fn style_over(text: &str, word: &str) -> Option<Style> {
        let out = styled(text);
        let at = out
            .text
            .find(word)
            .unwrap_or_else(|| panic!("no {word:?} in {}", out.text));
        out.spans
            .iter()
            .find(|(range, _)| range.contains(&at))
            .map(|(_, style)| *style)
    }

    #[test]
    fn a_colour_becomes_a_span_and_the_sequence_itself_disappears() {
        let text = format!("plain {ESC}[31mred{ESC}[0m plain");
        assert_eq!(styled(&text).text, "plain red plain");
        assert_eq!(
            style_over(&text, "red").unwrap().color,
            Some(Rgb::new(1, 0, 0))
        );
        assert!(style_over(&text, "plain").is_none());
    }

    #[test]
    fn bright_background_and_256_colour_forms_all_reach_the_same_table() {
        assert_eq!(
            style_over(&format!("{ESC}[91mx"), "x").unwrap().color,
            Some(Rgb::new(9, 0, 0))
        );
        assert_eq!(
            style_over(&format!("{ESC}[44mx"), "x").unwrap().background,
            Some(Rgb::new(4, 0, 0))
        );
        assert_eq!(
            style_over(&format!("{ESC}[38;5;1mx"), "x").unwrap().color,
            Some(Rgb::new(1, 0, 0))
        );
        assert_eq!(
            style_over(&format!("{ESC}[38;5;16mx"), "x").unwrap().color,
            Some(Rgb::new(0, 0, 0))
        );
        assert_eq!(
            style_over(&format!("{ESC}[38;5;231mx"), "x").unwrap().color,
            Some(Rgb::new(255, 255, 255))
        );
        assert_eq!(
            style_over(&format!("{ESC}[38;2;10;20;30mx"), "x")
                .unwrap()
                .color,
            Some(Rgb::new(10, 20, 30))
        );
    }

    #[test]
    fn everything_that_is_not_styling_is_dropped_rather_than_printed() {
        let text = format!("a{ESC}[2Jb{ESC}[Kc{ESC}]0;a title{BELL}d{ESC}=e");
        assert_eq!(styled(&text).text, "abcde");
    }

    #[test]
    fn a_carriage_return_rewrites_its_line_as_it_does_on_a_terminal() {
        assert_eq!(styled("10%\r50%\rdone\n").text, "done\n");
        assert_eq!(styled("kept\r\nfirst\rlast").text, "kept\nlast");
    }

    #[test]
    fn a_sequence_cut_off_mid_stream_takes_no_text_with_it() {
        assert_eq!(styled(&format!("text {ESC}[3")).text, "text ");
    }

    #[test]
    fn unstyled_text_costs_no_spans_at_all() {
        assert_eq!(styled("nothing to do here").spans.len(), 0);
        assert_eq!(styled(&format!("a{ESC}[2Jb")).spans.len(), 0);
    }
}