from pathlib import Path
import json, sys, argparse
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode
from docling.datamodel.accelerator_options import AcceleratorOptions, AcceleratorDevice
from docling.document_converter import DocumentConverter,PdfFormatOption
from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker
root=Path(__file__).resolve().parent
p=argparse.ArgumentParser();p.add_argument('--fast',action='store_true');p.add_argument('--out',default='.');p.add_argument('--artifacts');args=p.parse_args()
out=(root/args.out).resolve();out.mkdir(parents=True,exist_ok=True)
options=PdfPipelineOptions(do_ocr=False,artifacts_path=Path(args.artifacts).resolve() if args.artifacts else None)
options.accelerator_options=AcceleratorOptions(device=AcceleratorDevice.CPU,num_threads=2)
options.table_structure_options.mode=TableFormerMode.FAST if args.fast else TableFormerMode.ACCURATE
converter=DocumentConverter(format_options={InputFormat.PDF:PdfFormatOption(pipeline_options=options)})
document=converter.convert(root/'source.pdf').document
document.save_as_json(out/'document.json')
(out/'result.md').write_text(document.export_to_markdown(),encoding='utf-8')
for i,table in enumerate(document.tables,1):
    table.export_to_dataframe(doc=document).to_csv(out/f'table-{i}.csv',index=False)
    (out/f'table-{i}.html').write_text(table.export_to_html(doc=document))
chunks=[]
for chunk in HierarchicalChunker().chunk(dl_doc=document):
    chunks.append({'text':chunk.text,'headings':chunk.meta.headings,'items':[{'ref':item.self_ref,'pages':[p.page_no for p in item.prov]} for item in chunk.meta.doc_items],'source':'source.pdf'})
(out/'chunks.json').write_text(json.dumps(chunks,ensure_ascii=False,indent=2))