Corp/backend/tasks/media.py

import io
import logging
import shutil
import subprocess
import tempfile
from pathlib import Path

from django.core.files.base import ContentFile
from PIL import Image, ImageOps

from tasks.models import TaskAttachment


logger = logging.getLogger(__name__)

TASK_COVER_THUMBNAIL_SIZE = (640, 280)
IMAGE_EXTENSIONS = {
    ".avif",
    ".bmp",
    ".gif",
    ".heic",
    ".heif",
    ".jpeg",
    ".jpg",
    ".png",
    ".webp",
}
VIDEO_EXTENSIONS = {
    ".3gp",
    ".avi",
    ".m4v",
    ".mkv",
    ".mov",
    ".mp4",
    ".mpeg",
    ".mpg",
    ".ogv",
    ".webm",
}


def attachment_media_kind(attachment) -> str | None:
    mime_type = (attachment.mime_type or "").lower()
    suffix = Path(attachment.file_name or "").suffix.lower()
    if mime_type.startswith("image/") or suffix in IMAGE_EXTENSIONS:
        return "image"
    if mime_type.startswith("video/") or suffix in VIDEO_EXTENSIONS:
        return "video"
    return None


def task_attachment_media_kind(attachment: TaskAttachment) -> str | None:
    return attachment_media_kind(attachment)


def _render_image_thumbnail(attachment: TaskAttachment) -> bytes:
    with attachment.file.open("rb") as source:
        image = Image.open(source)
        image.seek(0)
        image = ImageOps.exif_transpose(image)
        if image.mode not in ("RGB", "RGBA"):
            image = image.convert("RGBA" if "transparency" in image.info else "RGB")
        if image.mode == "RGBA":
            background = Image.new("RGB", image.size, (20, 27, 36))
            background.paste(image, mask=image.getchannel("A"))
            image = background
        image = ImageOps.fit(
            image,
            TASK_COVER_THUMBNAIL_SIZE,
            method=Image.Resampling.LANCZOS,
        )
        output = io.BytesIO()
        image.save(output, format="WEBP", quality=84, method=5)
        return output.getvalue()


def _render_video_thumbnail(attachment: TaskAttachment) -> bytes:
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is not installed")

    input_suffix = Path(attachment.file_name or "").suffix or ".video"
    input_path = ""
    temporary_input = False
    output_path = ""
    try:
        try:
            input_path = attachment.file.path
        except (AttributeError, NotImplementedError):
            with tempfile.NamedTemporaryFile(
                suffix=input_suffix, delete=False
            ) as source_file:
                input_path = source_file.name
                temporary_input = True
                with attachment.file.open("rb") as source:
                    for chunk in iter(lambda: source.read(1024 * 1024), b""):
                        source_file.write(chunk)
        with tempfile.NamedTemporaryFile(suffix=".webp", delete=False) as output_file:
            output_path = output_file.name

        filter_value = "scale=640:280:force_original_aspect_ratio=increase,crop=640:280"
        command_prefix = [ffmpeg, "-hide_banner", "-loglevel", "error"]
        command_suffix = [
            "-map",
            "0:v:0",
            "-frames:v",
            "1",
            "-vf",
            filter_value,
            "-an",
            "-sn",
            "-dn",
            "-c:v",
            "libwebp",
            "-quality",
            "84",
            "-y",
            output_path,
        ]
        commands = [
            # Fast input seek works for ordinary MP4/WebM files.
            [*command_prefix, "-ss", "0.5", "-i", input_path, *command_suffix],
            # Accurate output seek handles fragmented MP4/MOV screen recordings
            # whose first keyframe cannot be reached by input seeking.
            [*command_prefix, "-i", input_path, "-ss", "0.5", *command_suffix],
            # Some recordings start on a non-zero timeline. In that case use the
            # first frame FFmpeg can decode instead of requiring a timestamp.
            [*command_prefix, "-i", input_path, *command_suffix],
        ]
        last_error = "FFmpeg did not produce a video thumbnail."
        for command in commands:
            try:
                result = subprocess.run(
                    command,
                    capture_output=True,
                    check=False,
                    timeout=30,
                )
            except subprocess.TimeoutExpired as error:
                last_error = str(error)
                continue
            if result.returncode == 0 and Path(output_path).stat().st_size > 0:
                break
            last_error = result.stderr.decode("utf-8", errors="replace")
        else:
            raise RuntimeError(last_error)
        return Path(output_path).read_bytes()
    finally:
        if temporary_input and input_path:
            Path(input_path).unlink(missing_ok=True)
        if output_path:
            Path(output_path).unlink(missing_ok=True)


def render_attachment_thumbnail(attachment) -> bytes:
    media_kind = attachment_media_kind(attachment)
    if media_kind == "image":
        return _render_image_thumbnail(attachment)
    if media_kind == "video":
        return _render_video_thumbnail(attachment)
    raise ValueError("Attachment is not supported media")


def ensure_task_attachment_thumbnail(attachment: TaskAttachment) -> bool:
    if attachment.thumbnail:
        return True
    if not task_attachment_media_kind(attachment):
        return False

    try:
        thumbnail = render_attachment_thumbnail(attachment)
        attachment.thumbnail.save(
            f"task-{attachment.task_id}-attachment-{attachment.id}.webp",
            ContentFile(thumbnail),
            save=True,
        )
        return True
    except Exception as error:
        logger.exception(
            "Task attachment thumbnail generation failed",
            extra={
                "event": "task.attachment.thumbnail_generation_failed",
                "task_id": attachment.task_id,
                "attachment_id": attachment.id,
                "error_type": type(error).__name__,
            },
        )
        return False