LED-mapper/background_image_loader.py

from __future__ import annotations

from dataclasses import dataclass
import math
from pathlib import Path

from PySide6.QtCore import QByteArray, QBuffer, QIODevice, QSize, Qt
from PySide6.QtGui import QImage, QImageIOHandler, QImageReader

try:
    from .models import BackgroundImage
except ImportError:
    from models import BackgroundImage


MAX_BACKGROUND_IMAGE_SIDE = 8192
MAX_BACKGROUND_IMAGE_PIXELS = 32_000_000
MAX_UNSCALED_SOURCE_PIXELS = 100_000_000


@dataclass(frozen=True)
class BackgroundImageLoadResult:
    image: QImage
    error: str = ""

    @property
    def loaded(self) -> bool:
        return not self.image.isNull()


def supported_image_filter() -> str:
    formats = sorted(
        {
            bytes(image_format).decode("ascii", errors="ignore").lower()
            for image_format in QImageReader.supportedImageFormats()
            if image_format
        }
    )
    wildcards: list[str] = []
    for image_format in formats:
        if not image_format or image_format == "pdf":
            continue
        wildcards.extend((f"*.{image_format}", f"*.{image_format.upper()}"))
    patterns = " ".join(dict.fromkeys(wildcards))
    return f"Изображения ({patterns});;Все файлы (*.*)"


def resolve_background_path(background: BackgroundImage, document_path: Path | None = None) -> Path:
    path = Path(background.path)
    if not path.is_absolute() and document_path is not None:
        path = document_path.parent / path
    return path


def background_cache_key(background: BackgroundImage, document_path: Path | None = None) -> tuple[object, ...]:
    if background.has_embedded_image():
        image_bytes = background.embedded_image_bytes()
        return ("embedded", background.embedded_image_sha256 or hash(image_bytes), len(image_bytes))
    path = resolve_background_path(background, document_path)
    try:
        stat = path.stat()
    except OSError:
        return ("file", str(path), None, None)
    return ("file", str(path), stat.st_mtime_ns, stat.st_size)


def _bounded_image_size(size: QSize, max_side: int) -> QSize:
    if not size.isValid() or max_side <= 0:
        return size
    side_factor = min(1.0, max_side / max(size.width(), size.height()))
    pixel_factor = min(1.0, math.sqrt(MAX_BACKGROUND_IMAGE_PIXELS / (size.width() * size.height())))
    factor = min(side_factor, pixel_factor)
    if factor >= 1.0:
        return size
    return QSize(max(1, math.floor(size.width() * factor)), max(1, math.floor(size.height() * factor)))


def load_background_image(
    background: BackgroundImage,
    document_path: Path | None = None,
    *,
    max_side: int = MAX_BACKGROUND_IMAGE_SIDE,
) -> BackgroundImageLoadResult:
    buffer: QBuffer | None = None
    if background.has_embedded_image():
        image_bytes = background.embedded_image_bytes()
        if not image_bytes:
            return BackgroundImageLoadResult(QImage(), "Встроенные данные изображения повреждены или пусты.")
        buffer = QBuffer()
        buffer.setData(QByteArray(image_bytes))
        if not buffer.open(QIODevice.OpenModeFlag.ReadOnly):
            return BackgroundImageLoadResult(QImage(), "Не удалось открыть встроенные данные изображения.")
        reader = QImageReader(buffer)
        source_name = background.embedded_image_name or "встроенное изображение"
    else:
        path = resolve_background_path(background, document_path)
        if not path.is_file():
            return BackgroundImageLoadResult(QImage(), f"Файл не найден: {path}")
        reader = QImageReader(str(path))
        source_name = path.name

    previous_allocation_limit: int | None = None
    try:
        reader.setAutoTransform(True)
        reader.setDecideFormatFromContent(True)
        size = reader.size()
        bounded_size = _bounded_image_size(size, max_side)
        if bounded_size != size:
            if not reader.supportsOption(QImageIOHandler.ImageOption.ScaledSize):
                source_pixels = size.width() * size.height()
                if source_pixels > MAX_UNSCALED_SOURCE_PIXELS:
                    return BackgroundImageLoadResult(
                        QImage(),
                        f"Изображение «{source_name}» слишком большое для безопасного декодирования "
                        f"({size.width()}×{size.height()} пикселей).",
                    )
                current_allocation_limit = QImageReader.allocationLimit()
                required_megabytes = math.ceil(source_pixels * 8 / (1024 * 1024)) + 16
                if current_allocation_limit > 0 and required_megabytes > current_allocation_limit:
                    previous_allocation_limit = current_allocation_limit
                    QImageReader.setAllocationLimit(required_megabytes)
            reader.setScaledSize(bounded_size)
        image = reader.read()
        if image.isNull():
            details = reader.errorString().strip()
            message = f"Не удалось декодировать «{source_name}»"
            if details:
                message += f": {details}"
            return BackgroundImageLoadResult(QImage(), message + ".")

        bounded_size = _bounded_image_size(image.size(), max_side)
        if bounded_size != image.size():
            image = image.scaled(
                bounded_size,
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation,
            )
        if image.format() not in (QImage.Format.Format_ARGB32, QImage.Format.Format_ARGB32_Premultiplied):
            image = image.convertToFormat(QImage.Format.Format_ARGB32_Premultiplied)
        if image.isNull():
            return BackgroundImageLoadResult(QImage(), f"Не удалось преобразовать «{source_name}» для отображения.")
        return BackgroundImageLoadResult(image)
    finally:
        if previous_allocation_limit is not None:
            QImageReader.setAllocationLimit(previous_allocation_limit)
        if buffer is not None:
            buffer.close()