toklen 0.2.0

use std::cell::RefCell;
use std::sync::atomic::{AtomicUsize, Ordering};

use fancy_regex::Regex;
use serde::Deserialize;

use super::Error;
use crate::pre_tokenized::PreTokenizedString;
use crate::pre_tokenized::PtSplit;

// Thread-local cache of previous Split results for incremental re-use.
//
// `const {}` enable a more efficient thread local implementation.
thread_local! {
    static SPLIT_CACHE: RefCell<SplitCache> = const {
        RefCell::new(SplitCache { split_id: 0, prev_input: Vec::new(), prev_matches: Vec::new() })
    };
}

struct SplitCache {
    split_id: usize,
    prev_input: Vec<u8>,
    // Use u32 instead of usize to
    // make the data more compact.
    prev_matches: Vec<(u32, u32)>,
}

/// Minimum shared prefix length (bytes) before incremental re-use kicks in.
const INCREMENTAL_MIN_PREFIX: usize = 4096;

/// Wrapper around a JIT-compiled PCRE2 regex.
struct Pcre2Regex(pcre2::bytes::Regex);

impl std::fmt::Debug for Pcre2Regex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("Pcre2Regex(...)")
    }
}

impl Clone for Pcre2Regex {
    fn clone(&self) -> Self {
        Self(
            pcre2::bytes::RegexBuilder::new()
                .utf(true)
                .ucp(true)
                .jit_if_available(true)
                .build(self.0.as_str())
                .expect("re-compile PCRE2 regex"),
        )
    }
}

/// A pattern for a Split pre-tokenizer.
#[derive(Clone, Debug, Deserialize)]
pub enum Pattern {
    String(String),
    Regex(String),
}

impl Pattern {
    fn source(&self) -> String {
        match self {
            Self::String(s) => fancy_regex::escape(s).to_string(),
            Self::Regex(r) => r.clone(),
        }
    }
}

/// How matched delimiters are handled in the output.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)]
pub enum SplitBehavior {
    Removed,
    #[default]
    Isolated,
    MergedWithPrevious,
    MergedWithNext,
    Contiguous,
}

/// Raw deserialization helper for [`Split`].
#[derive(Deserialize)]
struct SplitRaw {
    pattern: Pattern,
    #[serde(default)]
    behavior: SplitBehavior,
    #[serde(default)]
    invert: bool,
}

/// Monotonic counter for unique Split instance IDs.
static SPLIT_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);

/// A compiled Split pre-tokenizer.
#[derive(Clone, Debug, Deserialize)]
#[serde(try_from = "SplitRaw")]
pub struct Split {
    #[serde(skip)]
    id: usize,
    regex: Regex,
    behavior: SplitBehavior,
    invert: bool,
    pcre2_regex: Option<Pcre2Regex>,
}

/// Try to compile a PCRE2 JIT regex from `source`.
fn try_compile_pcre2(source: &str) -> Option<Pcre2Regex> {
    if source.contains("&&") {
        return None;
    }
    let re = pcre2::bytes::RegexBuilder::new()
        .utf(true)
        .ucp(true)
        .jit_if_available(true)
        .build(source)
        .ok()?;
    Some(Pcre2Regex(re))
}

impl TryFrom<SplitRaw> for Split {
    type Error = Error;

    fn try_from(raw: SplitRaw) -> Result<Self, Error> {
        let source = raw.pattern.source();
        Ok(Self {
            id: SPLIT_ID_COUNTER.fetch_add(1, Ordering::Relaxed),
            regex: Regex::new(&source)?,
            behavior: raw.behavior,
            invert: raw.invert,
            pcre2_regex: try_compile_pcre2(&source),
        })
    }
}

impl Split {
    /// Refine the splits of a [`PreTokenizedString`] in place.
    pub fn pre_tokenize(&self, pts: &mut PreTokenizedString) -> Result<(), Error> {
        if self.pcre2_regex.is_some()
            && self.behavior == SplitBehavior::Isolated
            && !self.invert
            && pts.splits.len() == 1
            && pts.splits[0].token_id.is_none()
        {
            return self.pre_tokenize_pcre2_isolated(pts);
        }

        let mut new_splits = Vec::with_capacity(pts.splits.len() << 1);

        for split in &pts.splits {
            if split.token_id.is_some() {
                new_splits.push(split.clone());
                continue;
            }

            let text = pts.split_text(split);
            if text.is_empty() {
                continue;
            }

            let base = split.range.start;
            let segments = self.find_segments(text)?;
            let ranges = self.apply_behavior(&segments);
            for (s, e) in ranges {
                if s < e {
                    new_splits.push(PtSplit {
                        range: (base + s as usize)..(base + e as usize),
                        token_id: None,
                    });
                }
            }
        }

        pts.splits = new_splits;
        Ok(())
    }

    /// Fast path for PCRE2 JIT + Isolated behavior.
    fn pre_tokenize_pcre2_isolated(&self, pts: &mut PreTokenizedString) -> Result<(), Error> {
        let buffer = pts.buffer.as_str();
        let bytes = buffer.as_bytes();
        let pcre2 = self.pcre2_regex.as_ref().unwrap();

        let split = &pts.splits[0];
        let base = split.range.start;
        let text = &buffer[split.range.clone()];

        let split_id = self.id;
        let (mut matches, restart_pos) = SPLIT_CACHE.with(|c| {
            let mut cache = c.borrow_mut();
            if cache.split_id != split_id {
                cache.split_id = split_id;
                cache.prev_input.clear();
                cache.prev_matches.clear();
                return (Vec::new(), 0u32);
            }
            let common_len = common_prefix_len(&cache.prev_input, bytes);

            if common_len >= INCREMENTAL_MIN_PREFIX && !cache.prev_matches.is_empty() {
                let common_len = common_len as u32;
                let reuse_count = cache
                    .prev_matches
                    .partition_point(|&(_, end)| end < common_len);
                let restart = if reuse_count > 0 {
                    cache.prev_matches[reuse_count - 1].1
                } else {
                    0u32
                };
                let mut m = std::mem::take(&mut cache.prev_matches);
                m.truncate(reuse_count);
                (m, restart)
            } else {
                (Vec::new(), 0u32)
            }
        });

        // Run PCRE2 on the portion after the reusable prefix.
        let suffix = &text[restart_pos as usize..];
        if !suffix.is_empty() {
            let base_index = base + restart_pos as usize;
            let suffix_matches = find_matches_pcre2(suffix, base_index, pcre2)?;
            matches.extend(suffix_matches);
        }

        // Build splits from matches.
        let text_len = text.len();
        let mut new_splits = Vec::with_capacity(matches.len() << 1);
        let mut prev = base as u32;
        for &(s, e) in &matches {
            if s > prev {
                new_splits.push(PtSplit {
                    range: (prev as usize)..(s as usize),
                    token_id: None,
                });
            }
            new_splits.push(PtSplit {
                range: (s as usize)..(e as usize),
                token_id: None,
            });
            prev = e;
        }
        let end = (base + text_len) as u32;
        if prev < end {
            new_splits.push(PtSplit {
                range: (prev as usize)..(end as usize),
                token_id: None,
            });
        }

        // Update the cache.
        SPLIT_CACHE.with(|c| {
            let mut cache = c.borrow_mut();
            let input_buf = std::mem::take(&mut cache.prev_input);
            if input_buf.len() == bytes.len() {
                cache.prev_input = input_buf;
                cache.prev_input.copy_from_slice(bytes);
            } else {
                cache.prev_input = bytes.to_vec();
            }
            cache.prev_matches = matches;
        });

        pts.splits = new_splits;
        Ok(())
    }

    fn find_segments(&self, input: &str) -> Result<Vec<(u32, u32, bool)>, Error> {
        if let Some(pcre2) = &self.pcre2_regex {
            let matches = find_matches_pcre2(input, 0, pcre2)?;
            return Ok(matches_to_segments(
                &matches,
                input.len() as u32,
                self.invert,
            ));
        }
        self.find_segments_seq(input)
    }

    /// Sequential regex matching (fancy_regex fallback).
    fn find_segments_seq(&self, input: &str) -> Result<Vec<(u32, u32, bool)>, Error> {
        let mut segments: Vec<(u32, u32, bool)> = Vec::new();
        let mut prev_end = 0usize;

        for m in self.regex.find_iter(input) {
            let m = m?;
            if m.start() == m.end() {
                continue;
            }
            if m.start() > prev_end {
                segments.push((prev_end as u32, m.start() as u32, false));
            }
            segments.push((m.start() as u32, m.end() as u32, true));
            prev_end = m.end();
        }
        if prev_end < input.len() {
            segments.push((prev_end as u32, input.len() as u32, false));
        }

        if self.invert {
            for seg in &mut segments {
                seg.2 = !seg.2;
            }
        }

        Ok(segments)
    }

    /// Phase 2: merge / remove / isolate segments according to behavior.
    fn apply_behavior(&self, segments: &[(u32, u32, bool)]) -> Vec<(u32, u32)> {
        match self.behavior {
            SplitBehavior::Removed => segments
                .iter()
                .filter(|&&(_, _, is_match)| !is_match)
                .map(|&(s, e, _)| (s, e))
                .collect(),

            SplitBehavior::Isolated => segments.iter().map(|&(s, e, _)| (s, e)).collect(),

            SplitBehavior::Contiguous => {
                let mut result: Vec<(u32, u32)> = Vec::new();
                let mut prev_match = None;
                for &(s, e, is_match) in segments {
                    if prev_match == Some(is_match) {
                        if let Some(last) = result.last_mut() {
                            last.1 = e;
                        }
                    } else {
                        result.push((s, e));
                    }
                    prev_match = Some(is_match);
                }
                result
            }

            SplitBehavior::MergedWithPrevious => {
                let mut result: Vec<(u32, u32)> = Vec::new();
                let mut prev_was_match = false;
                for &(s, e, is_match) in segments {
                    if is_match && !prev_was_match {
                        if let Some(last) = result.last_mut() {
                            last.1 = e;
                        } else {
                            result.push((s, e));
                        }
                    } else {
                        result.push((s, e));
                    }
                    prev_was_match = is_match;
                }
                result
            }

            SplitBehavior::MergedWithNext => {
                let mut result: Vec<(u32, u32)> = Vec::new();
                let mut prev_was_match = false;
                for &(s, e, is_match) in segments.iter().rev() {
                    if is_match && !prev_was_match {
                        if let Some(last) = result.last_mut() {
                            last.0 = s;
                        } else {
                            result.push((s, e));
                        }
                    } else {
                        result.push((s, e));
                    }
                    prev_was_match = is_match;
                }
                result.reverse();
                result
            }
        }
    }
}

fn matches_to_segments(
    matches: &[(u32, u32)],
    input_len: u32,
    invert: bool,
) -> Vec<(u32, u32, bool)> {
    let mut segments = Vec::with_capacity((matches.len() << 1) + 1);
    let mut prev = 0u32;
    for &(s, e) in matches {
        if s > prev {
            segments.push((prev, s, invert));
        }
        segments.push((s, e, !invert));
        prev = e;
    }
    if prev < input_len {
        segments.push((prev, input_len, invert));
    }
    segments
}

/// Find the length of the common prefix between two byte slices.
fn common_prefix_len(a: &[u8], b: &[u8]) -> usize {
    let min_len = a.len().min(b.len());
    let chunks = min_len >> 3; // min_len / 8
    for i in 0..chunks {
        let off = i << 3; // i * 8
        let wa = u64::from_ne_bytes(a[off..off + 8].try_into().unwrap());
        let wb = u64::from_ne_bytes(b[off..off + 8].try_into().unwrap());
        if wa != wb {
            let diff = wa ^ wb;
            // off + (diff.trailing_zeros() / 8) as usize
            return off + (diff.trailing_zeros() >> 3) as usize;
        }
    }
    let tail_start = chunks << 3; // chunk * 8
    for i in tail_start..min_len {
        if a[i] != b[i] {
            return i;
        }
    }
    min_len
}

/// Find all pattern matches using PCRE2 JIT.
fn find_matches_pcre2(
    input: &str,
    base: usize,
    regex: &Pcre2Regex,
) -> Result<Vec<(u32, u32)>, Error> {
    let mut matches = Vec::with_capacity(input.len() / 3);
    let bytes = input.as_bytes();
    let mut pos = 0;
    while pos < bytes.len() {
        match regex.0.find_at(bytes, pos) {
            Ok(Some(m)) => {
                if m.start() == m.end() {
                    pos = m.end() + 1;
                    continue;
                }
                matches.push(((base + m.start()) as u32, (base + m.end()) as u32));
                pos = m.end();
            }
            Ok(None) => break,
            Err(e) => return Err(Error::Unsupported(format!("PCRE2: {e}"))),
        }
    }
    Ok(matches)
}