streamrip-go/internal/artwork/artwork.go

package artwork

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

	"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{}{}
)

var (
	folderLocksMu sync.Mutex
	folderLocks   = map[string]*sync.Mutex{}
)

// lockFolder serializes Prepare calls that target the same album folder and
// returns the unlock function.
//
// ripTrack calls Prepare once per track with the same albumFolder for every
// worker in a playlist rip, and the cover handling below is check-then-write.
// Without this, two workers both saw a missing cover.jpg, both created it, and
// the one that failed removed the file the other had just finished writing.
func lockFolder(folder string) func() {
	folderLocksMu.Lock()
	mu, ok := folderLocks[folder]
	if !ok {
		mu = &sync.Mutex{}
		folderLocks[folder] = mu
	}
	folderLocksMu.Unlock()

	mu.Lock()
	return mu.Unlock
}

// fetchAtomic downloads sourceURL to destPath via a unique temp file in the
// same directory, then renames it into place. Downloading straight to destPath
// meant a failed attempt could delete a file another worker had completed.
func fetchAtomic(ctx context.Context, dl Downloader, sourceURL, destPath string, maxWidth int) bool {
	tmpPath := fmt.Sprintf("%s.part-%d-%d", destPath, os.Getpid(), atomic.AddUint64(&tmpCounter, 1))
	if err := dl.FileNoProgress(ctx, sourceURL, tmpPath); err != nil {
		_ = os.Remove(tmpPath)
		return false
	}
	if maxWidth > 0 {
		_ = downscaleImage(tmpPath, maxWidth)
	}
	if err := os.Rename(tmpPath, destPath); err != nil {
		_ = os.Remove(tmpPath)
		return false
	}
	return true
}

var tmpCounter uint64

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
	}

	// Everything below is check-then-write against shared paths.
	unlock := lockFolder(folder)
	defer unlock()

	result := Result{}
	if saveArtwork && largestURL != "" {
		savedPath := filepath.Join(folder, "cover.jpg")
		if fileExists(savedPath) || fetchAtomic(ctx, dl, largestURL, savedPath, cfg.SavedMaxWidth) {
			result.SavedPath = savedPath
		}
	}

	if cfg.Embed && embedURL != "" {
		embedDir := sessionEmbedDir(folder)
		// 0o700: predictable path under the shared system temp dir.
		if err := os.MkdirAll(embedDir, 0o700); err == nil {
			registerTempDir(embedDir)
			embedPath := filepath.Join(embedDir, embedFilename(embedURL))
			if fileExists(embedPath) || fetchAtomic(ctx, dl, embedURL, 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})
}