"""Ingest layer: convert uploaded bytes (PDF/IMG) into a list of numpy images."""

from __future__ import annotations

import io
from dataclasses import dataclass
from typing import Any

import fitz  # PyMuPDF
import numpy as np
from PIL import Image

from ocr_sprint.schemas.document import SourceKind

# Generic alias used across the pipeline. We don't constrain dtype/shape because
# OpenCV operations accept multiple dtypes and numpy generics are still rough.
NDArrayU8 = np.ndarray[Any, Any]

PDF_MAGIC = b"%PDF-"
PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
JPEG_MAGIC = b"\xff\xd8\xff"
TIFF_MAGIC_LE = b"II*\x00"
TIFF_MAGIC_BE = b"MM\x00*"


@dataclass(frozen=True)
class IngestedPage:
    """One page worth of image data ready for preprocessing."""

    image: NDArrayU8  # HxWx3 BGR uint8 (OpenCV convention)
    page_index: int


def detect_source_kind(content: bytes) -> SourceKind:
    """Best-effort sniff of an uploaded payload."""
    if content.startswith(PDF_MAGIC):
        return SourceKind.PDF
    if content.startswith((PNG_MAGIC, JPEG_MAGIC, TIFF_MAGIC_LE, TIFF_MAGIC_BE)):
        return SourceKind.IMAGE
    return SourceKind.UNKNOWN


def _pil_to_bgr(img: Image.Image) -> NDArrayU8:
    """Convert PIL image to OpenCV BGR numpy array."""
    if img.mode != "RGB":
        img = img.convert("RGB")
    arr = np.asarray(img, dtype=np.uint8)
    # RGB to BGR
    return arr[:, :, ::-1].copy()


def ingest_pdf(content: bytes, target_dpi: int = 300) -> list[IngestedPage]:
    """Render every page of a PDF to a numpy image at the target DPI.

    Uses PyMuPDF (no poppler dependency). DPI is enforced via a transform matrix:
    fitz's default is 72 DPI, so the zoom factor is target_dpi / 72.
    """
    pages: list[IngestedPage] = []
    zoom = target_dpi / 72.0
    matrix = fitz.Matrix(zoom, zoom)
    with fitz.open(stream=content, filetype="pdf") as doc:
        for idx, page in enumerate(doc):
            pix = page.get_pixmap(matrix=matrix, alpha=False)
            img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
            pages.append(IngestedPage(image=_pil_to_bgr(img), page_index=idx))
    return pages


def ingest_image(content: bytes) -> list[IngestedPage]:
    """Decode a single image into a one-element page list."""
    img = Image.open(io.BytesIO(content))
    return [IngestedPage(image=_pil_to_bgr(img), page_index=0)]


def ingest(content: bytes, kind: SourceKind, target_dpi: int = 300) -> list[IngestedPage]:
    """Dispatch to the right loader based on declared source kind."""
    if kind == SourceKind.PDF:
        return ingest_pdf(content, target_dpi=target_dpi)
    if kind == SourceKind.IMAGE:
        return ingest_image(content)
    raise ValueError(f"Unsupported source kind: {kind}")
