streamrip-go/internal/artwork/artwork.go

package artwork

import (
	"context"
	"crypto/sha1"
	"fmt"
	"image"
	"image/jpeg"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"golang.org/x/image/draw"

	"streamrip-go/internal/config"
)

type Downloader interface {
	File(ctx context.Context, sourceURL, outputPath string) error
	FileNoProgress(ctx context.Context, sourceURL, outputPath string) error
}

type Result struct {
	EmbedPath string
	SavedPath string
}

var (
	tempDirsMu sync.Mutex
	tempDirs   = map[string]struct{}{}
)

func Prepare(ctx context.Context, dl Downloader, folder string, albumMeta map[string]any, cfg config.ArtworkConfig, forPlaylist bool) (Result, error) {
	saveArtwork := cfg.SaveArtwork
	if forPlaylist {
		saveArtwork = false
	}
	if !(cfg.Embed || saveArtwork) {
		return Result{}, nil
	}

	imageMap, ok := albumMeta["image"].(map[string]any)
	if !ok {
		return Result{}, nil
	}

	largestURL := pickLargestURL(imageMap)
	embedURL := pickEmbedURL(imageMap, cfg.EmbedSize)
	if embedURL == "" {
		embedURL = largestURL
	}

	result := Result{}
	if saveArtwork && largestURL != "" {
		savedPath := filepath.Join(folder, "cover.jpg")
		if fileExists(savedPath) {
			result.SavedPath = savedPath
		} else if err := dl.FileNoProgress(ctx, largestURL, savedPath); err == nil {
			if cfg.SavedMaxWidth > 0 {
				_ = downscaleImage(savedPath, cfg.SavedMaxWidth)
			}
			result.SavedPath = savedPath
		}
	}

	if cfg.Embed && embedURL != "" {
		embedDir := sessionEmbedDir(folder)
		if err := os.MkdirAll(embedDir, 0o755); err == nil {
			registerTempDir(embedDir)
			embedPath := filepath.Join(embedDir, embedFilename(embedURL))
			if fileExists(embedPath) {
				result.EmbedPath = embedPath
			} else if err := dl.FileNoProgress(ctx, embedURL, embedPath); err == nil {
				if cfg.EmbedMaxWidth > 0 {
					_ = downscaleImage(embedPath, cfg.EmbedMaxWidth)
				}
				result.EmbedPath = embedPath
			}
		}
	}

	return result, nil
}

func CleanupTempDirs() {
	tempDirsMu.Lock()
	defer tempDirsMu.Unlock()

	for dir := range tempDirs {
		_ = os.RemoveAll(dir)
		delete(tempDirs, dir)
	}
}

func registerTempDir(path string) {
	tempDirsMu.Lock()
	defer tempDirsMu.Unlock()
	tempDirs[path] = struct{}{}
}

func fileExists(path string) bool {
	st, err := os.Stat(path)
	if err != nil {
		return false
	}
	return !st.IsDir()
}

func pickLargestURL(image map[string]any) string {
	for _, key := range []string{"original", "mega", "extralarge", "large", "small", "thumbnail"} {
		if v := stringAny(image[key]); v != "" {
			return v
		}
	}
	return ""
}

func pickEmbedURL(image map[string]any, size string) string {
	size = strings.ToLower(strings.TrimSpace(size))
	if size == "" {
		size = "large"
	}
	for _, key := range []string{size, "large", "extralarge", "small", "thumbnail", "original"} {
		if v := stringAny(image[key]); v != "" {
			return v
		}
	}
	return ""
}

func embedFilename(url string) string {
	s := sha1.Sum([]byte(url))
	return fmt.Sprintf("cover%x.jpg", s[:8])
}

func sessionEmbedDir(folder string) string {
	key := sha1.Sum([]byte(folder))
	return filepath.Join(os.TempDir(), "streamrip-go-artwork", fmt.Sprintf("%x", key[:8]))
}

func stringAny(v any) string {
	s, _ := v.(string)
	return s
}

func downscaleImage(path string, maxDimension int) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	img, _, err := image.Decode(f)
	_ = f.Close()
	if err != nil {
		return err
	}

	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	if w <= 0 || h <= 0 || maxDimension <= 0 {
		return nil
	}
	if w <= maxDimension && h <= maxDimension {
		return nil
	}

	newW, newH := w, h
	if w > h {
		newW = maxDimension
		newH = int(float64(h) * (float64(maxDimension) / float64(w)))
	} else {
		newH = maxDimension
		newW = int(float64(w) * (float64(maxDimension) / float64(h)))
	}
	if newW <= 0 {
		newW = 1
	}
	if newH <= 0 {
		newH = 1
	}

	dst := image.NewRGBA(image.Rect(0, 0, newW, newH))
	draw.CatmullRom.Scale(dst, dst.Bounds(), img, b, draw.Over, nil)

	out, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() { _ = out.Close() }()
	return jpeg.Encode(out, dst, &jpeg.Options{Quality: 92})
}