// This file is part of ICU4X. For terms of use, please see the file

// called LICENSE at the top level of the ICU4X source tree

// (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).

use crate::{impl_display_with_writeable, LengthHint, Writeable};

use core::fmt;

/// A [`Writeable`] adapter that replaces occurrences of a needle with a replacement.

///

/// This adapter performs the replacement in a streaming fashion during `write_to`,

/// requiring zero allocations.

///

/// # Examples

///

/// ```

/// use writeable::adapters::Replace;

/// use writeable::assert_writeable_eq;

/// use writeable::concat_writeable;

///

/// let source = concat_writeable!("I 💖 🦀", " and 🦀 loves me!");

/// let replace = Replace {

///     source,

///     needle: "🦀",

///     replacement: "Rust",

/// };

///

/// assert_writeable_eq!(replace, "I 💖 Rust and Rust loves me!");

/// ```

#[derive(Debug)]

#[allow(clippy::exhaustive_structs)] // designed for nesting

pub struct Replace<A, B, C> {

    /// The source writeable.

    pub source: A,

    /// The needle to search for.

    pub needle: B,

    /// The replacement writeable.

    pub replacement: C,

}

// Computes the Knuth-Morris-Pratt (KMP) prefix function (failure function) value

// for the character prefix ending at byte index `matched_bytes` in `needle`.

//

// Returns the byte length of the longest proper prefix of `needle[0..matched_bytes]`

// that is also a suffix of `needle[0..matched_bytes]`.

//

// This is computed on the fly without allocation by iterating over char boundaries.

fn get_pi_bytes(needle: &str, matched_bytes: usize) -> usize {

    let s = match needle.get(0..matched_bytes) {

        Some(s) => s,

        None => return 0,

    };

    // char_indices() gives us the byte offsets of character starts.

    // These offsets correspond to the byte lengths of all possible prefixes.

    // We want to iterate them in reverse order, excluding the first one (0)

    // because we want proper prefixes.

    for k in s

        .char_indices()

        .map(|(idx, _)| idx)

        .rev()

        .filter(|&idx| idx > 0)

    {

        // Compare the prefix of length `k` with the suffix of length `k`.

        if let Some(suffix) = s.as_bytes().get(s.len() - k..) {

            if s.as_bytes().starts_with(suffix) {

                return k;

            }

        }

    }

    0

}

// A writer wrapper that performs streaming replacement.

// It intercepts characters written to it, matches them against `needle` using KMP

// (tracking progress by storing the remaining unmatched suffix of the needle),

// and writes `replacement` when a full match is found, or the original characters otherwise.

struct ReplaceWriter<'a, W: ?Sized, C> {

    // The underlying sink to write to.

    sink: &'a mut W,

    // The needle we are searching for.

    needle: &'a str,

    // The replacement to write when the needle is matched.

    replacement: &'a C,

    // The remaining unmatched suffix of the needle.

    // This is always a suffix of `needle` starting at a character boundary.

    remaining_needle: &'a str,

}

impl<'a, W, C> ReplaceWriter<'a, W, C>

where

    W: fmt::Write + ?Sized,

    C: Writeable,

{

    fn new(sink: &'a mut W, needle: &'a str, replacement: &'a C) -> Self {

        Self {

            sink,

            needle,

            replacement,

            remaining_needle: needle,

        }

    }

    // Helper to get the length of the prefix matched so far.

    fn matched_len(&self) -> usize {

        self.needle.len() - self.remaining_needle.len()

    }

    // Finalizes the writer, flushing any partially matched prefix to the sink.

    fn finalize(&mut self) -> fmt::Result {

        let matched = self.matched_len();

        if matched > 0 {

            let slice = self.needle.get(0..matched).ok_or(fmt::Error)?;

            self.sink.write_str(slice)?;

            self.remaining_needle = self.needle;

        }

        Ok(())

    }

}

impl<'a, W, C> fmt::Write for ReplaceWriter<'a, W, C>

where

    W: fmt::Write + ?Sized,

    C: Writeable,

{

    fn write_str(&mut self, s: &str) -> fmt::Result {

        for c in s.chars() {

            self.write_char(c)?;

        }

        Ok(())

    }

    fn write_char(&mut self, c: char) -> fmt::Result {

        // If the needle is empty, we just pass through the characters.

        if self.needle.is_empty() {

            return self.sink.write_char(c);

        }

        let mut matched = self.matched_len();

        // KMP State Transition:

        // While we have a mismatch and we are not at the start of the needle,

        // backtrack using the prefix function.

        while matched > 0 && !self.remaining_needle.starts_with(c) {

            let old_j = matched;

            matched = get_pi_bytes(self.needle, old_j);

            // Since we backtracked, the prefix of length `old_j - j` is no longer

            // part of the potential match. We write it to the sink as a single slice.

            let slice = self.needle.get(0..(old_j - matched)).ok_or(fmt::Error)?;

            self.sink.write_str(slice)?;

            // Update remaining_needle to reflect the new matched length.

            self.remaining_needle = self.needle.get(matched..).ok_or(fmt::Error)?;

        }

        // If the character matches the next character in the needle, advance the match state.

        if self.remaining_needle.starts_with(c) {

            // Advance remaining_needle by the matched character.

            self.remaining_needle = self

                .remaining_needle

                .get(c.len_utf8()..)

                .ok_or(fmt::Error)?;

            if self.remaining_needle.is_empty() {

                // Full match found! Write the replacement instead of the needle.

                self.replacement.write_to(self.sink)?;

                // Reset match state.

                self.remaining_needle = self.needle;

            }

        } else {

            // Mismatch at the very beginning of the needle. Write the character as is.

            self.sink.write_char(c)?;

        }

        Ok(())

    }

}

impl<A, C> Writeable for Replace<A, &str, C>

where

    A: Writeable,

    C: Writeable,

{

    // We do not implement writeable_borrow because it is meant to be a constant-time O(1)

    // operation, but determining if a replacement occurred would require O(N) scanning.

    fn write_to<W: fmt::Write + ?Sized>(&self, sink: &mut W) -> fmt::Result {

        let mut writer = ReplaceWriter::new(sink, self.needle, &self.replacement);

        self.source.write_to(&mut writer)?;

        writer.finalize()

    }

    fn writeable_length_hint(&self) -> LengthHint {

        let source_hint = self.source.writeable_length_hint();

        let needle_len = self.needle.len();

        let replacement_hint = self.replacement.writeable_length_hint();

        // If needle and replacement have same exact length, length is unchanged.

        if let Some(r_upper) = replacement_hint.1 {

            if replacement_hint.0 == r_upper && needle_len == r_upper {

                return source_hint;

            }

        }

        let mut lower = 0;

        let mut upper = None;

        // If replacement is always larger than or equal to needle:

        // New length is at least the source length.

        if replacement_hint.0 >= needle_len {

            lower = source_hint.0;

        }

        // If replacement is always smaller than or equal to needle:

        // New length is at most the source length.

        if let Some(r_upper) = replacement_hint.1 {

            if r_upper <= needle_len {

                upper = source_hint.1;

            }

        }

        LengthHint(lower, upper)

    }

}

impl_display_with_writeable!(Replace<A, &'a str, C>, #[cfg(feature = "alloc")], where 'a, A: Writeable, C: Writeable);

#[test]

fn test_replace() {

    use crate::assert_writeable_eq;

    use crate::concat::Concat;

    // Basic replacement

    let replace1 = Replace {

        source: Concat("Hello", " 10 22 1101 33"),

        needle: "10",

        replacement: Concat("4", "4"),

    };

    assert_writeable_eq!(replace1, "Hello 44 22 1441 33");

    // Empty needle (should just write source)

    let replace2 = Replace {

        source: "Hello World",

        needle: "",

        replacement: "X",

    };

    assert_writeable_eq!(replace2, "Hello World");

    // Empty replacement

    let replace3 = Replace {

        source: "Hello 10 World 10",

        needle: "10",

        replacement: "",

    };

    assert_writeable_eq!(replace3, "Hello  World ");

    // Needle not found

    let replace4 = Replace {

        source: "Hello World",

        needle: "10",

        replacement: "X",

    };

    assert_writeable_eq!(replace4, "Hello World");

    // Needle at the beginning

    let replace5 = Replace {

        source: "10 Hello World",

        needle: "10",

        replacement: "X",

    };

    assert_writeable_eq!(replace5, "X Hello World");

    // Needle at the end

    let replace6 = Replace {

        source: "Hello World 10",

        needle: "10",

        replacement: "X",

    };

    assert_writeable_eq!(replace6, "Hello World X");

    // Overlapping needles (should consume and not match again)

    let replace7 = Replace {

        source: "ababa",

        needle: "aba",

        replacement: "X",

    };

    assert_writeable_eq!(replace7, "Xba");

    // Self-overlap but no match

    let replace8 = Replace {

        source: "aab",

        needle: "aac",

        replacement: "X",

    };

    assert_writeable_eq!(replace8, "aab");

    // Multi-byte UTF-8

    let replace9 = Replace {

        source: "🚀 🛸 🚀🚀 🚁",

        needle: "🚀",

        replacement: "星",

    };

    assert_writeable_eq!(replace9, "星 🛸 星星 🚁");

    // Multi-byte UTF-8 with partial match

    let replace10 = Replace {

        source: "🚀🚁",

        needle: "🚀🛸",

        replacement: "星",

    };

    assert_writeable_eq!(replace10, "🚀🚁");

    // Multi-byte UTF-8 with backtracking (no match)

    let replace11 = Replace {

        source: "🚀🚀🚁",

        needle: "🚀🚀🛸",

        replacement: "星",

    };

    assert_writeable_eq!(replace11, "🚀🚀🚁");

    // Multi-byte UTF-8 with backtracking (match)

    let replace12 = Replace {

        source: "🚀🚀🚀🛸",

        needle: "🚀🚀🛸",

        replacement: "星",

    };

    assert_writeable_eq!(replace12, "🚀星");

}