warpbox-dev/backend/libs/jobs/thumbnails.go

package jobs

import (
	"bytes"
	"context"
	"image"
	_ "image/gif"
	"image/jpeg"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"strings"
	"time"

	_ "golang.org/x/image/webp"
	"warpbox.dev/backend/libs/config"
	"warpbox.dev/backend/libs/services"
)

type thumbnailJobResult struct {
	Scanned   int
	Generated int
	Failed    int
}

func GenerateThumbnailsForBoxAsync(uploadService *services.UploadService, logger *slog.Logger, boxID string) {
	go func() {
		box, err := uploadService.GetBox(boxID)
		if err != nil {
			logger.Warn("thumbnail box lookup failed", "source", "thumbnail", "severity", "warn", "code", 4204, "box_id", boxID, "error", err.Error())
			return
		}

		result, err := generateMissingThumbnailsForBox(uploadService, logger, box)
		if err != nil {
			logger.Warn("thumbnail one-shot job failed", "source", "thumbnail", "severity", "warn", "code", 4205, "box_id", boxID, "error", err.Error())
			return
		}
		if result.Generated > 0 || result.Failed > 0 {
			logger.Info("thumbnail one-shot job complete", "source", "thumbnail", "severity", "user_activity", "code", 2205, "box_id", boxID, "generated", result.Generated, "failed", result.Failed)
		}
	}()
}

func newThumbnailsJob(cfg config.Config, logger *slog.Logger, uploadService *services.UploadService) job {
	return job{
		name:     "thumbnail",
		enabled:  cfg.ThumbnailEnabled,
		interval: cfg.ThumbnailEvery,
		run: func() {
			result, err := generateMissingThumbnails(uploadService, logger)
			if err != nil {
				logger.Warn("thumbnail job failed", "source", "thumbnail", "severity", "warn", "code", 4203, "error", err.Error())
				return
			}
			if result.Generated > 0 || result.Failed > 0 {
				logger.Info("thumbnail job complete", "source", "thumbnail", "severity", "user_activity", "code", 2204, "generated", result.Generated, "failed", result.Failed)
			}
		},
	}
}

func generateMissingThumbnails(uploadService *services.UploadService, logger *slog.Logger) (thumbnailJobResult, error) {
	boxes, err := uploadService.ListBoxes(0)
	if err != nil {
		return thumbnailJobResult{}, err
	}

	var result thumbnailJobResult
	now := time.Now().UTC()
	for _, box := range boxes {
		if !box.ExpiresAt.After(now) {
			continue
		}

		boxResult, err := generateMissingThumbnailsForBox(uploadService, logger, box)
		result.Scanned += boxResult.Scanned
		result.Generated += boxResult.Generated
		result.Failed += boxResult.Failed
		if err != nil {
			return result, err
		}
	}

	return result, nil
}

func generateMissingThumbnailsForBox(uploadService *services.UploadService, logger *slog.Logger, box services.Box) (thumbnailJobResult, error) {
	var result thumbnailJobResult
	if !box.ExpiresAt.After(time.Now().UTC()) {
		return result, nil
	}

	changed := false
	for i := range box.Files {
		file := &box.Files[i]
		if file.Thumbnail != "" || !needsThumbnail(*file) {
			continue
		}
		result.Scanned++

		thumbnail, err := generateThumbnail(uploadService, box, *file)
		if err != nil {
			logger.Warn("thumbnail generation failed", "source", "thumbnail", "severity", "warn", "code", 4101, "file_id", file.ID, "error", err.Error())
			result.Failed++
			continue
		}
		if thumbnail == "" {
			result.Failed++
			continue
		}

		file.Thumbnail = thumbnail
		changed = true
		result.Generated++
	}

	if changed {
		if err := uploadService.SaveBox(box); err != nil {
			return result, err
		}
	}
	return result, nil
}

func needsThumbnail(file services.File) bool {
	return file.PreviewKind == "image" || file.PreviewKind == "video"
}

func generateThumbnail(uploadService *services.UploadService, box services.Box, file services.File) (string, error) {
	thumbnailName := "@thumb@" + file.ID + ".jpg"
	object, err := uploadService.OpenFileObject(context.Background(), box, file)
	if err != nil {
		return "", err
	}
	defer object.Body.Close()

	switch {
	case strings.HasPrefix(file.ContentType, "image/"):
		data, err := createImageThumbnail(object.Body)
		if err != nil {
			return "", err
		}
		_, err = uploadService.PutThumbnailObject(context.Background(), box, thumbnailName, bytes.NewReader(data), int64(len(data)), "image/jpeg")
		return thumbnailName, err
	case strings.HasPrefix(file.ContentType, "video/"):
		data, err := createVideoThumbnail(object.Body)
		if err != nil {
			return "", err
		}
		_, err = uploadService.PutThumbnailObject(context.Background(), box, thumbnailName, bytes.NewReader(data), int64(len(data)), "image/jpeg")
		return thumbnailName, err
	default:
		return "", nil
	}
}

func createImageThumbnail(source io.Reader) ([]byte, error) {
	img, _, err := image.Decode(source)
	if err != nil {
		return nil, err
	}

	thumb := resizeNearest(img, 360, 240)
	var target bytes.Buffer
	err = jpeg.Encode(&target, thumb, &jpeg.Options{Quality: 82})
	if err != nil {
		return nil, err
	}
	return target.Bytes(), nil
}

func createVideoThumbnail(source io.Reader) ([]byte, error) {
	sourceFile, err := os.CreateTemp("", "warpbox-video-*")
	if err != nil {
		return nil, err
	}
	defer os.Remove(sourceFile.Name())
	if _, err := io.Copy(sourceFile, source); err != nil {
		sourceFile.Close()
		return nil, err
	}
	if err := sourceFile.Close(); err != nil {
		return nil, err
	}
	targetFile, err := os.CreateTemp("", "warpbox-thumb-*.jpg")
	if err != nil {
		return nil, err
	}
	targetPath := targetFile.Name()
	targetFile.Close()
	defer os.Remove(targetPath)
	if err := exec.Command("ffmpeg", "-y", "-loglevel", "error", "-ss", "00:00:01", "-i", sourceFile.Name(), "-frames:v", "1", "-vf", "scale=360:-1", targetPath).Run(); err != nil {
		return nil, err
	}
	return os.ReadFile(targetPath)
}

func resizeNearest(src image.Image, maxWidth, maxHeight int) *image.RGBA {
	bounds := src.Bounds()
	width := bounds.Dx()
	height := bounds.Dy()
	if width <= 0 || height <= 0 {
		return image.NewRGBA(image.Rect(0, 0, 1, 1))
	}

	scale := min(float64(maxWidth)/float64(width), float64(maxHeight)/float64(height))
	if scale > 1 {
		scale = 1
	}
	targetWidth := max(1, int(float64(width)*scale))
	targetHeight := max(1, int(float64(height)*scale))
	dst := image.NewRGBA(image.Rect(0, 0, targetWidth, targetHeight))

	for y := 0; y < targetHeight; y++ {
		for x := 0; x < targetWidth; x++ {
			srcX := bounds.Min.X + int(float64(x)/scale)
			srcY := bounds.Min.Y + int(float64(y)/scale)
			dst.Set(x, y, src.At(srcX, srcY))
		}
	}

	return dst
}