import json
import torch
from pathlib import Path
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling_core.types.doc import PictureItem

# 1. Hardware Monkeypatch for Apple Silicon M3 CPU routing
torch.backends.mps.is_available = lambda: False
torch.backends.mps.is_built = lambda: False
torch.set_default_device("cpu")

def convert_to_json_and_images(pdf_name: str, output_json_name: str, image_folder_name: str = "extracted_images"):
    source = f"./{pdf_name}"
    output_dir = Path(image_folder_name)
    output_dir.mkdir(exist_ok=True)

    # 2. Configure pipeline to capture images
    pipeline_options = PdfPipelineOptions()
    pipeline_options.images_scale = 2.0             # 2.0 = roughly 144 DPI for crisp extraction
    pipeline_options.generate_picture_images = True # Required to isolate figure items

    # 3. Initialize converter with hardware-isolated options
    converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
        }
    )

    print(f"🚀 Processing {pdf_name} and building document schema...")
    result = converter.convert(source)

    # 4. Loop through the document items and save physical picture blocks
    picture_counter = 0
    for element, _level in result.document.iterate_items():
        if isinstance(element, PictureItem):
            picture_counter += 1
            image_path = output_dir / f"figure_{picture_counter}.png"
            
            # Save the raw image block cleanly using Pillow
            element.get_image(result.document).save(image_path, "PNG")

    print(f"✨ Compilation finished! Serializing to JSON schema...")
    
    # 5. Export the Document object directly into a standard Python dictionary matrix
    doc_dict = result.document.export_to_dict()
    
    # Write cleanly structured, indented JSON payload to disk
    with open(output_json_name, "w", encoding="utf-8") as f:
        json.dump(doc_dict, f, ensure_ascii=False, indent=2)
    
    print(f"✅ Success! Saved {output_json_name} and extracted {picture_counter} figures to ./{image_folder_name}/")

# Run it for your target wiring rules document
convert_to_json_and_images(
    "1-AS_NZS_3000_2018_Wiring_Rules_Standards_Australia.pdf", 
    "1-AS_NZS_3000_2018_Wiring_Rules_Standards_Australia.json"
)