leech/thumbnail/thumbnail.go

package thumbnail

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"leech/config"
	"os"
	"path/filepath"
	"slices"
	"sync"

	"github.com/disintegration/imaging"
	"github.com/fogleman/fauxgl"
	"github.com/gen2brain/go-fitz"
	"github.com/gofiber/fiber/v2"
	ffmpeg "github.com/u2takey/ffmpeg-go"
)

var SupportedFileTypes = []string{".png", ".PNG", ".jpg", ".JPG", ".jpeg", ".JPEG", ".webp", ".webP", ".WEBP", ".pdf", ".PDF", ".mp4", ".MP4", ".webm", ".WEBM", ".mkv", ".MKV", ".obj", ".OBJ", ".stl", ".STL", ".ply", ".PLY", ".3ds", ".3DS"}
var FileTypesMap = map[string]imaging.Format{
	".png":  imaging.PNG,
	".PNG":  imaging.PNG,
	".jpg":  imaging.JPEG,
	".JPG":  imaging.JPEG,
	".jpeg": imaging.JPEG,
	".JPEG": imaging.JPEG,
	".pdf":  imaging.PNG,
	".PDF":  imaging.PNG,
	".webp": imaging.PNG,
	".WEBP": imaging.PNG,
	".mp4":  imaging.PNG,
	".MP4":  imaging.PNG,
	".webm": imaging.PNG,
	".WEBM": imaging.PNG,
	".mkv":  imaging.PNG,
	".MKV":  imaging.PNG,
	".obj":  imaging.PNG,
	".OBJ":  imaging.PNG,
	".stl":  imaging.PNG,
	".STL":  imaging.PNG,
	".ply":  imaging.PNG,
	".PLY":  imaging.PNG,
	".3ds":  imaging.PNG,
	".3DS":  imaging.PNG,
}

var FFMPEGFormats = []string{
	".webp",
	".webP",
	".WEBP",
	".mp4",
	".MP4",
	".webm",
	".WEBM",
	".mkv",
	".MKV",
}

var FauxGLFormats = []string{
	".obj",
	".OBJ",
	".stl",
	".STL",
	".ply",
	".PLY",
	".3ds",
	".3DS",
}

var thumbnailSize = 48
var AA = 2

var thumbnailCache = map[string][]byte{}
var thumbnailCacheMutex = &sync.RWMutex{}
var memLimiterMutex = &sync.RWMutex{}
var jobCounter = 0
var checkAgain = make(chan bool, 5)

func IsSupportedFileType(completePath string) bool {
	fileExt := filepath.Ext(completePath)
	return slices.Contains(SupportedFileTypes, fileExt)
}

func WaitForAvailable() {
	memLimiterMutex.RLock()
	for config.Config.ThumbnailJobLimit == jobCounter {
		memLimiterMutex.RUnlock()
		<-checkAgain
		memLimiterMutex.RLock()
	}
	memLimiterMutex.RUnlock()
	memLimiterMutex.Lock()
	jobCounter += 1
	memLimiterMutex.Unlock()
}

func Free() {
	memLimiterMutex.Lock()
	jobCounter -= 1
	memLimiterMutex.Unlock()
	select {
	case checkAgain <- true:
		return
	default:
		return
	}
}

func GetThumbnail(c *fiber.Ctx, completePath string) {
	c.Set(fiber.HeaderContentType, "image")
	thumbnailCacheMutex.RLock()
	bytesThumb, ok := thumbnailCache[completePath]
	thumbnailCacheMutex.RUnlock()
	if ok {
		c.Write(bytesThumb)
		return
	}
	WaitForAvailable()
	/* While waiting some other thread could have generated the thumbnail */
	thumbnailCacheMutex.RLock()
	bytesThumb, ok = thumbnailCache[completePath]
	thumbnailCacheMutex.RUnlock()
	if ok {
		c.Write(bytesThumb)
		Free()
		return
	}

	fileExt := filepath.Ext(completePath)
	if !slices.Contains(SupportedFileTypes, fileExt) {
		Free()
		return
	}
	f, err := os.Open(completePath)
	if err != nil {
		Free()
		return
	}
	var img image.Image
	if fileExt == ".pdf" || fileExt == ".PDF" {
		doc, err := fitz.New(completePath)
		if err != nil {
			Free()
			return
		}
		img, err = doc.Image(0)
		if err != nil {
			Free()
			return
		}
		doc.Close()
	} else if slices.Contains(FFMPEGFormats, fileExt) {
		buf := bytes.NewBuffer(nil)
		err := ffmpeg.Input(completePath).
			Filter("select", ffmpeg.Args{fmt.Sprintf("gte(n,%d)", 0)}).
			Output("pipe:", ffmpeg.KwArgs{"vframes": 1, "format": "image2", "vcodec": "mjpeg"}).
			WithOutput(buf, nil). //os.Stdout
			Silent(true).
			Run()
		if err != nil {
			Free()
			return
		}
		img, err = imaging.Decode(buf)
		if err != nil {
			Free()
			return
		}
	} else if slices.Contains(FauxGLFormats, fileExt) {
		// load a mesh
		var mesh *fauxgl.Mesh
		if fileExt == ".obj" || fileExt == ".OBJ" {
			mesh, err = fauxgl.LoadOBJ(completePath)
		} else if fileExt == ".stl" || fileExt == ".STL" {
			mesh, err = fauxgl.LoadSTL(completePath)
		} else if fileExt == ".ply" || fileExt == ".PLY" {
			mesh, err = fauxgl.LoadPLY(completePath)
		} else {
			mesh, err = fauxgl.Load3DS(completePath)
		}
		if err != nil {
			Free()
			return
		}

		// fit mesh in a bi-unit cube centered at the origin
		mesh.BiUnitCube()

		// smooth the normals
		mesh.SmoothNormalsThreshold(fauxgl.Radians(30))

		// create a rendering context
		context := fauxgl.NewContext(thumbnailSize*AA, thumbnailSize*AA)
		/*context.ClearColorBufferWith(fauxgl.HexColor("#FFF8E3"))*/

		// create transformation matrix and light direction
		aspect := 1.
		matrix := fauxgl.LookAt(fauxgl.V(-4, 1.5, -2) /* eye */, fauxgl.V(0, -0.07, 0) /* center */, fauxgl.V(0, 1, 0) /* up */).Perspective(33 /*fovy*/, aspect, 1 /* near */, 30 /* far */)

		// use builtin phong shader
		shader := fauxgl.NewPhongShader(matrix, fauxgl.V(-0.75, 1, 0.25).Normalize() /* light */, fauxgl.V(-4, 1.5, -2) /* eye */)
		shader.ObjectColor = fauxgl.HexColor("#9e5272")
		context.Shader = shader

		// render
		context.DrawMesh(mesh)
		img = context.Image()
	} else {
		img, _, err = image.Decode(f)
		if err != nil {
			Free()
			return
		}
		f.Close()
	}
	// load images and make 64x64 thumbnails of them
	thumbnail := imaging.Thumbnail(img, thumbnailSize, thumbnailSize, imaging.CatmullRom)
	// create a new blank image
	dst := imaging.New(thumbnailSize, thumbnailSize, color.NRGBA{0, 0, 0, 0})
	// paste thumbnails into the new image side by side
	dst = imaging.Paste(dst, thumbnail, image.Pt(0, 0))
	// save the combined image to buffer
	var buf bytes.Buffer
	if FileTypesMap[fileExt] == imaging.PNG {
		err = png.Encode(&buf, dst)
	} else {
		err = imaging.Encode(&buf, dst, FileTypesMap[fileExt])
	}
	if err != nil {
		Free()
		return
	}
	thumbnailCacheMutex.Lock()
	thumbnailCache[completePath] = buf.Bytes()
	thumbnailCacheMutex.Unlock()
	c.Write(buf.Bytes())

	Free()
}