src/chunker.py

import re


def chunk_text(text: str, chunk_size: int = 512, overlap: int = 50) -> list[str]:
    """Cắt text thành chunks theo số ký tự có overlap."""
    if len(text) <= chunk_size:
        return [text]
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks


def chunk_markdown(text: str, chunk_size: int = 1000, overlap: int = 100) -> list[str]:
    """Cắt Markdown theo section (heading), giữ nguyên bảng."""
    # Tách theo heading
    sections = re.split(r'(?=^#{1,3} )', text, flags=re.MULTILINE)
    sections = [s.strip() for s in sections if s.strip()]

    chunks = []
    for section in sections:
        if len(section) <= chunk_size:
            chunks.append(section)
        else:
            # Section quá dài → cắt nhỏ hơn nhưng không cắt giữa bảng
            sub_chunks = _split_preserving_tables(section, chunk_size, overlap)
            chunks.extend(sub_chunks)

    return chunks


def _split_preserving_tables(text: str, chunk_size: int, overlap: int) -> list[str]:
    """Cắt text nhưng không cắt giữa bảng Markdown."""
    # Tách block: bảng vs text thường
    blocks = re.split(r'(\n\|.+?\n(?:\|.+?\n)*)', text, flags=re.DOTALL)
    result = []
    current = ""

    for block in blocks:
        if len(current) + len(block) <= chunk_size:
            current += block
        else:
            if current.strip():
                result.append(current.strip())
            if len(block) > chunk_size:
                # Block quá lớn → cắt cứng
                result.extend(chunk_text(block, chunk_size, overlap))
                current = ""
            else:
                current = block

    if current.strip():
        result.append(current.strip())

    return result


def chunk_pages(pages: list[dict], chunk_size: int = 1000, overlap: int = 100) -> list[dict]:
    """Nhận list pages, trả về list chunks với metadata."""
    chunks = []
    for page in pages:
        parser = page["metadata"].get("parser", "plain")

        if parser == "docling":
            text_chunks = chunk_markdown(page["text"], chunk_size, overlap)
        else:
            text_chunks = chunk_text(page["text"], chunk_size, overlap)

        for i, chunk in enumerate(text_chunks):
            chunks.append({
                "text": chunk,
                "metadata": {
                    **page["metadata"],
                    "chunk_index": i,
                }
            })
    return chunks