magi-code 0.96.1

use super::*;

pub(crate) struct PromptTextNormalizer {
    pub(super) normalized: String,
    max_bytes: usize,
    pending_cr: bool,
    overflowed: bool,
}

impl PromptTextNormalizer {
    pub(crate) fn new() -> Self {
        Self::with_limit(MAX_PROMPT_BYTES)
    }

    fn with_limit(max_bytes: usize) -> Self {
        Self {
            normalized: String::new(),
            max_bytes,
            pending_cr: false,
            overflowed: false,
        }
    }

    pub(crate) fn push_str(&mut self, text: &str) {
        for ch in text.chars() {
            self.push_char(ch);
        }
    }

    pub(crate) fn push_char(&mut self, ch: char) {
        if self.pending_cr {
            self.pending_cr = false;
            if ch == '\n' {
                return;
            }
        }

        if ch == '\r' {
            self.pending_cr = true;
            self.push_normalized_char('\n');
        } else if ch == '\t'
            || ch == '\n'
            || !(('\u{0}'..='\u{1f}').contains(&ch) || ch == '\u{7f}')
        {
            self.push_normalized_char(ch);
        }
    }

    fn push_normalized_char(&mut self, ch: char) {
        if self.overflowed {
            return;
        }
        let Some(new_len) = self.normalized.len().checked_add(ch.len_utf8()) else {
            self.overflowed = true;
            return;
        };
        if new_len > self.max_bytes {
            self.overflowed = true;
        } else {
            self.normalized.push(ch);
        }
    }

    pub(crate) fn finish(self) -> Result<String, PromptEditResult> {
        if self.overflowed {
            Err(PromptEditResult::RejectedTooLarge)
        } else {
            Ok(self.normalized)
        }
    }
}

pub(crate) fn normalize_prompt_text_bounded(text: &str) -> Result<String, PromptEditResult> {
    let mut normalizer = PromptTextNormalizer::new();
    normalizer.push_str(text);
    normalizer.finish()
}

pub(super) fn normalize_prompt_text(text: &str) -> String {
    let mut normalizer = PromptTextNormalizer::with_limit(usize::MAX);
    normalizer.push_str(text);
    normalizer
        .finish()
        .expect("unbounded prompt normalization cannot exceed usize::MAX")
}

pub(super) fn normalized_cursor_offset(original: &str, cursor: usize, normalized: &str) -> usize {
    if cursor >= original.len() {
        return normalized.len();
    }
    let cursor = clamp_char_boundary(original, cursor);
    normalize_prompt_text(&original[..cursor]).len()
}