#!/usr/bin/env python3
"""Semantic chunking strategies: compare fixed-size vs structure-aware splitting.

Demonstrates three chunkers and shows why structure-aware splitting preserves
meaningful units (headings, paragraphs, sentences) vs naive fixed char counts.
"""
from __future__ import annotations

import re
from typing import Callable


SAMPLE_DOC = """# Introduction

Retrieval-augmented generation (RAG) combines information retrieval with
language models. It grounds answers in retrieved documents.

## Architecture

The pipeline has three stages: indexing, retrieval, and generation. Each
stage has distinct failure modes.

## Chunking

Chunking splits documents into embeddable units. Fixed-size chunks are a
fallback, not a strategy. Semantic boundaries preserve meaning.
"""


def fixed_size_chunks(text: str, size: int = 80, overlap: int = 10) -> list[str]:
    """Naive fixed-char chunking with overlap (the common anti-pattern)."""
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + size, len(text))
        chunks.append(text[start:end])
        if end == len(text):
            break
        start = end - overlap
    return chunks


def paragraph_chunks(text: str) -> list[str]:
    """Split on blank lines (paragraph boundaries) -- a real semantic unit."""
    return [p.strip() for p in text.split("\n\n") if p.strip()]


def heading_chunks(text: str) -> list[str]:
    """Split on markdown headings, keeping heading + following content together."""
    parts = re.split(r'(?=^#{1,6}\s)', text, flags=re.MULTILINE)
    return [p.strip() for p in parts if p.strip()]


def sentence_chunks(text: str, max_sentences: int = 2) -> list[str]:
    """Group consecutive sentences, ending at sentence boundaries."""
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    chunks = []
    for i in range(0, len(sentences), max_sentences):
        chunk = " ".join(sentences[i:i + max_sentences]).strip()
        if chunk:
            chunks.append(chunk)
    return chunks


def report(name: str, chunks: list[str]) -> None:
    print(f"\n=== {name} ({len(chunks)} chunks) ===")
    for i, c in enumerate(chunks):
        print(f"  [{i}] ({len(c)} chars) {c[:70]}{'...' if len(c) > 70 else ''}")


def main() -> None:
    report("Fixed-size (80 chars, overlap 10)", fixed_size_chunks(SAMPLE_DOC))
    report("Paragraph", paragraph_chunks(SAMPLE_DOC))
    report("Heading + section", heading_chunks(SAMPLE_DOC))
    report("Sentence (max 2 per chunk)", sentence_chunks(SAMPLE_DOC))

    print("\n=== Notes ===")
    print("- Fixed-size cuts mid-sentence/mid-heading -> embeds badly.")
    print("- Heading chunks keep a coherent 'answerable unit' together.")
    print("- Paragraph/sentence chunks honor natural boundaries.")
    print("- Attach metadata (source, heading path) for filtering + citation.")


if __name__ == "__main__":
    main()