# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "docling>=2.0.0",
#   "docling-core[chunking]>=2.0.0",
#   "transformers>=4.40.0",
# ]
# ///
"""
Minimal HTML chunking script using docling's HybridChunker.

Usage:
  python slowchunk.py /path/to/file.html
"""

import sys

from docling.document_converter import DocumentConverter
from docling.datamodel.base_models import InputFormat
from docling.chunking import HybridChunker
from docling_core.transforms.chunker.tokenizer.huggingface import HuggingFaceTokenizer
from transformers import AutoTokenizer

EMBED_MODEL_ID = "ibm-granite/granite-embedding-30m-english"
MAX_TOKENS = 512


def main():
    if len(sys.argv) < 2:
        print("Usage: python slowchunk.py <html_file>", file=sys.stderr)
        sys.exit(1)

    file_path = sys.argv[1]

    print(f"Reading {file_path}...", file=sys.stderr)
    with open(file_path, "r", encoding="utf-8") as f:
        content = f.read()
    print(f"Read {len(content)} chars.", file=sys.stderr)

    print("Initializing tokenizer...", file=sys.stderr)
    tokenizer = HuggingFaceTokenizer(
        tokenizer=AutoTokenizer.from_pretrained(EMBED_MODEL_ID),
        max_tokens=MAX_TOKENS,
    )

    print("Initializing document converter...", file=sys.stderr)
    converter = DocumentConverter()

    print("Converting HTML to document...", file=sys.stderr)
    result = converter.convert_string(content=content, format=InputFormat.HTML)
    doc = result.document

    print("Initializing chunker...", file=sys.stderr)
    chunker = HybridChunker(tokenizer=tokenizer, merge_peers=True)

    print("Chunking...", file=sys.stderr)
    chunks = list(chunker.chunk(doc))
    print(f"Produced {len(chunks)} chunks.", file=sys.stderr)

    print("Contextualizing and counting tokens...", file=sys.stderr)
    total_tokens = 0
    oversized = 0
    for chunk in chunks:
        text = chunker.contextualize(chunk)
        n = tokenizer.count_tokens(text)
        total_tokens += n
        if n > MAX_TOKENS:
            oversized += 1

    print(
        f"Done. {len(chunks)} chunks, {total_tokens} total tokens, "
        f"avg {total_tokens // max(len(chunks), 1)} tokens/chunk, "
        f"{oversized} oversized.",
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()