//! ANSI-aware string replacement.
//!
//! This crate provides utilities for performing find-and-replace operations on
//! strings that contain ANSI escape sequences (e.g., terminal colors and styles).
//! Replacements are matched against the visible text only, while preserving all
//! embedded ANSI sequences in their original positions.
//!
//! # Example
//!
//! ```
//! use ansi_replace::AnsiExt;
//!
//! let colored = "\x1b[31mhello world\x1b[0m";
//! let result = colored.ansi_replace("hello", "goodbye");
//! assert_eq!(result, "\x1b[31mgoodbye world\x1b[0m");
//! ```

use core::fmt::Write as _;
use core::ops::Not;

use stable_pattern::Searcher as _;

use crate::replacer::Writable;

pub mod pattern;
pub mod replacer;

#[cfg(test)]
mod tests;

/// An iterator over the chunks of a string containing ANSI escape sequences.
///
/// Each chunk is either a contiguous block of visible text or an ANSI escape
/// sequence. Use [`AnsiExt::ansi_chunks`] to create an instance.
pub struct AnsiIterator<'a> {
    text: &'a str,
    visible_offset: usize,
}

fn consume_csi_sequence(text: &str) -> usize {
    // CSI sequence: ESC [ ... <final byte (0x40-0x7E)>
    let bytes = text.as_bytes();
    let mut pos = 2; // Start after ESC [
    while pos < bytes.len() && !(0x40..=0x7E).contains(&bytes[pos]) {
        pos += 1;
    }
    if pos < bytes.len() {
        pos += 1; // consume final byte
    }
    pos
}

fn consume_osc_sequence(text: &str) -> usize {
    // OSC sequence: ESC ] ... (terminated by BEL or ST)
    let bel = text.find('\x07').map(|pos| pos + 1).unwrap_or(text.len());
    let st = text.find("\x1b\\").map(|pos| pos + 2).unwrap_or(text.len());
    bel.min(st)
}

fn consume_other_ansi_sequence(text: &str) -> usize {
    // For other ANSI sequences that start with ESC but don't match CSI or OSC
    for i in 1..text.len() {
        if text.is_char_boundary(i) {
            return i; // Next ESC character indicates the end of the current sequence
        }
    }
    text.len() // If we can't find another ESC, consume the rest of the string
}

fn consume_ansi_sequence(text: &str) -> usize {
    if text.starts_with("\x1b[") {
        return consume_csi_sequence(text);
    }
    if text.starts_with("\x1b]") {
        return consume_osc_sequence(text);
    }
    consume_other_ansi_sequence(text)
}

fn consume_visible_text(text: &str) -> usize {
    text.find('\x1b').unwrap_or(text.len())
}

fn consume_chunk(text: &str) -> usize {
    if text.starts_with('\x1b') {
        consume_ansi_sequence(text)
    } else {
        consume_visible_text(text)
    }
}

/// A single chunk from an ANSI-aware string split.
///
/// The chunk is either visible text or an ANSI escape sequence.
pub struct AnsiChunk<'a> {
    /// The raw content of this chunk (visible text or escape sequence).
    pub content: &'a str,
    /// The byte offset of this chunk within the visible (stripped) text.
    /// For ANSI escape sequences, this is the offset of the surrounding visible text.
    pub visible_offset: usize,
}

impl AnsiChunk<'_> {
    /// Returns true if this chunk is an ANSI escape sequence, false if it's visible text.
    pub fn is_ansi(&self) -> bool {
        self.content.starts_with('\x1b')
    }
}

impl<'a> Iterator for AnsiIterator<'a> {
    type Item = AnsiChunk<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.text.is_empty() {
            return None;
        }

        let n = consume_chunk(self.text);

        if n == 0 {
            return None; // Should not happen, but just in case
        }

        let (content, rest) = self.text.split_at(n);
        let visible_offset = self.visible_offset;

        self.text = rest;
        if !content.starts_with('\x1b') {
            // Visible text, update visible offset
            self.visible_offset += content.len();
        }

        Some(AnsiChunk {
            content,
            visible_offset,
        })
    }
}

/// Extension trait for ANSI-aware string operations.
///
/// This trait is implemented for [`str`] and provides methods for stripping,
/// iterating, and replacing text while preserving ANSI escape sequences.
pub trait AnsiExt {
    /// Returns an iterator over the ANSI chunks of this string.
    ///
    /// Each chunk is either a span of visible text or an ANSI escape sequence.
    fn ansi_chunks(&'_ self) -> AnsiIterator<'_>;

    /// Returns a new string with all ANSI escape sequences removed.
    fn ansi_strip(&self) -> String {
        self.ansi_chunks()
            .filter_map(|chunk| chunk.is_ansi().not().then_some(chunk.content))
            .collect()
    }

    /// Replaces at most `count` non-overlapping matches of `from` with `to`,
    /// matching against visible text only while preserving ANSI escape sequences.
    fn ansi_replacen(
        &self,
        from: impl crate::pattern::Pattern,
        mut to: impl crate::replacer::Replacer,
        count: usize,
    ) -> String {
        let stripped = self.ansi_strip();

        let mut searcher = from.into_searcher(&stripped);
        let mut searcher = core::iter::from_fn(move || searcher.next_match())
            .take(count)
            .map(|(s, e)| s..e)
            .chain(core::iter::once(usize::MAX..usize::MAX));

        let mut m = searcher.next().unwrap();
        let mut result = String::new();
        let writable = Writable::new(&mut result);

        for chunk in self.ansi_chunks() {
            if chunk.is_ansi() {
                let _ = writable.write_str(chunk.content);
                continue;
            }

            let chunk_start = chunk.visible_offset;
            let chunk_end = chunk_start + chunk.content.len();
            let mut pos = chunk_start;

            loop {
                // Emit replacement for any matches that ended at or before current pos
                while m.end <= pos {
                    to.replace(&stripped[m.start..m.end], m.start..m.end, writable);
                    m = searcher.next().unwrap();
                }

                if pos >= chunk_end {
                    break;
                }

                if pos < m.start {
                    // Copy visible text up to the match start (or chunk end)
                    let end = m.start.min(chunk_end);
                    let _ =
                        writable.write_str(&chunk.content[pos - chunk_start..end - chunk_start]);
                    pos = end;
                } else {
                    // Skip visible text inside the match
                    pos = m.end.min(chunk_end);
                }
            }
        }

        result
    }

    /// Replaces all non-overlapping matches of `from` with `to`, matching
    /// against visible text only while preserving ANSI escape sequences.
    fn ansi_replace(
        &self,
        from: impl crate::pattern::Pattern,
        to: impl crate::replacer::Replacer,
    ) -> String {
        self.ansi_replacen(from, to, usize::MAX)
    }
}

impl AnsiExt for str {
    fn ansi_chunks(&'_ self) -> AnsiIterator<'_> {
        AnsiIterator {
            text: self,
            visible_offset: 0,
        }
    }
}
