from __future__ import annotations

import math
from pathlib import Path
from types import ModuleType
from typing import Any

from .errors import InferenceError


class PDFPageRenderer:
    """Bounded, page-at-a-time PDFium renderer."""

    def __init__(
        self,
        pdf_path: Path,
        work_dir: Path,
        *,
        dpi: int,
        start_page: int,
        max_pages: int,
        max_page_pixels: int,
        max_total_pixels: int,
        max_page_rendered_bytes: int,
        max_rendered_bytes: int,
    ) -> None:
        self.pdf_path = pdf_path
        self.work_dir = work_dir
        self.dpi = dpi
        self.start_page = start_page
        self.max_pages = max_pages
        self.max_page_pixels = max_page_pixels
        self.max_total_pixels = max_total_pixels
        self.max_page_rendered_bytes = max_page_rendered_bytes
        self.max_rendered_bytes = max_rendered_bytes
        self._pdfium: ModuleType | None = None
        self._document: Any = None
        self.total_pages = 0
        self._rendered_pixels = 0
        self._rendered_bytes = 0

    def __enter__(self) -> PDFPageRenderer:
        try:
            import pypdfium2 as pdfium
        except ImportError as exc:
            raise InferenceError("pypdfium2 is missing; reinstall unlimited-ocr-rdna4") from exc

        self._pdfium = pdfium
        try:
            self._document = pdfium.PdfDocument(str(self.pdf_path))
            self.total_pages = len(self._document)
        except Exception as exc:
            raise InferenceError(f"could not open PDF {self.pdf_path}: {exc}") from exc

        if self.total_pages < 1:
            self.close()
            raise InferenceError(f"PDF has no pages: {self.pdf_path}")
        if self.start_page > self.total_pages:
            self.close()
            raise InferenceError(f"start page {self.start_page} exceeds PDF page count {self.total_pages}")
        return self

    def __exit__(self, _exc_type: object, _exc: object, _traceback: object) -> None:
        self.close()

    @property
    def page_numbers(self) -> range:
        last_page = min(self.total_pages, self.start_page - 1 + self.max_pages)
        return range(self.start_page, last_page + 1)

    def render(self, page_number: int) -> Path:
        if self._document is None:
            raise InferenceError("PDF renderer is not open")
        page = None
        bitmap = None
        image = None
        output = self.work_dir / f"page-{page_number:06d}.png"
        try:
            page = self._document[page_number - 1]
            width_points, height_points = page.get_size()
            scale = self.dpi / 72
            width_pixels = math.ceil(width_points * scale)
            height_pixels = math.ceil(height_points * scale)
            pixels = width_pixels * height_pixels
            if pixels <= 0 or pixels > self.max_page_pixels:
                raise InferenceError(
                    f"PDF page {page_number} would render {pixels:,} pixels; limit is {self.max_page_pixels:,}"
                )
            if self._rendered_pixels + pixels > self.max_total_pixels:
                raise InferenceError(
                    f"PDF rendered-pixel total would exceed {self.max_total_pixels:,} on page {page_number}"
                )

            bitmap = page.render(scale=scale)
            image = bitmap.to_pil()
            image.save(output, format="PNG", optimize=True)
            rendered_bytes = output.stat().st_size
            if rendered_bytes > self.max_page_rendered_bytes:
                output.unlink(missing_ok=True)
                raise InferenceError(
                    f"PDF page {page_number} rendered to {rendered_bytes:,} bytes; "
                    f"limit is {self.max_page_rendered_bytes:,}"
                )
            if self._rendered_bytes + rendered_bytes > self.max_rendered_bytes:
                output.unlink(missing_ok=True)
                raise InferenceError(
                    f"PDF rendered-byte total would exceed {self.max_rendered_bytes:,} on page {page_number}"
                )
            self._rendered_pixels += pixels
            self._rendered_bytes += rendered_bytes
            return output
        except InferenceError:
            raise
        except Exception as exc:
            output.unlink(missing_ok=True)
            raise InferenceError(f"could not render PDF page {page_number}: {exc}") from exc
        finally:
            if image is not None:
                image.close()
            if bitmap is not None:
                bitmap.close()
            if page is not None:
                page.close()

    def close(self) -> None:
        if self._document is not None:
            self._document.close()
            self._document = None