infrasynth-backend-kit/infrasynth/files/processing.py

import io
import logging
from collections.abc import Callable
from typing import Any

from celery import shared_task
from django.utils import timezone

logger = logging.getLogger(__name__)


class PipelineStepRegistry:
    """Registry of processing steps so apps add steps without editing the kit.

    Register in ``apps.py:ready()``::

        from infrasynth.files.processing import PipelineStepRegistry

        def strip_exif(data: bytes, mime_type: str, params: dict):
            ...
            return data, mime_type

        PipelineStepRegistry.register("strip_exif", strip_exif)

    Or decorate: ``@pipeline_step("strip_exif")``. A registered step takes
    ``(data, mime_type, params)`` and returns ``(data, mime_type)``.
    """

    _steps: dict[str, Callable] = {}

    @classmethod
    def register(cls, name: str, func: Callable | None = None):
        if func is None:

            def decorator(inner: Callable) -> Callable:
                cls._steps[name] = inner
                return inner

            return decorator
        cls._steps[name] = func
        return func

    @classmethod
    def get(cls, name: str) -> Callable | None:
        return cls._steps.get(name)

    @classmethod
    def all(cls) -> dict[str, Callable]:
        return dict(cls._steps)


def pipeline_step(name: str):
    """Decorator form of :meth:`PipelineStepRegistry.register`."""
    return PipelineStepRegistry.register(name)


def get_pipeline_executor() -> Any:
    """Returns an executor instance (``INFRASYNTH_FILES["PIPELINE_EXECUTOR"]``)."""
    from django.utils.module_loading import import_string

    from infrasynth.shared.settings_utils import get_setting

    path = get_setting("INFRASYNTH_FILES", "PIPELINE_EXECUTOR", "infrasynth.files.processing.PipelineExecutor")
    return import_string(path)()


class PipelineExecutor:
    """Executes a processing pipeline over a stored file, step by step."""

    def execute(self, execution):
        from .models import PipelineExecution

        execution.status = PipelineExecution.Status.RUNNING
        execution.started_at = timezone.now()
        execution.save(update_fields=["status", "started_at"])

        try:
            data, mime_type = self._read_source(execution.file)
            pipeline = execution.pipeline
            steps = pipeline.steps if pipeline else []

            for step in steps:
                step_type = step.get("type")
                params = step.get("params", {})
                registered = PipelineStepRegistry.get(step_type)
                handler = registered or getattr(self, f"_step_{step_type}", None)
                if handler is None:
                    raise ValueError(f"Unknown pipeline step type: '{step_type}'")
                data, mime_type = handler(data, mime_type, params)

            output = self._store_output(execution, data, mime_type)
            execution.status = PipelineExecution.Status.COMPLETED
            execution.output_file = output
            execution.completed_at = timezone.now()
            execution.error = ""
            execution.save(update_fields=["status", "output_file", "completed_at", "error"])
        except Exception as exc:  # noqa: BLE001
            logger.exception("Pipeline execution %s failed", execution.pk)
            execution.status = PipelineExecution.Status.FAILED
            execution.completed_at = timezone.now()
            execution.error = str(exc)
            execution.save(update_fields=["status", "completed_at", "error"])

        from .signals import file_processed

        file_processed.send(
            sender=PipelineExecution,
            file_id=execution.file_id,
            pipeline_name=execution.pipeline.slug if execution.pipeline else "",
            output_file_id=execution.output_file_id,
            status=execution.status,
        )
        return execution

    def _read_source(self, stored_file):
        from .storage import get_storage_backend

        backend = get_storage_backend(stored_file.storage_backend)
        fh = backend.open(stored_file.storage_key, "rb")
        return fh.read(), stored_file.mime_type

    def _step_resize(self, data: bytes, mime_type: str, params: dict):
        from PIL import Image

        img = Image.open(io.BytesIO(data))
        width = int(params.get("width", 800))
        height = params.get("height")
        if height:
            img.thumbnail((width, int(height)))
        else:
            img.thumbnail((width, width))
        output = io.BytesIO()
        img.save(output, format=img.format or "PNG")
        return output.getvalue(), mime_type

    def _step_optimize(self, data: bytes, mime_type: str, params: dict):
        from PIL import Image

        quality = int(params.get("quality", 80))
        img = Image.open(io.BytesIO(data))
        fmt = img.format or "PNG"
        if fmt.upper() == "PNG":
            img = img.convert("P", palette=Image.Palette.ADAPTIVE, colors=256)  # type: ignore[assignment]
            output = io.BytesIO()
            img.save(output, format="PNG", optimize=True)
        else:
            output = io.BytesIO()
            img.save(output, format=fmt, quality=quality, optimize=True)
        return output.getvalue(), mime_type

    def _step_watermark(self, data: bytes, mime_type: str, params: dict):
        from PIL import Image, ImageDraw, ImageFont

        text = params.get("text", "Confidential")
        img = Image.open(io.BytesIO(data)).convert("RGBA")
        layer = Image.new("RGBA", img.size, (0, 0, 0, 0))
        draw = ImageDraw.Draw(layer)
        try:
            font = ImageFont.load_default(size=48)
        except TypeError:
            font = ImageFont.load_default()
        width, height = img.size
        draw.text((width // 4, height // 2), text, fill=(255, 255, 255, 120), font=font)
        out = Image.alpha_composite(img, layer)
        output = io.BytesIO()
        out.save(output, format="PNG")
        return output.getvalue(), "image/png"

    def _step_scan(self, data: bytes, mime_type: str, params: dict):
        from .scanner import get_scanner

        result = get_scanner().scan(data)
        if not result.clean:
            raise ValueError(f"Virus scan rejected the file: {result.threat}")
        logger.info("Virus scan passed for %d bytes (scanner=%s)", len(data), result.scanner)
        return data, mime_type

    def _store_output(self, execution, data: bytes, mime_type: str):
        from django.core.files.base import ContentFile

        from .services import FileService

        original = execution.file
        name = f"processed_{execution.pipeline.slug}_{original.original_filename}"
        content_file = ContentFile(data, name=name)
        content_file.content_type = mime_type  # type: ignore[attr-defined]

        service = FileService()
        output = service.upload(
            content_file,
            filename=name,
            user=original.uploaded_by,
            metadata={"source_file_id": original.id, "pipeline": execution.pipeline.slug},
        )
        return output


@shared_task(name="infrasynth.files.run_pipeline_execution", bind=True, max_retries=3)
def run_pipeline_execution(self, execution_id, tenant_id=None):
    """Celery task wrapper around the pipeline executor."""
    from infrasynth.tenancy.context import tenant_context
    from infrasynth.tenancy.models import Tenant

    from .models import PipelineExecution

    try:
        execution = PipelineExecution.all_objects.select_related("file", "pipeline").get(pk=execution_id)
    except PipelineExecution.DoesNotExist:
        logger.warning("Pipeline execution %s not found", execution_id)
        return None

    tenant = Tenant.objects.filter(pk=tenant_id).first() if tenant_id else execution.tenant
    with tenant_context(tenant):
        return get_pipeline_executor().execute(execution)