vamp/src/terminal_colors.rs

use std::{
    collections::VecDeque,
    io::{self, Write},
    time::{Duration, Instant},
};

use crossterm::event::{Event, KeyCode, KeyEventKind, KeyEventState, KeyModifiers};

const QUERY_TIMEOUT: Duration = Duration::from_millis(250);
// TODO: Remove periodic polling once a released Herdr fix reliably delivers palette-ready
// DEC 2031 notifications, including changes with unchanged light/dark preference.
// Verify notification-triggered queries return fresh colors before removing this fallback.
// Upstream: https://github.com/herdrdev/herdr/issues/4183
const REFRESH_INTERVAL: Duration = Duration::from_secs(1);
const ESCAPE_TIMEOUT: Duration = Duration::from_millis(25);
const HEADERS: [&[u8]; 3] = [b"11;", b"4;1;", b"4;2;"];
const MAX_BODY: usize = 32;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct ThemeColors {
    pub background: [u8; 3],
    pub red: [u8; 3],
    pub green: [u8; 3],
}

struct Query {
    deadline: Instant,
    colors: [Option<[u8; 3]>; 3],
}

#[derive(Default)]
enum Capture {
    #[default]
    Idle,
    Escape(Instant),
    Report(Instant),
    Terminator,
}

/// Adapts crossterm's key-shaped OSC replies without reading from stdin itself.
/// Only complete color reports are consumed; rejected candidates retain their
/// original events, including modifiers, kind, and state.
#[derive(Default)]
pub(crate) struct TerminalColors {
    pub colors: Option<ThemeColors>,
    query: Option<Query>,
    refresh_at: Option<Instant>,
    capture: Capture,
    original: Vec<Event>,
    body: [u8; MAX_BODY],
    len: usize,
    ready: VecDeque<Event>,
}

impl TerminalColors {
    pub fn refresh_due(&self, now: Instant) -> bool {
        self.query.is_none() && self.refresh_at.is_none_or(|deadline| now >= deadline)
    }

    pub fn query(&mut self, output: &mut impl Write, now: Instant) -> io::Result<()> {
        self.expire(now);
        self.refresh_at = Some(now + REFRESH_INTERVAL);
        self.query = None;
        if let Err(error) = output
            .write_all(b"\x1b]11;?\x1b\\\x1b]4;1;?;2;?\x1b\\")
            .and_then(|()| output.flush())
        {
            self.colors = None;
            return Err(error);
        }
        self.query = Some(Query {
            deadline: now + QUERY_TIMEOUT,
            colors: [None; 3],
        });
        Ok(())
    }

    pub fn push(&mut self, event: Event, now: Instant) {
        self.expire(now);
        self.accept(event, now);
    }

    pub fn expire(&mut self, now: Instant) {
        if self
            .query
            .as_ref()
            .is_some_and(|query| now >= query.deadline)
        {
            self.query = None;
            self.colors = None;
        }
        let ambiguous_expired = match self.capture {
            Capture::Escape(deadline) => now >= deadline,
            Capture::Report(deadline) => {
                now >= deadline
                    && !HEADERS
                        .iter()
                        .any(|header| self.body[..self.len].starts_with(header))
            }
            _ => false,
        };
        if ambiguous_expired {
            self.replay();
        }
        // Retain a bounded, recognized OSC candidate across query expiry: a
        // late remainder must still be swallowed, not inserted into the editor.
        // Neither late replies nor partial triples can update the palette.
    }

    pub fn pop(&mut self) -> Option<Event> {
        self.ready.pop_front()
    }

    fn accept(&mut self, event: Event, now: Instant) {
        let token = token(&event);
        match self.capture {
            Capture::Idle => match token {
                Token::Start => {
                    self.original.push(event);
                    self.capture = Capture::Report(now + QUERY_TIMEOUT);
                }
                Token::Escape => {
                    self.original.push(event);
                    self.capture = Capture::Escape(now + ESCAPE_TIMEOUT);
                }
                _ => self.ready.push_back(event),
            },
            Capture::Escape(_) if matches!(token, Token::Byte(b']')) => {
                self.original.push(event);
                self.capture = Capture::Report(now + QUERY_TIMEOUT);
            }
            Capture::Report(_) if matches!(token, Token::Escape) => {
                self.original.push(event);
                self.capture = Capture::Terminator;
            }
            Capture::Report(_) if matches!(token, Token::End) => {
                self.original.push(event);
                self.finish();
            }
            Capture::Terminator if matches!(token, Token::Byte(b'\\')) => {
                self.original.push(event);
                self.finish();
            }
            Capture::Report(_) => {
                if let Token::Byte(byte) = token
                    && self.len < MAX_BODY
                {
                    self.body[self.len] = byte;
                    if report_prefix(&self.body[..self.len + 1]) {
                        self.len += 1;
                        self.original.push(event);
                        return;
                    }
                }
                self.replay();
                self.accept(event, now);
            }
            _ => {
                self.replay();
                self.accept(event, now);
            }
        }
    }

    fn finish(&mut self) {
        let Some((slot, color)) = parse_report(&self.body[..self.len]) else {
            self.replay();
            return;
        };
        self.original.clear();
        self.len = 0;
        self.capture = Capture::Idle;
        let Some(query) = self.query.as_mut() else {
            return;
        };
        query.colors[slot] = Some(color);
        if let [Some(background), Some(red), Some(green)] = query.colors {
            self.colors = Some(ThemeColors {
                background,
                red,
                green,
            });
            self.query = None;
        }
    }

    fn replay(&mut self) {
        self.ready.extend(self.original.drain(..));
        self.len = 0;
        self.capture = Capture::Idle;
    }
}

#[derive(Clone, Copy)]
enum Token {
    Start,
    End,
    Escape,
    Byte(u8),
    Other,
}

fn token(event: &Event) -> Token {
    let Event::Key(key) = event else {
        return Token::Other;
    };
    if key.kind != KeyEventKind::Press || key.state != KeyEventState::NONE {
        return Token::Other;
    }
    match (key.code, key.modifiers) {
        (KeyCode::Char(']'), KeyModifiers::ALT) => Token::Start,
        (KeyCode::Char('\\'), KeyModifiers::ALT) | (KeyCode::Char('g'), KeyModifiers::CONTROL) => {
            Token::End
        }
        (KeyCode::Esc, KeyModifiers::NONE) => Token::Escape,
        (KeyCode::Char(ch), modifiers)
            if ch.is_ascii()
                && (modifiers == KeyModifiers::NONE
                    || (ch.is_ascii_uppercase() && modifiers == KeyModifiers::SHIFT)) =>
        {
            Token::Byte(ch as u8)
        }
        _ => Token::Other,
    }
}

fn report_prefix(body: &[u8]) -> bool {
    HEADERS.iter().any(|header| {
        header.starts_with(body) || body.strip_prefix(*header).is_some_and(color_prefix)
    })
}

fn color_prefix(color: &[u8]) -> bool {
    if b"rgb:".starts_with(color) {
        return true;
    }
    if let Some(hex) = color.strip_prefix(b"#") {
        return hex.len() <= 6 && hex.iter().all(u8::is_ascii_hexdigit);
    }
    let Some(rgb) = color.strip_prefix(b"rgb:") else {
        return false;
    };
    let mut components = rgb.split(|&byte| byte == b'/');
    for index in 0..3 {
        let Some(component) = components.next() else {
            return true;
        };
        if component.len() > 4 || !component.iter().all(u8::is_ascii_hexdigit) {
            return false;
        }
        if component.is_empty() && (index < 2 && components.clone().next().is_some()) {
            return false;
        }
    }
    components.next().is_none()
}

fn parse_report(body: &[u8]) -> Option<(usize, [u8; 3])> {
    for (slot, header) in HEADERS.iter().enumerate() {
        if let Some(color) = body.strip_prefix(*header) {
            return parse_color(color).map(|color| (slot, color));
        }
    }
    None
}

fn parse_color(color: &[u8]) -> Option<[u8; 3]> {
    let mut result = [0; 3];
    if let Some(hex) = color.strip_prefix(b"#") {
        if hex.len() != 6 {
            return None;
        }
        for (channel, component) in result.iter_mut().zip(hex.chunks_exact(2)) {
            *channel = hex_value(component)? as u8;
        }
        return Some(result);
    }
    let mut components = color.strip_prefix(b"rgb:")?.split(|&byte| byte == b'/');
    for channel in &mut result {
        let component = components.next()?;
        if component.is_empty() || component.len() > 4 {
            return None;
        }
        let value = hex_value(component)?;
        let maximum = (1_u32 << (component.len() * 4)) - 1;
        *channel = ((value * 255 + maximum / 2) / maximum) as u8;
    }
    components.next().is_none().then_some(result)
}

fn hex_value(component: &[u8]) -> Option<u32> {
    component.iter().try_fold(0, |value, &byte| {
        (byte as char).to_digit(16).map(|digit| value * 16 + digit)
    })
}

#[cfg(test)]
#[path = "../tests/unit/terminal_colors.rs"]
mod tests;