"""
PDF and Image handling utilities for OCR space.
Supports multi-page PDF rendering from paths or in-memory bytes with guaranteed file handle closure.
"""

import os
from typing import List, Union, Tuple, Optional
from PIL import Image
import io
import logging

logger = logging.getLogger("PDFUtils")


def is_pdf_file(file_input: Union[str, bytes, None]) -> bool:
    """Check if the input represents a PDF."""
    if isinstance(file_input, str):
        return file_input.lower().endswith(".pdf")
    if isinstance(file_input, bytes):
        return file_input.startswith(b"%PDF")
    return False


def convert_pdf_bytes_to_images(pdf_bytes: bytes, max_pages: int = 10, dpi: int = 200) -> List[Image.Image]:
    """Converts in-memory PDF bytes into a list of PIL Images (one per page)."""
    images = []

    # 1. Try pypdfium2 (in-memory, no temp file)
    try:
        import pypdfium2 as pdfium
        pdf = pdfium.PdfDocument(pdf_bytes)
        try:
            total_pages = min(len(pdf), max_pages)
            scale = dpi / 72.0
            for i in range(total_pages):
                page = pdf[i]
                image = page.render(scale=scale).to_pil()
                images.append(image.convert("RGB"))
            if images:
                return images
        finally:
            pdf.close()
    except Exception as e:
        logger.debug(f"pypdfium2 byte render fallback: {e}")

    # 2. Try fitz (PyMuPDF in-memory)
    try:
        import fitz
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        try:
            total_pages = min(len(doc), max_pages)
            for i in range(total_pages):
                page = doc[i]
                pix = page.get_pixmap(dpi=dpi)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                images.append(img)
            if images:
                return images
        finally:
            doc.close()
    except Exception as e:
        logger.debug(f"fitz byte render fallback: {e}")

    if not images:
        raise RuntimeError("Could not decode in-memory PDF with pypdfium2 or fitz.")
    return images


def convert_pdf_to_images(pdf_path: str, max_pages: int = 10, dpi: int = 200) -> List[Image.Image]:
    """Converts a PDF file from a path into a list of PIL Images."""
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"PDF not found: {pdf_path}")

    with open(pdf_path, "rb") as f:
        pdf_bytes = f.read()

    return convert_pdf_bytes_to_images(pdf_bytes, max_pages=max_pages, dpi=dpi)


def load_input_image_or_pdf(
    file_input: Union[str, bytes, Image.Image, None], page_index: int = 0
) -> Tuple[Optional[Image.Image], int, List[Image.Image]]:
    """
    Helper to process upload file (image, PDF path, or PDF bytes) and return:
      - (selected_image, total_pages, all_page_images)
    """
    if file_input is None:
        return None, 0, []

    if isinstance(file_input, Image.Image):
        return file_input.convert("RGB"), 1, [file_input.convert("RGB")]

    if isinstance(file_input, bytes):
        try:
            img = Image.open(io.BytesIO(file_input)).convert("RGB")
            return img, 1, [img]
        except Exception:
            pages = convert_pdf_bytes_to_images(file_input)
            total = len(pages)
            if total == 0:
                raise ValueError("PDF contains no readable pages.")
            idx = max(0, min(page_index, total - 1))
            return pages[idx], total, pages

    if isinstance(file_input, str):
        if not os.path.exists(file_input):
            raise FileNotFoundError(f"File not found: {file_input}")

        if is_pdf_file(file_input):
            pages = convert_pdf_to_images(file_input)
            total = len(pages)
            if total == 0:
                raise ValueError("PDF contains no readable pages.")
            idx = max(0, min(page_index, total - 1))
            return pages[idx], total, pages
        else:
            img = Image.open(file_input).convert("RGB")
            return img, 1, [img]

    return None, 0, []