import json
import logging
from pathlib import Path
from typing import List

import yaml

# MAKE SURE THIS IMPORT IS PRESENT
from docling.backend.docling_parse_v4_backend import DoclingParseV4DocumentBackend
from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import OcrOptions
from docling.datamodel.pipeline_options import PipelineOptions
from docling.document_converter import DocumentConverter, FormatOption
from docling.pipeline.simple_pipeline import SimplePipeline
from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline

_log = logging.getLogger(__name__)

class OCR_PROCESSOR():
    def __init__(self):
        print("Init OCR_PROCESSOR")

    def process_documents(
        self, input_paths: List[Path], output_dir: Path, ocr_options: OcrOptions
    ):
        """Processes documents using a custom-configured DocumentConverter."""

        if not output_dir.exists():
            output_dir.mkdir(parents=True, exist_ok=True)

        pipeline_opts_with_ocr = PipelineOptions(ocr_options=ocr_options)

        doc_converter = DocumentConverter(
            allowed_formats=[
                InputFormat.PDF,
                InputFormat.IMAGE,
                InputFormat.DOCX,
                InputFormat.CSV,
                InputFormat.XLSX,
            ],
            format_options={
                InputFormat.PDF: FormatOption(
                    pipeline_cls=StandardPdfPipeline,
                    backend=PyPdfiumDocumentBackend,
                    pipeline_options=pipeline_opts_with_ocr,
                ),
                # --- THIS IS THE CORRECTED SECTION ---
                InputFormat.IMAGE: FormatOption(
                    pipeline_cls=SimplePipeline,
                    # THIS IS THE LINE THAT MUST BE ADDED
                    backend=DoclingParseV4DocumentBackend,
                    pipeline_options=pipeline_opts_with_ocr,
                ),
            },
        )

        _log.info(f"Starting conversion for {len(input_paths)} files...")
        conv_results = doc_converter.convert_all(input_paths)

        for res in conv_results:
            if res.errors:
                _log.error(f"Failed to convert {res.input.file.name}: {res.error}")
                continue

            base_name = res.input.file.stem
            print(f"Successfully processed: {res.input.file.name}")

            md_path = output_dir / f"{base_name}.md"
            with md_path.open("w", encoding="utf-8") as fp:
                fp.write(res.document.export_to_markdown())
            print(f"   -> Saved Markdown to: {md_path}")

            json_path = output_dir / f"{base_name}.json"
            with json_path.open("w", encoding="utf-8") as fp:
                fp.write(json.dumps(res.document.export_to_dict(), indent=2))
            print(f"   -> Saved JSON to: {json_path}")

            yaml_path = output_dir / f"{base_name}.yaml"
            with yaml_path.open("w", encoding="utf-8") as fp:
                fp.write(yaml.safe_dump(res.document.export_to_dict(), allow_unicode=True))
            print(f"   -> Saved YAML to: {yaml_path}")