WarpBox/lib/boxstore/thumbnails.go

package boxstore

import (
	"image"
	"image/color"
	"image/draw"
	_ "image/gif"
	"image/jpeg"
	_ "image/png"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"warpbox/lib/helpers"
	"warpbox/lib/models"
)

const (
	thumbnailDir     = ".thumbnails"
	thumbnailMaxSize = 160
)

type thumbnailTask struct {
	BoxID  string
	FileID string
	Name   string
}

func StartThumbnailWorker(batchSize int, interval time.Duration) {
	if batchSize < 1 {
		batchSize = 10
	}

	if interval <= 0 {
		interval = 30 * time.Second
	}

	go func() {
		for {
			ProcessThumbnailBatch(batchSize)
			time.Sleep(interval)
		}
	}()
}

func ProcessThumbnailBatch(batchSize int) int {
	tasks := collectThumbnailTasks(batchSize)
	for _, task := range tasks {
		if err := generateThumbnail(task); err != nil {
			markThumbnailFailed(task.BoxID, task.FileID)
			continue
		}
	}

	return len(tasks)
}

func ThumbnailFilePath(boxID string, fileID string) (string, bool) {
	if !helpers.ValidLowerHexID(fileID, 16) {
		return "", false
	}

	return helpers.SafeChildPath(filepath.Join(BoxPath(boxID), thumbnailDir), fileID+".jpg")
}

func collectThumbnailTasks(batchSize int) []thumbnailTask {
	entries, err := os.ReadDir(uploadRoot)
	if err != nil {
		return nil
	}

	tasks := make([]thumbnailTask, 0, batchSize)
	for _, entry := range entries {
		if !entry.IsDir() || !ValidBoxID(entry.Name()) {
			continue
		}

		tasks = append(tasks, collectBoxThumbnailTasks(entry.Name(), batchSize-len(tasks))...)
		if len(tasks) >= batchSize {
			return tasks
		}
	}

	return tasks
}

func collectBoxThumbnailTasks(boxID string, remaining int) []thumbnailTask {
	if remaining <= 0 {
		return nil
	}

	manifestMu.Lock()
	defer manifestMu.Unlock()

	manifest, err := readManifestUnlocked(boxID)
	if err != nil || IsExpired(manifest) || manifest.OneTimeDownload {
		return nil
	}

	tasks := make([]thumbnailTask, 0, remaining)
	changed := false
	for index, file := range manifest.Files {
		if len(tasks) >= remaining {
			break
		}

		if file.Status != models.FileStatusReady || file.ThumbnailPath != nil || file.ThumbnailStatus != "" {
			continue
		}

		if !canGenerateThumbnail(file) {
			manifest.Files[index].ThumbnailStatus = models.ThumbnailStatusUnsupported
			changed = true
			continue
		}

		tasks = append(tasks, thumbnailTask{
			BoxID:  boxID,
			FileID: file.ID,
			Name:   file.Name,
		})
	}

	if changed {
		writeManifestUnlocked(boxID, manifest)
	}

	return tasks
}

func canGenerateThumbnail(file models.BoxFile) bool {
	if strings.HasPrefix(file.MimeType, "image/") {
		return true
	}

	extension := strings.ToLower(filepath.Ext(file.Name))
	return extension == ".jpg" || extension == ".jpeg" || extension == ".png" || extension == ".gif"
}

func generateThumbnail(task thumbnailTask) error {
	sourcePath, ok := SafeBoxFilePath(task.BoxID, task.Name)
	if !ok {
		return os.ErrInvalid
	}
	if err := ensureRegularFile(sourcePath); err != nil {
		return err
	}

	source, err := os.Open(sourcePath)
	if err != nil {
		return err
	}
	defer source.Close()

	src, _, err := image.Decode(source)
	if err != nil {
		return err
	}

	thumb := resizeImage(src, thumbnailMaxSize)
	if err := os.MkdirAll(filepath.Join(BoxPath(task.BoxID), thumbnailDir), 0755); err != nil {
		return err
	}

	path, ok := ThumbnailFilePath(task.BoxID, task.FileID)
	if !ok {
		return os.ErrInvalid
	}

	target, tempPath, err := createTempSibling(path)
	if err != nil {
		return err
	}
	committed := false
	defer func() {
		target.Close()
		if !committed {
			os.Remove(tempPath)
		}
	}()

	if err := jpeg.Encode(target, thumb, &jpeg.Options{Quality: 82}); err != nil {
		return err
	}
	if err := target.Close(); err != nil {
		return err
	}
	if err := os.Rename(tempPath, path); err != nil {
		return err
	}
	committed = true

	return markThumbnailReady(task.BoxID, task.FileID)
}

func resizeImage(src image.Image, maxSize int) image.Image {
	bounds := src.Bounds()
	width := bounds.Dx()
	height := bounds.Dy()
	if width <= 0 || height <= 0 {
		return src
	}

	targetWidth := width
	targetHeight := height
	if width > maxSize || height > maxSize {
		if width >= height {
			targetWidth = maxSize
			targetHeight = height * maxSize / width
		} else {
			targetHeight = maxSize
			targetWidth = width * maxSize / height
		}
	}

	if targetWidth < 1 {
		targetWidth = 1
	}
	if targetHeight < 1 {
		targetHeight = 1
	}

	dst := image.NewRGBA(image.Rect(0, 0, targetWidth, targetHeight))
	draw.Draw(dst, dst.Bounds(), image.NewUniform(color.White), image.Point{}, draw.Src)
	for y := 0; y < targetHeight; y++ {
		for x := 0; x < targetWidth; x++ {
			srcX := bounds.Min.X + x*width/targetWidth
			srcY := bounds.Min.Y + y*height/targetHeight
			dst.Set(x, y, src.At(srcX, srcY))
		}
	}

	return dst
}

func markThumbnailReady(boxID string, fileID string) error {
	path := "/box/" + boxID + "/thumbnails/" + url.PathEscape(fileID)
	return updateThumbnailState(boxID, fileID, &path, models.ThumbnailStatusReady)
}

func markThumbnailFailed(boxID string, fileID string) {
	updateThumbnailState(boxID, fileID, nil, models.ThumbnailStatusFailed)
}

func updateThumbnailState(boxID string, fileID string, thumbnailPath *string, status string) error {
	manifestMu.Lock()
	defer manifestMu.Unlock()

	manifest, err := readManifestUnlocked(boxID)
	if err != nil {
		return err
	}

	for index, file := range manifest.Files {
		if file.ID != fileID {
			continue
		}

		manifest.Files[index].ThumbnailPath = thumbnailPath
		manifest.Files[index].ThumbnailStatus = status
		return writeManifestUnlocked(boxID, manifest)
	}

	return os.ErrNotExist
}