cp-parser 0.3.1

//! Text chunking for embedding preparation

use cp_core::{CPError, Chunk, Result};
use text_splitter::{ChunkConfig as TSChunkConfig, MarkdownSplitter};
use uuid::Uuid;

/// Configuration for chunking
#[derive(Debug, Clone)]
pub struct ChunkConfig {
    /// Target chunk size in characters
    pub chunk_size: usize,
    /// Overlap between chunks in characters
    pub overlap: usize,
}

impl Default for ChunkConfig {
    fn default() -> Self {
        Self {
            chunk_size: 1000, // ~250 tokens
            overlap: 200,     // ~50 tokens
        }
    }
}

/// Chunker for splitting text into overlapping segments
pub struct Chunker {
    config: ChunkConfig,
}

impl Default for Chunker {
    fn default() -> Self {
        Self::new(ChunkConfig::default())
    }
}

impl Chunker {
    /// Create a new chunker with the given config
    pub fn new(config: ChunkConfig) -> Self {
        Self { config }
    }

    /// Split text into overlapping chunks
    pub fn chunk(&self, doc_id: Uuid, text: &str) -> Result<Vec<Chunk>> {
        if text.is_empty() {
            return Ok(Vec::new());
        }

        let ts_config = TSChunkConfig::new(self.config.chunk_size)
            .with_overlap(self.config.overlap)
            .map_err(|e| CPError::Parse(format!("Invalid chunk config: {e}")))?
            .with_trim(true);
        let splitter = MarkdownSplitter::new(ts_config);

        let chunks: Vec<Chunk> = splitter
            .chunk_indices(text)
            .enumerate()
            .map(|(seq, (byte_offset, chunk_text))| {
                Chunk::new(doc_id, chunk_text, byte_offset as u64, seq as u32)
            })
            .collect();

        Ok(chunks)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_text() {
        let chunker = Chunker::default();
        let chunks = chunker.chunk(Uuid::new_v4(), "").unwrap();
        assert!(chunks.is_empty());
    }

    #[test]
    fn test_short_text() {
        let chunker = Chunker::default();
        let chunks = chunker.chunk(Uuid::new_v4(), "Short text.").unwrap();
        assert_eq!(chunks.len(), 1);
    }

    #[test]
    fn test_long_text_chunking() {
        let chunker = Chunker::new(ChunkConfig {
            chunk_size: 100,
            overlap: 20,
        });

        let text = "A".repeat(250);
        let chunks = chunker.chunk(Uuid::new_v4(), &text).unwrap();

        assert!(chunks.len() > 1);

        // Check sequence numbers
        for (i, chunk) in chunks.iter().enumerate() {
            assert_eq!(chunk.sequence, i as u32);
        }
    }

    #[test]
    fn test_sentence_boundary() {
        let chunker = Chunker::new(ChunkConfig {
            chunk_size: 20,
            overlap: 5,
        });

        let text = "First sentence. Second sentence. Third sentence.";
        let chunks = chunker.chunk(Uuid::new_v4(), text).unwrap();

        // Should produce multiple chunks (text is longer than chunk_size)
        assert!(chunks.len() > 1);
    }

    #[test]
    fn test_byte_offsets_valid() {
        let chunker = Chunker::new(ChunkConfig {
            chunk_size: 50,
            overlap: 10,
        });

        let text = "# Heading\n\nFirst paragraph with some text.\n\n## Subheading\n\nSecond paragraph with more text here.";
        let chunks = chunker.chunk(Uuid::new_v4(), text).unwrap();

        for chunk in &chunks {
            let offset = chunk.byte_offset as usize;
            // The byte offset should be within the original text
            assert!(
                offset <= text.len(),
                "byte_offset {} exceeds text len {}",
                offset,
                text.len()
            );
        }
    }

    #[test]
    fn test_overlap_shared_text() {
        let chunker = Chunker::new(ChunkConfig {
            chunk_size: 30,
            overlap: 10,
        });

        let text = "Word one. Word two. Word three. Word four. Word five. Word six.";
        let chunks = chunker.chunk(Uuid::new_v4(), text).unwrap();

        if chunks.len() >= 2 {
            // With overlap, there should be some shared content between adjacent chunks
            // (the text-splitter handles overlap at semantic boundaries, so just verify
            // we get multiple chunks and they have valid sequences)
            for (i, chunk) in chunks.iter().enumerate() {
                assert_eq!(chunk.sequence, i as u32);
            }
        }
    }
}