wijtransferen/backend/internal/media/processor.go

package media

import (
	"bytes"
	"context"
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	_ "image/png"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"sync"

	"github.com/example/sndit/backend/internal/storage"
	"github.com/google/uuid"
)

type Processor struct {
	repository *Repository
	storage    storage.Storage
	jobs       chan uuid.UUID
	stop       chan struct{}
	waitGroup  sync.WaitGroup
}

func NewProcessor(repository *Repository, objectStorage storage.Storage, workers int) *Processor {
	if workers < 1 {
		workers = 1
	}
	processor := &Processor{
		repository: repository,
		storage:    objectStorage,
		jobs:       make(chan uuid.UUID, 256),
		stop:       make(chan struct{}),
	}
	for index := 0; index < workers; index++ {
		processor.waitGroup.Add(1)
		go processor.worker()
	}
	return processor
}

func (p *Processor) Enqueue(mediaID uuid.UUID) {
	select {
	case p.jobs <- mediaID:
	case <-p.stop:
	}
}

func (p *Processor) Close() {
	close(p.stop)
	p.waitGroup.Wait()
}

func (p *Processor) worker() {
	defer p.waitGroup.Done()
	for {
		select {
		case mediaID := <-p.jobs:
			p.process(mediaID)
		case <-p.stop:
			return
		}
	}
}

func (p *Processor) process(mediaID uuid.UUID) {
	ctx := context.Background()
	item, err := p.repository.GetByID(ctx, mediaID)
	if err != nil {
		return
	}

	if item.ExternalURL != "" {
		_ = p.repository.MarkReady(ctx, mediaID, item.StorageKey, item.StorageKey, item.Width, item.Height)
		return
	}

	if IsImage(item) {
		p.processImage(ctx, item)
	} else if IsVideo(item) {
		p.processVideo(ctx, item)
	} else {
		_ = p.repository.MarkReady(ctx, mediaID, item.StorageKey, item.StorageKey, item.Width, item.Height)
	}
}

func (p *Processor) processImage(ctx context.Context, item Record) {
	object, err := p.storage.Get(ctx, item.StorageKey)
	if err != nil {
		_ = p.repository.MarkFailed(ctx, item.ID, err.Error())
		return
	}
	defer object.Close()

	decoded, _, err := image.Decode(io.LimitReader(object, 100<<20))
	if err != nil {
		_ = p.repository.MarkReady(ctx, item.ID, item.StorageKey, item.StorageKey, 0, 0)
		return
	}

	width := decoded.Bounds().Dx()
	height := decoded.Bounds().Dy()
	previewKey := variantKey(item, "preview.jpg")
	thumbnailKey := variantKey(item, "thumbnail.jpg")

	preview, err := encodeJPEG(resize(decoded, 2400), 88)
	if err != nil {
		_ = p.repository.MarkFailed(ctx, item.ID, err.Error())
		return
	}
	thumbnail, err := encodeJPEG(resize(decoded, 640), 84)
	if err != nil {
		_ = p.repository.MarkFailed(ctx, item.ID, err.Error())
		return
	}
	if err := p.storage.Put(ctx, previewKey, bytes.NewReader(preview), int64(len(preview)), "image/jpeg"); err != nil {
		_ = p.repository.MarkFailed(ctx, item.ID, err.Error())
		return
	}
	if err := p.storage.Put(ctx, thumbnailKey, bytes.NewReader(thumbnail), int64(len(thumbnail)), "image/jpeg"); err != nil {
		_ = p.repository.MarkFailed(ctx, item.ID, err.Error())
		return
	}
	_ = p.repository.MarkReady(ctx, item.ID, previewKey, thumbnailKey, width, height)
}

func (p *Processor) processVideo(ctx context.Context, item Record) {
	thumbnailKey := variantKey(item, "thumbnail.jpg")
	previewKey := variantKey(item, "preview.mp4")

	if err := p.generateVideoThumbnail(ctx, item, thumbnailKey); err != nil {
		p.log("video thumbnail failed for %s: %s — using fallback", item.ID, err)
		p.generateFallbackThumbnail(ctx, item, thumbnailKey)
	}
	if err := p.generateVideoPreview(ctx, item, previewKey); err != nil {
		p.log("video preview failed for %s: %s — using original", item.ID, err)
		_ = p.repository.MarkReady(ctx, item.ID, item.StorageKey, thumbnailKey, 0, 0)
		return
	}
	_ = p.repository.MarkReady(ctx, item.ID, previewKey, thumbnailKey, 0, 0)
}

func (p *Processor) generateVideoThumbnail(ctx context.Context, item Record, key string) error {
	tmp := filepath.Join(os.TempDir(), fmt.Sprintf("sndit-%s", uuid.New()))
	defer os.Remove(tmp)

	if err := p.downloadToFile(ctx, item.StorageKey, tmp); err != nil {
		return err
	}

	var stderr bytes.Buffer
	cmd := exec.CommandContext(ctx, "ffmpeg",
		"-ss", "0",
		"-i", tmp,
		"-vframes", "1",
		"-q:v", "6",
		"-f", "image2pipe",
		"-c:v", "mjpeg",
		"pipe:1",
	)
	cmd.Stderr = &stderr

	output, err := cmd.Output()
	if err != nil {
		return fmt.Errorf("ffmpeg thumbnail: %w, stderr: %s", err, stderr.String())
	}

	return p.storage.Put(ctx, key, bytes.NewReader(output), int64(len(output)), "image/jpeg")
}

func (p *Processor) generateVideoPreview(ctx context.Context, item Record, key string) error {
	tmpInput := filepath.Join(os.TempDir(), fmt.Sprintf("sndit-in-%s", uuid.New()))
	tmpOutput := filepath.Join(os.TempDir(), fmt.Sprintf("sndit-out-%s.mp4", uuid.New()))
	defer os.Remove(tmpInput)
	defer os.Remove(tmpOutput)

	if err := p.downloadToFile(ctx, item.StorageKey, tmpInput); err != nil {
		return err
	}

	var stderr bytes.Buffer
	cmd := exec.CommandContext(ctx, "ffmpeg",
		"-i", tmpInput,
		"-vf", "scale='min(720,iw)':-1:force_original_aspect_ratio=decrease,pad='ceil(iw/2)*2':'ceil(ih/2)*2':-1:-1:black",
		"-c:v", "libx264",
		"-preset", "fast",
		"-crf", "28",
		"-movflags", "+faststart",
		"-c:a", "aac",
		"-b:a", "64k",
		"-y", tmpOutput,
	)
	cmd.Stderr = &stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("ffmpeg preview: %w, stderr: %s", err, stderr.String())
	}

	data, err := os.ReadFile(tmpOutput)
	if err != nil {
		return fmt.Errorf("read output file: %w", err)
	}

	return p.storage.Put(ctx, key, bytes.NewReader(data), int64(len(data)), "video/mp4")
}

func (p *Processor) downloadToFile(ctx context.Context, storageKey, tmp string) error {
	object, err := p.storage.Get(ctx, storageKey)
	if err != nil {
		return fmt.Errorf("get original: %w", err)
	}
	defer object.Close()

	f, err := os.Create(tmp)
	if err != nil {
		return fmt.Errorf("create temp file: %w", err)
	}

	if _, err := io.Copy(f, object); err != nil {
		f.Close()
		os.Remove(tmp)
		return fmt.Errorf("copy to temp file: %w", err)
	}
	if err := f.Sync(); err != nil {
		f.Close()
		os.Remove(tmp)
		return fmt.Errorf("sync temp file: %w", err)
	}
	f.Close()
	return nil
}

func variantKey(item Record, filename string) string {
	return fmt.Sprintf("galleries/%s/%s/%s", item.GalleryID, item.ID, filename)
}

func resize(source image.Image, maxSide int) image.Image {
	bounds := source.Bounds()
	width, height := bounds.Dx(), bounds.Dy()
	if width <= maxSide && height <= maxSide {
		return source
	}

	scale := float64(maxSide) / float64(width)
	if height > width {
		scale = float64(maxSide) / float64(height)
	}
	newWidth := int(float64(width) * scale)
	newHeight := int(float64(height) * scale)
	destination := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
	for y := 0; y < newHeight; y++ {
		for x := 0; x < newWidth; x++ {
			sourceX := bounds.Min.X + x*width/newWidth
			sourceY := bounds.Min.Y + y*height/newHeight
			destination.Set(x, y, source.At(sourceX, sourceY))
		}
	}
	return destination
}

func encodeJPEG(source image.Image, quality int) ([]byte, error) {
	var output bytes.Buffer
	if err := jpeg.Encode(&output, source, &jpeg.Options{Quality: quality}); err != nil {
		return nil, fmt.Errorf("encode preview: %w", err)
	}
	return output.Bytes(), nil
}

func (p *Processor) log(format string, args ...any) {
	log.Printf("[media-processor] "+format, args...)
}

func (p *Processor) generateFallbackThumbnail(ctx context.Context, item Record, key string) {
	img := image.NewRGBA(image.Rect(0, 0, 640, 480))
	bg := color.RGBA{20, 20, 20, 255}
	draw.Draw(img, img.Bounds(), &image.Uniform{bg}, image.Point{}, draw.Src)

	data, err := encodeJPEG(img, 60)
	if err != nil {
		p.log("fallback thumbnail encode failed: %s", err)
		_ = p.repository.MarkReady(ctx, item.ID, item.StorageKey, item.StorageKey, 0, 0)
		return
	}
	if err := p.storage.Put(ctx, key, bytes.NewReader(data), int64(len(data)), "image/jpeg"); err != nil {
		p.log("fallback thumbnail upload failed: %s", err)
		_ = p.repository.MarkReady(ctx, item.ID, item.StorageKey, item.StorageKey, 0, 0)
		return
	}
	_ = p.repository.MarkReady(ctx, item.ID, item.StorageKey, key, 0, 0)
}