rs-chunks 0.6.4

use serde_json::json;

use super::common::{
    current_section_heading, extract_heading_text, heading_level_txt, heading_path_strings,
    parse_txt_blocks, split_sentences, update_heading_stack, ChunkRecordInput, ContentType,
};

struct IndexedSentence {
    text: String,
    paragraph_index: usize,
    section_heading: Option<String>,
    heading_path: Vec<String>,
}

pub fn build_sentence_chunks(
    bytes: &[u8],
    sentences_per_chunk: usize,
) -> Result<Vec<ChunkRecordInput>, String> {
    if sentences_per_chunk == 0 {
        return Err("sentences_per_chunk must be greater than 0".to_string());
    }

    let text = crate::text_encoding::decode_text(bytes).0;
    // Empty input is not a failure. A blank or whitespace-only document parsed
    // perfectly well; it simply has nothing to chunk, so it returns `[]` like
    // docx/ppt/xlsx always have (TECH_DEBT T6). Reserving errors for genuine
    // parse failures is also what lets `epub::extract` stop swallowing them.
    if text.trim().is_empty() {
        return Ok(Vec::new());
    }

    let blocks = parse_txt_blocks(&text);
    let total = blocks.len();
    let mut result: Vec<ChunkRecordInput> = Vec::new();
    let mut heading_stack: Vec<(u8, String)> = Vec::new();
    let mut sentences: Vec<IndexedSentence> = Vec::new();
    let mut para_index = 0usize;
    let mut chunk_index = 0usize;

    let flush_sentences = |sentences: &mut Vec<IndexedSentence>,
                           result: &mut Vec<ChunkRecordInput>,
                           ci: &mut usize,
                           spc: usize,
                           total: usize| {
        let mut i = 0usize;
        while i < sentences.len() {
            let end = (i + spc).min(sentences.len());
            let window = &sentences[i..end];
            let content = window
                .iter()
                .map(|s| s.text.as_str())
                .collect::<Vec<_>>()
                .join(" ");
            if !content.is_empty() {
                result.push(ChunkRecordInput {
                    content_type: ContentType::Sentence,
                    content,
                    metadata: json!({
                        "sentences_per_chunk":    spc,
                        "actual_sentence_count":  window.len(),
                        "chunk_index":            *ci,
                        "source_paragraph_index": window[0].paragraph_index,
                        "section_heading":        window[0].section_heading,
                        "heading_path":           window[0].heading_path,
                        "document_metadata": { "source_type": "txt", "total_input_blocks": total }
                    }),
                });
                *ci += 1;
            }
            i = end;
        }
        sentences.clear();
    };

    let is_prose = |ct: ContentType| {
        matches!(
            ct,
            ContentType::PlainParagraph
                | ContentType::LongSingleParagraph
                | ContentType::ShortDisconnectedParagraph
        )
    };

    for block in &blocks {
        if block.content_type == ContentType::HeadingSection {
            flush_sentences(
                &mut sentences,
                &mut result,
                &mut chunk_index,
                sentences_per_chunk,
                total,
            );
            let level = heading_level_txt(&block.content);
            let text = extract_heading_text(&block.content);
            update_heading_stack(&mut heading_stack, level, text.clone());
            result.push(ChunkRecordInput {
                content_type: ContentType::HeadingSection,
                content: text.clone(),
                metadata: json!({
                    "section_heading": text,
                    "section_level":   level,
                    "heading_path":    heading_path_strings(&heading_stack),
                    "chunk_index":     chunk_index,
                    "document_metadata": { "source_type": "txt", "total_input_blocks": total }
                }),
            });
            chunk_index += 1;
        } else if block.content_type == ContentType::CodeBlock
            || block.content_type == ContentType::Table
        {
            flush_sentences(
                &mut sentences,
                &mut result,
                &mut chunk_index,
                sentences_per_chunk,
                total,
            );
            result.push(ChunkRecordInput {
                content_type: block.content_type,
                content: block.content.clone(),
                metadata: json!({
                    "section_heading": current_section_heading(&heading_stack),
                    "heading_path":    heading_path_strings(&heading_stack),
                    "chunk_index":     chunk_index,
                    "document_metadata": { "source_type": "txt", "total_input_blocks": total }
                }),
            });
            chunk_index += 1;
            para_index += 1;
        } else if block.content_type == ContentType::BulletNumberedList {
            // List items become individual "sentences".
            let sh = current_section_heading(&heading_stack);
            let hp = heading_path_strings(&heading_stack);
            let pi = para_index;
            for line in block.content.lines() {
                let t = line.trim().to_string();
                if !t.is_empty() {
                    sentences.push(IndexedSentence {
                        text: t,
                        paragraph_index: pi,
                        section_heading: sh.clone(),
                        heading_path: hp.clone(),
                    });
                }
            }
            para_index += 1;
        } else if is_prose(block.content_type) {
            let sh = current_section_heading(&heading_stack);
            let hp = heading_path_strings(&heading_stack);
            let pi = para_index;
            for s in split_sentences(&block.content) {
                sentences.push(IndexedSentence {
                    text: s,
                    paragraph_index: pi,
                    section_heading: sh.clone(),
                    heading_path: hp.clone(),
                });
            }
            para_index += 1;
        }
    }
    flush_sentences(
        &mut sentences,
        &mut result,
        &mut chunk_index,
        sentences_per_chunk,
        total,
    );
    // Empty is not a failure (TECH_DEBT T6): the document parsed, this mode
    // simply produced nothing. Returning `[]` keeps every mode consistent with
    // docx/ppt/xlsx and lets epub distinguish an empty chapter from a broken
    // one without swallowing errors (L14).
    if result.is_empty() {
        return Ok(Vec::new());
    }
    Ok(result)
}