import logging

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
    ThreadedPdfPipelineOptions,
    TableStructureV2Options,
    TableFormerMode,
    EasyOcrOptions,
)
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
from docling_core.transforms.serializer.markdown import (
    MarkdownParams,
    MarkdownTableSerializer,
)

from docling_core.transforms.chunker.hierarchical_chunker import (
    ChunkingDocSerializer,
    ChunkingSerializerProvider,
)
import tiktoken

log.info("Starting bug reproduction script...")
log.info("Configuring document converter with CPU acceleration...")
pdf_pipline_options = ThreadedPdfPipelineOptions(
    accelerator_options=AcceleratorOptions(
        device=AcceleratorDevice.CPU
    ),  # Use CPU acceleration
    threads=4,
    max_concurrent_results=4,
    ocr_batch_size=4,  # default 4
    layout_batch_size=4,  # default 4
    table_batch_size=4,  # currently not using GPU batching
    ocr_options=EasyOcrOptions(
        lang=[
            "en",
            "de",
            "fr",
            "es",
            "it",
            "pt",
        ],  # Add more languages as needed
    ),
)
pdf_pipline_options.enable_remote_services = True
pdf_pipline_options.allow_external_plugins = True
pdf_pipline_options.do_ocr = True
pdf_pipline_options.do_table_structure = True
pdf_pipline_options.do_code_enrichment = False
pdf_pipline_options.do_formula_enrichment = False
pdf_pipline_options.do_chart_extraction = False
pdf_pipline_options.table_structure_options = TableStructureV2Options(
    do_cell_matching=True, mode=TableFormerMode.ACCURATE
)
doc_converter = DocumentConverter(
    allowed_formats=[
        InputFormat.PDF,
        # Plain Text formats
        InputFormat.HTML,
        InputFormat.MD,
        InputFormat.CSV,
        # Office formats
        InputFormat.DOCX,
        InputFormat.XLSX,
        InputFormat.PPTX,
    ],
    format_options={
        InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_pipline_options)
    },
)

log.info("Document converter configured. Starting conversion of test document...")
conv_result = doc_converter.convert(
    "./f2552dc4-ad0d-4086-bc3c-cb4cfc2a8aa5.pdf"
)
tokenizer = OpenAITokenizer(
    tokenizer=tiktoken.encoding_for_model("text-embedding-3-small"),
    max_tokens=8192,
)
log.info("Document conversion completed. Starting chunking of converted document...")


class SerializerProvider(ChunkingSerializerProvider):
    def get_serializer(self, doc):
        return ChunkingDocSerializer(
            doc=doc,
            table_serializer=MarkdownTableSerializer(),
        )

chunker = HybridChunker(
    tokenizer=tokenizer,
    max_tokens=8192,
    repeat_table_headers=True,
    merge_peers=True,
    serializer_provider=SerializerProvider(),
)
chunk_iter = chunker.chunk(dl_doc=conv_result.document)
for i, chunk in enumerate(chunk_iter):
    log.debug(f"Chunk {i}: {chunk}")
    enriched_text = chunker.contextualize(chunk=chunk)
    if tokenizer.count_tokens(enriched_text) > 8192:
        log.warning(f"Chunk {i} exceeds token limit after contextualization: {tokenizer.count_tokens(enriched_text)} tokens")