import inspect
import json
import logging
from typing import Optional

from PIL.Image import Image
from pydantic import BaseModel

from docling.backend.abstract_backend import PaginatedDocumentBackend
from docling.backend.docling_parse_backend import ThreadedDoclingParseDocumentBackend
from docling.backend.pdf_backend import PdfDocumentBackend
from docling.datamodel.base_models import (
    ConversionStatus,
    DoclingComponentType,
    ErrorItem,
    FailureCategory,
    VlmStopReason,
)
from docling.datamodel.document import InputDocument
from docling.datamodel.extraction import (
    ExtractedPageData,
    ExtractionResult,
    ExtractionTemplateType,
)
from docling.datamodel.pipeline_options import (
    PipelineOptions,
    VlmExtractionPipelineOptions,
)
from docling.datamodel.settings import settings
from docling.models.extraction.transformers_extraction_model import (
    TransformersExtractionModel,
)
from docling.pipeline.base_extraction_pipeline import BaseExtractionPipeline
from docling.utils.accelerator_utils import decide_device

_log = logging.getLogger(__name__)


def _raise_if_unsupported_threaded_backend(
    backend: PaginatedDocumentBackend, pipeline_name: str
) -> None:
    if isinstance(backend, ThreadedDoclingParseDocumentBackend):
        raise RuntimeError(
            f"{pipeline_name} does not support ThreadedDoclingParseDocumentBackend yet. "
            "It still requires ordered/random page access via load_page() and cannot "
            "consume iterator-only or out-of-order page delivery. Use StandardPdfPipeline instead."
        )


class ExtractionVlmPipeline(BaseExtractionPipeline):
    def __init__(self, pipeline_options: VlmExtractionPipelineOptions):
        super().__init__(pipeline_options)

        self.accelerator_options = pipeline_options.accelerator_options
        self.pipeline_options: VlmExtractionPipelineOptions

        self.vlm_model = TransformersExtractionModel(
            enabled=True,
            artifacts_path=self.artifacts_path,
            accelerator_options=self.accelerator_options,
            vlm_options=pipeline_options.vlm_options,
            prompt_style=pipeline_options.extraction_prompt_style,
        )

    def _extract_data(
        self,
        ext_res: ExtractionResult,
        template: Optional[ExtractionTemplateType] = None,
    ) -> ExtractionResult:
        """Extract data using the VLM model."""
        backend = ext_res.input._backend
        if isinstance(backend, PdfDocumentBackend):
            _raise_if_unsupported_threaded_backend(backend, self.__class__.__name__)

        try:
            # Get images from input document using the backend
            images = self._get_images_from_input(ext_res.input)
            if not images:
                ext_res.status = ConversionStatus.FAILURE
                ext_res.errors.append(
                    ErrorItem(
                        component_type=DoclingComponentType.PIPELINE,
                        module_name=self.__class__.__name__,
                        error_message="No images found in document",
                        category=FailureCategory.BACKEND_FAILURE,
                    )
                )
                return ext_res

            # Use provided template or default prompt
            if template is not None:
                prompt = self._serialize_template(template)
            else:
                prompt = "Extract all text and structured information from this document. Return as JSON."

            # Process all images with VLM model
            start_page, end_page = ext_res.input.limits.page_range
            for i, image in enumerate(images):
                # Calculate the actual page number based on the filtered range
                page_number = start_page + i
                try:
                    predictions = list(self.vlm_model.process_images([image], prompt))

                    if predictions:
                        # Parse the extracted text as JSON if possible, otherwise use as-is
                        extracted_text = predictions[0].text
                        extracted_data = None
                        vlm_stop_reason: VlmStopReason = predictions[0].stop_reason
                        if (
                            vlm_stop_reason == VlmStopReason.LENGTH
                            or vlm_stop_reason == VlmStopReason.STOP_SEQUENCE
                        ):
                            ext_res.status = ConversionStatus.PARTIAL_SUCCESS

                        try:
                            extracted_data = json.loads(extracted_text)
                        except (json.JSONDecodeError, ValueError):
                            # If not valid JSON, keep extracted_data as None
                            pass

                        # Create page data with proper structure
                        page_data = ExtractedPageData(
                            page_no=page_number,
                            extracted_data=extracted_data,
                            raw_text=extracted_text,  # Always populate raw_text
                        )
                        ext_res.pages.append(page_data)
                    else:
                        # Add error page data
                        page_data = ExtractedPageData(
                            page_no=page_number,
                            extracted_data=None,
                            errors=["No extraction result from VLM model"],
                        )
                        ext_res.pages.append(page_data)

                except Exception as e:
                    _log.error(f"Error processing page {page_number}: {e}")
                    page_data = ExtractedPageData(
                        page_no=page_number, extracted_data=None, errors=[str(e)]
                    )
                    ext_res.pages.append(page_data)

        except Exception as e:
            _log.error(f"Error during extraction: {e}")
            ext_res.errors.append(
                ErrorItem(
                    component_type=DoclingComponentType.PIPELINE,
                    module_name=self.__class__.__name__,
                    error_message=str(e),
                    category=FailureCategory.UNKNOWN,
                )
            )

        return ext_res

    def _determine_status(self, ext_res: ExtractionResult) -> ConversionStatus:
        """Determine the status based on extraction results."""
        if ext_res.pages and not any(page.errors for page in ext_res.pages):
            return (
                ConversionStatus.PARTIAL_SUCCESS
                if ext_res.status == ConversionStatus.PARTIAL_SUCCESS
                else ConversionStatus.SUCCESS
            )
        else:
            return ConversionStatus.FAILURE

    def _get_images_from_input(self, input_doc: InputDocument) -> list[Image]:
        """Extract images from input document using the backend."""
        images = []

        try:
            backend = input_doc._backend

            assert isinstance(backend, PdfDocumentBackend)
            # Use the backend's pagination interface
            page_count = backend.page_count()

            # Respect page range limits, following the same pattern as PaginatedPipeline
            start_page, end_page = input_doc.limits.page_range
            _log.info(
                f"Processing pages {start_page}-{end_page} of {page_count} total pages for extraction"
            )

            for page_num in range(page_count):
                # Only process pages within the specified range (0-based indexing)
                if start_page - 1 <= page_num <= end_page - 1:
                    try:
                        page_backend = backend.load_page(page_num)
                        if page_backend.is_valid():
                            # Get page image at a reasonable scale
                            page_image = page_backend.get_page_image(
                                scale=self.pipeline_options.vlm_options.scale
                            )
                            images.append(page_image)
                        else:
                            _log.warning(f"Page {page_num + 1} backend is not valid")
                    except Exception as e:
                        _log.error(f"Error loading page {page_num + 1}: {e}")

        except Exception as e:
            _log.error(f"Error getting images from input document: {e}")

        return images

    def _serialize_template(self, template: ExtractionTemplateType) -> str:
        """Serialize template to string based on its type."""
        if isinstance(template, str):
            return template
        elif isinstance(template, dict):
            return json.dumps(template, indent=2)
        elif isinstance(template, BaseModel):
            return template.model_dump_json(indent=2)
        elif inspect.isclass(template) and issubclass(template, BaseModel):
            from polyfactory.factories.pydantic_factory import ModelFactory

            class ExtractionTemplateFactory(ModelFactory[template]):  # type: ignore
                __use_examples__ = True  # prefer Field(examples=...) when present
                __use_defaults__ = True  # use field defaults instead of random values
                __check_model__ = (
                    True  # setting the value to avoid deprecation warnings
                )

            return ExtractionTemplateFactory.build().model_dump_json(indent=2)  # type: ignore
        else:
            raise ValueError(f"Unsupported template type: {type(template)}")

    @classmethod
    def get_default_options(cls) -> PipelineOptions:
        return VlmExtractionPipelineOptions()
