cml-rs 0.4.0

//! Profile-aware semantic chunking for CML v0.2 documents.

use crate::id_generator::ElementId;
use crate::types::*;
use std::collections::HashMap;

/// Maximum tokens per chunk before splitting (configurable).
pub const MAX_CHUNK_TOKENS: usize = 512;

/// Overlap tokens between consecutive sub-chunks.
pub const CHUNK_OVERLAP_TOKENS: usize = 50;

/// Estimated characters per token (rough heuristic).
const CHARS_PER_TOKEN: usize = 4;

/// A semantic chunk with hierarchical metadata.
#[derive(Debug, Clone, PartialEq)]
pub struct Chunk {
    /// Unique chunk ID (element ID or element ID + sub-chunk index).
    pub id: String,

    /// Parent chunk ID (for sub-chunks or nested elements).
    pub parent_id: Option<String>,

    /// Content hash for validation.
    pub content_hash: String,

    /// CML profile (e.g. code:api, legal:constitution).
    pub profile: String,

    /// Element type (section, paragraph, code, etc.).
    pub element_type: String,

    /// Actual text content.
    pub content: String,

    /// Estimated token count.
    pub token_count: usize,

    /// Additional metadata.
    pub metadata: HashMap<String, String>,
}

impl Chunk {
    /// Create a new chunk from an element.
    fn from_element(
        id: ElementId,
        profile: &str,
        element_type: &str,
        content: String,
        metadata: HashMap<String, String>,
        parent_id: Option<String>,
    ) -> Self {
        let token_count = estimate_tokens(&content);
        Self {
            id: id.id.clone(),
            parent_id,
            content_hash: id.content_hash,
            profile: profile.to_string(),
            element_type: element_type.to_string(),
            content,
            token_count,
            metadata,
        }
    }

    /// Create a sub-chunk from a parent chunk.
    fn sub_chunk(
        parent_id: &str,
        index: usize,
        profile: &str,
        element_type: &str,
        content: String,
        metadata: HashMap<String, String>,
    ) -> Self {
        let id = format!("{}#{}", parent_id, index);
        let token_count = estimate_tokens(&content);
        Self {
            id,
            parent_id: Some(parent_id.to_string()),
            content_hash: String::new(),
            profile: profile.to_string(),
            element_type: element_type.to_string(),
            content,
            token_count,
            metadata,
        }
    }
}

/// Estimate token count from text (rough heuristic: 1 token per 4 chars).
fn estimate_tokens(text: &str) -> usize {
    text.len() / CHARS_PER_TOKEN
}

/// Split large text into overlapping chunks.
fn split_into_subchunks(text: &str, max_tokens: usize, overlap_tokens: usize) -> Vec<String> {
    let max_chars = max_tokens * CHARS_PER_TOKEN;
    let overlap_chars = overlap_tokens * CHARS_PER_TOKEN;

    if text.len() <= max_chars {
        return vec![text.to_string()];
    }

    let mut chunks = Vec::new();
    let mut start = 0;

    while start < text.len() {
        let end = (start + max_chars).min(text.len());
        let chunk_text = &text[start..end];
        chunks.push(chunk_text.to_string());

        if end >= text.len() {
            break;
        }

        start = end.saturating_sub(overlap_chars);
    }

    chunks
}

/// Profile-aware chunker for CML v0.2 documents.
pub struct CmlChunker;

impl CmlChunker {
    /// Chunk a CML v0.2 document into hierarchical chunks.
    pub fn chunk_document(doc: &CmlDocument) -> Vec<Chunk> {
        let mut chunks = Vec::new();
        let base_id = doc
            .id
            .as_deref()
            .map(sanitize_segment)
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| sanitize_segment(&doc.header.title));
        let base_id = if base_id.is_empty() {
            "doc".to_string()
        } else {
            base_id
        };

        for (index, block) in doc.body.blocks.iter().enumerate() {
            chunk_block(
                block,
                &doc.profile,
                &base_id,
                &[],
                None,
                index + 1,
                &mut chunks,
            );
        }

        chunks
    }
}

fn chunk_block(
    block: &BlockElement,
    profile: &str,
    base_id: &str,
    parent_path: &[String],
    parent_id: Option<String>,
    index: usize,
    chunks: &mut Vec<Chunk>,
) {
    let (element_type, explicit_id) = block_identity(block);
    let segment = make_segment(explicit_id, element_type, index);
    let mut path = parent_path.to_vec();
    path.push(segment);

    let content = block_text(block);
    let mut metadata = HashMap::new();
    add_block_metadata(block, &mut metadata);

    let element_id = ElementId::new(join_path(base_id, &path), &content);
    let chunk = Chunk::from_element(
        element_id,
        profile,
        element_type,
        content,
        metadata,
        parent_id.clone(),
    );

    let chunk_id = chunk.id.clone();
    push_chunk_with_subchunks(chunks, chunk);

    match block {
        BlockElement::Section(section) => {
            for (child_index, child) in section.content.iter().enumerate() {
                chunk_block(
                    child,
                    profile,
                    base_id,
                    &path,
                    Some(chunk_id.clone()),
                    child_index + 1,
                    chunks,
                );
            }
        }
        BlockElement::Aside(aside) => {
            for (child_index, child) in aside.content.iter().enumerate() {
                chunk_block(
                    child,
                    profile,
                    base_id,
                    &path,
                    Some(chunk_id.clone()),
                    child_index + 1,
                    chunks,
                );
            }
        }
        BlockElement::Quote(quote) => {
            for (child_index, child) in quote.content.iter().enumerate() {
                chunk_block(
                    child,
                    profile,
                    base_id,
                    &path,
                    Some(chunk_id.clone()),
                    child_index + 1,
                    chunks,
                );
            }
        }
        BlockElement::List(list) => {
            for (item_index, item) in list.items.iter().enumerate() {
                chunk_list_item(
                    item,
                    profile,
                    base_id,
                    &path,
                    &chunk_id,
                    item_index + 1,
                    chunks,
                );
            }
        }
        BlockElement::Table(table) => {
            chunk_table_rows(table, profile, base_id, &path, &chunk_id, chunks);
        }
        BlockElement::Paragraph(_)
        | BlockElement::Heading(_)
        | BlockElement::Code(_)
        | BlockElement::Break(_)
        | BlockElement::Figure(_) => {}
    }
}

fn chunk_list_item(
    item: &ListItem,
    profile: &str,
    base_id: &str,
    parent_path: &[String],
    parent_id: &str,
    index: usize,
    chunks: &mut Vec<Chunk>,
) {
    let segment = make_segment(item.id.as_ref(), "item", index);
    let mut path = parent_path.to_vec();
    path.push(segment);

    match &item.content {
        ListItemContent::Inline(inline) => {
            let content = inline_texts(inline);
            let mut metadata = HashMap::new();
            metadata.insert("list_index".to_string(), index.to_string());
            if let Some(id) = &item.id {
                metadata.insert("source_id".to_string(), id.clone());
            }
            let element_id = ElementId::new(join_path(base_id, &path), &content);
            let chunk = Chunk::from_element(
                element_id,
                profile,
                "list_item",
                content,
                metadata,
                Some(parent_id.to_string()),
            );
            push_chunk_with_subchunks(chunks, chunk);
        }
        ListItemContent::Block(blocks) => {
            for (child_index, block) in blocks.iter().enumerate() {
                chunk_block(
                    block,
                    profile,
                    base_id,
                    &path,
                    Some(parent_id.to_string()),
                    child_index + 1,
                    chunks,
                );
            }
        }
    }
}

fn chunk_table_rows(
    table: &Table,
    profile: &str,
    base_id: &str,
    parent_path: &[String],
    parent_id: &str,
    chunks: &mut Vec<Chunk>,
) {
    for (row_index, row) in table.body.rows.iter().enumerate() {
        let row_segment = format!("row-{}", row_index + 1);
        let mut row_path = parent_path.to_vec();
        row_path.push(row_segment);

        let row_content = row
            .columns
            .iter()
            .map(|col| cell_text(&col.cell))
            .collect::<Vec<String>>()
            .join(" | ");
        let mut metadata = HashMap::new();
        metadata.insert("row_index".to_string(), (row_index + 1).to_string());
        let element_id = ElementId::new(join_path(base_id, &row_path), &row_content);
        let row_chunk = Chunk::from_element(
            element_id,
            profile,
            "table_row",
            row_content,
            metadata,
            Some(parent_id.to_string()),
        );
        let row_chunk_id = row_chunk.id.clone();
        push_chunk_with_subchunks(chunks, row_chunk);

        for (col_index, col) in row.columns.iter().enumerate() {
            let cell_segment = format!("cell-{}", col_index + 1);
            let mut cell_path = row_path.clone();
            cell_path.push(cell_segment);
            let cell_content = cell_text(&col.cell);
            let mut cell_meta = HashMap::new();
            cell_meta.insert("row_index".to_string(), (row_index + 1).to_string());
            cell_meta.insert("col_index".to_string(), (col_index + 1).to_string());
            let element_id = ElementId::new(join_path(base_id, &cell_path), &cell_content);
            let cell_chunk = Chunk::from_element(
                element_id,
                profile,
                "table_cell",
                cell_content,
                cell_meta,
                Some(row_chunk_id.clone()),
            );
            push_chunk_with_subchunks(chunks, cell_chunk);
        }
    }
}

fn push_chunk_with_subchunks(chunks: &mut Vec<Chunk>, chunk: Chunk) {
    let content = chunk.content.clone();
    let element_type = chunk.element_type.clone();
    let profile = chunk.profile.clone();
    let metadata = chunk.metadata.clone();
    let parent_id = chunk.id.clone();
    let token_count = chunk.token_count;
    chunks.push(chunk);

    if token_count <= MAX_CHUNK_TOKENS {
        return;
    }

    let subchunks = split_into_subchunks(&content, MAX_CHUNK_TOKENS, CHUNK_OVERLAP_TOKENS);
    if subchunks.len() <= 1 {
        return;
    }

    for (index, sub) in subchunks.iter().enumerate() {
        let mut sub_metadata = metadata.clone();
        sub_metadata.insert("subchunk_index".to_string(), (index + 1).to_string());
        let sub_chunk = Chunk::sub_chunk(
            &parent_id,
            index + 1,
            &profile,
            &element_type,
            sub.clone(),
            sub_metadata,
        );
        chunks.push(sub_chunk);
    }
}

fn block_identity(block: &BlockElement) -> (&'static str, Option<&String>) {
    match block {
        BlockElement::Section(section) => ("section", section.id.as_ref()),
        BlockElement::Paragraph(paragraph) => ("paragraph", paragraph.id.as_ref()),
        BlockElement::Heading(heading) => ("heading", heading.id.as_ref()),
        BlockElement::Aside(aside) => ("aside", aside.id.as_ref()),
        BlockElement::Quote(quote) => ("quote", quote.id.as_ref()),
        BlockElement::List(list) => ("list", list.id.as_ref()),
        BlockElement::Table(table) => ("table", table.id.as_ref()),
        BlockElement::Code(code) => ("code", code.id.as_ref()),
        BlockElement::Break(_) => ("break", None),
        BlockElement::Figure(figure) => ("figure", figure.id.as_ref()),
    }
}

fn add_block_metadata(block: &BlockElement, metadata: &mut HashMap<String, String>) {
    match block {
        BlockElement::Section(section) => {
            insert_opt(metadata, "section_type", section.section_type.as_ref());
            insert_opt(metadata, "reference", section.reference.as_ref());
        }
        BlockElement::Paragraph(paragraph) => {
            insert_opt(
                metadata,
                "paragraph_type",
                paragraph.paragraph_type.as_ref(),
            );
            insert_opt(metadata, "source_id", paragraph.id.as_ref());
        }
        BlockElement::Heading(heading) => {
            insert_opt(metadata, "heading_type", heading.heading_type.as_ref());
            metadata.insert("size".to_string(), heading.size.to_string());
            insert_opt(metadata, "source_id", heading.id.as_ref());
        }
        BlockElement::Aside(aside) => {
            insert_opt(metadata, "aside_type", aside.aside_type.as_ref());
            metadata.insert(
                "side".to_string(),
                match aside.side {
                    Side::Left => "left".to_string(),
                    Side::Right => "right".to_string(),
                },
            );
            insert_opt(metadata, "source_id", aside.id.as_ref());
        }
        BlockElement::Quote(quote) => {
            insert_opt(metadata, "reference", quote.reference.as_ref());
            insert_opt(metadata, "source", quote.source.as_ref());
            insert_opt(metadata, "source_id", quote.id.as_ref());
        }
        BlockElement::List(list) => {
            if let Some(list_type) = &list.list_type {
                let value = match list_type {
                    ListType::Ordered => "ordered",
                    ListType::Unordered => "unordered",
                };
                metadata.insert("list_type".to_string(), value.to_string());
            }
            if let Some(style) = &list.style {
                let value = match style {
                    ListStyle::Numeric => "numeric",
                    ListStyle::Roman => "roman",
                    ListStyle::Alpha => "alpha",
                    ListStyle::Symbolic => "symbolic",
                };
                metadata.insert("style".to_string(), value.to_string());
            }
            insert_opt(metadata, "source_id", list.id.as_ref());
        }
        BlockElement::Table(table) => {
            insert_opt(metadata, "table_type", table.table_type.as_ref());
            let row_count = table.body.rows.len();
            let col_count = table
                .body
                .rows
                .iter()
                .map(|row| row.columns.len())
                .max()
                .unwrap_or(0);
            metadata.insert("rows".to_string(), row_count.to_string());
            metadata.insert("columns".to_string(), col_count.to_string());
            insert_opt(metadata, "source_id", table.id.as_ref());
        }
        BlockElement::Code(code) => {
            insert_opt(metadata, "lang", code.lang.as_ref());
            if let Some(copyable) = code.copyable {
                metadata.insert("copyable".to_string(), copyable.to_string());
            }
            insert_opt(metadata, "source_id", code.id.as_ref());
        }
        BlockElement::Break(break_el) => {
            insert_opt(metadata, "break_type", break_el.break_type.as_ref());
        }
        BlockElement::Figure(figure) => {
            insert_opt(metadata, "figure_type", figure.figure_type.as_ref());
            insert_opt(metadata, "reference", figure.reference.as_ref());
            insert_opt(metadata, "source_id", figure.id.as_ref());
        }
    }
}

fn insert_opt(map: &mut HashMap<String, String>, key: &str, value: Option<&String>) {
    if let Some(value) = value {
        if !value.is_empty() {
            map.insert(key.to_string(), value.clone());
        }
    }
}

fn block_text(block: &BlockElement) -> String {
    match block {
        BlockElement::Section(section) => blocks_text(&section.content),
        BlockElement::Paragraph(paragraph) => inline_texts(&paragraph.content),
        BlockElement::Heading(heading) => inline_texts(&heading.content),
        BlockElement::Aside(aside) => blocks_text(&aside.content),
        BlockElement::Quote(quote) => blocks_text(&quote.content),
        BlockElement::List(list) => list_text(list),
        BlockElement::Table(table) => table_text(table),
        BlockElement::Code(code) => code.content.clone(),
        BlockElement::Break(_) => String::new(),
        BlockElement::Figure(_) => String::new(),
    }
}

fn blocks_text(blocks: &[BlockElement]) -> String {
    blocks
        .iter()
        .map(block_text)
        .filter(|text| !text.is_empty())
        .collect::<Vec<String>>()
        .join("\n")
}

fn list_text(list: &List) -> String {
    list.items
        .iter()
        .map(|item| match &item.content {
            ListItemContent::Inline(inline) => inline_texts(inline),
            ListItemContent::Block(blocks) => blocks_text(blocks),
        })
        .filter(|text| !text.is_empty())
        .collect::<Vec<String>>()
        .join("\n")
}

fn table_text(table: &Table) -> String {
    table
        .body
        .rows
        .iter()
        .map(|row| {
            row.columns
                .iter()
                .map(|col| cell_text(&col.cell))
                .collect::<Vec<String>>()
                .join(" | ")
        })
        .filter(|text| !text.is_empty())
        .collect::<Vec<String>>()
        .join("\n")
}

fn cell_text(cell: &TableCell) -> String {
    inline_texts(&cell.content)
}

fn inline_texts(elements: &[InlineElement]) -> String {
    elements
        .iter()
        .map(inline_text)
        .collect::<Vec<String>>()
        .join("")
}

fn inline_text(element: &InlineElement) -> String {
    match element {
        InlineElement::Text(text) => text.clone(),
        InlineElement::Em(em) => inline_texts(&em.content),
        InlineElement::Bo(bo) => inline_texts(&bo.content),
        InlineElement::Un(un) => inline_texts(&un.content),
        InlineElement::St(st) => inline_texts(&st.content),
        InlineElement::Snip(snip) => snip.content.clone(),
        InlineElement::Key(key) => key.content.clone(),
        InlineElement::Rf(rf) => {
            if rf.content.is_empty() {
                rf.reference.clone()
            } else {
                rf.content.clone()
            }
        }
        InlineElement::Tg(tg) => {
            if tg.content.is_empty() {
                tg.reference.clone()
            } else {
                tg.content.clone()
            }
        }
        InlineElement::Lk(lk) => {
            if lk.content.is_empty() {
                lk.reference.clone()
            } else {
                lk.content.clone()
            }
        }
        InlineElement::Curr(curr) => format!("{} {}", curr.value, curr.currency_type),
        InlineElement::End(_) => "\n".to_string(),
    }
}

fn make_segment(explicit_id: Option<&String>, fallback: &str, index: usize) -> String {
    match explicit_id {
        Some(id) => sanitize_segment(id),
        None => format!("{}-{}", fallback, index),
    }
}

fn join_path(base_id: &str, path: &[String]) -> String {
    if path.is_empty() {
        base_id.to_string()
    } else {
        format!("{}.{}", base_id, path.join("."))
    }
}

fn sanitize_segment(raw: &str) -> String {
    let mut output = String::new();
    for ch in raw.chars() {
        if ch.is_ascii_alphanumeric() {
            output.push(ch.to_ascii_lowercase());
        } else if ch == '-' || ch == '_' {
            output.push(ch);
        } else {
            output.push('-');
        }
    }
    while output.contains("--") {
        output = output.replace("--", "-");
    }
    output.trim_matches('-').to_string()
}