from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np
from PIL import Image

IMAGE_EXTENSIONS = {
    ".jpg",
    ".jpeg",
    ".png",
    ".bmp",
    ".tiff",
    ".tif",
    ".webp",
    ".ppm",
    ".pgm",
    ".pbm",
}


def load_image(path: str | Path) -> np.ndarray:
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Image not found: {path}")

    ext = path.suffix.lower()
    if ext == ".pdf":
        return _load_pdf(path)

    img = cv2.imread(str(path))
    if img is not None:
        return img

    pil_img = Image.open(path)
    arr = np.array(pil_img)
    if arr.ndim == 2:
        return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
    if arr.shape[-1] == 4:
        return cv2.cvtColor(arr, cv2.COLOR_RGBA2BGR)
    return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)


def load_pdf_pages(path: Path, scale: float = 2.0) -> list[np.ndarray]:
    import pypdfium2 as pdfium

    pdf = pdfium.PdfDocument(str(path))
    pages = []
    for page in pdf:
        bitmap = page.render(scale=scale)
        pil = bitmap.to_pil()
        arr = np.array(pil)
        pages.append(cv2.cvtColor(arr, cv2.COLOR_RGB2BGR))
    return pages


def _load_pdf(path: Path) -> np.ndarray:
    pages = load_pdf_pages(path)
    if pages:
        return pages[0]
    raise ValueError(f"PDF has no pages: {path}")