git/intake/main
internal/artwork/artwork.go
262 lines6 KB last change raw · history
package artwork

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"intake/internal/audio"
)

type Source string

const (
	SourceEmbedded    Source = "embedded"
	SourceFolder      Source = "folder"
	SourceCoverArt    Source = "coverartarchive"
	SourceDeezer      Source = "deezer"
	SourceITunes      Source = "itunes"
	SourceLocalFolder Source = "local"
	SourcePlaceholder Source = "placeholder"
)

type Candidate struct {
	Source        Source
	Label         string
	Path          string
	URL           string
	Width, Height int
	Bytes         int64
	Preferred     bool
	EmbeddedCount int
	MatchScore    int
	ProviderTitle string
}

func (c Candidate) ExactMatch() bool { return c.MatchScore >= 3 }

func (c Candidate) Describe() string {
	dims := "unknown size"
	if c.Width > 0 && c.Height > 0 {
		dims = fmt.Sprintf("%dx%d", c.Width, c.Height)
	}
	return fmt.Sprintf("%s, %s", dims, humanBytes(c.Bytes))
}

func (c Candidate) IsLowResolution() bool {
	if c.Width == 0 || c.Height == 0 {
		return false
	}
	return c.Width < 500 || c.Height < 500
}

type Options struct {
	ExcludeHashes    []string
	StagingDir       string
	ArchiveDir       string
	LocalCoverFolder string
	Sources          []Source
	Offline          bool
	CacheDir         string
	Progress         func(string)
}

func Discover(ctx context.Context, album AlbumInfo, tracks []audio.Track, opts Options) ([]Candidate, []string, error) {
	var (
		out      []Candidate
		warnings []string
	)
	progress := opts.Progress
	if progress == nil {
		progress = func(string) {}
	}

	if embedded := embeddedCandidates(ctx, tracks); len(embedded) > 0 {
		out = append(out, embedded...)
	}

	for _, dir := range []string{opts.ArchiveDir, opts.StagingDir} {
		if dir == "" {
			continue
		}
		if found := folderCandidates(ctx, dir, opts.StagingDir); len(found) > 0 {
			out = append(out, found...)
			break
		}
	}

	if opts.LocalCoverFolder != "" {
		if found := localCandidates(ctx, opts.LocalCoverFolder, album); len(found) > 0 {
			out = append(out, found...)
		}
	}

	if !opts.Offline && enabled(opts.Sources, SourceCoverArt) {
		progress("querying MusicBrainz and Cover Art Archive ...")
		got, err := coverArtArchive(ctx, album, opts)
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("cover art archive: %v", err))
		}
		out = append(out, got...)
	}
	if !opts.Offline && enabled(opts.Sources, SourceDeezer) {
		progress("querying Deezer ...")
		got, err := deezer(ctx, album, opts)
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("deezer: %v", err))
		}
		out = append(out, got...)
	}
	if !opts.Offline && enabled(opts.Sources, SourceITunes) {
		progress("querying iTunes ...")
		got, err := itunes(ctx, album, opts)
		if err != nil {
			warnings = append(warnings, fmt.Sprintf("itunes: %v", err))
		}
		out = append(out, got...)
	}

	if len(opts.ExcludeHashes) > 0 {
		excluded := map[string]bool{}
		for _, h := range opts.ExcludeHashes {
			excluded[h] = true
		}
		kept := out[:0]
		for _, c := range out {
			if h, err := hashFile(c.Path); err == nil && excluded[h] {
				continue
			}
			kept = append(kept, c)
		}
		out = kept
	}

	out = dedupe(out)
	sort.SliceStable(out, func(i, j int) bool {
		a, b := out[i], out[j]
		if a.Preferred != b.Preferred {
			return a.Preferred
		}
		ap, bp := a.Width*a.Height, b.Width*b.Height
		if ap != bp {
			return ap > bp
		}
		return a.Bytes > b.Bytes
	})
	return out, warnings, nil
}

type AlbumInfo struct {
	Artist string
	Album  string
	Year   string
	Tracks int
}

func enabled(list []Source, s Source) bool {
	if len(list) == 0 {
		return true
	}
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}

func dedupe(in []Candidate) []Candidate {
	type kept struct {
		bytes  string
		visual visual
		hashed bool
	}
	var (
		out []Candidate
		vis []kept
	)
	for _, c := range in {
		// Prefer the visual hash, which catches one picture at several sizes or
		// re-encoded by another provider. The byte hash is the fallback for an
		// image that cannot be decoded.
		dupe := -1
		if vh, ok := visualHash(c.Path); ok {
			for i := range vis {
				if vis[i].hashed && vis[i].visual.near(vh) {
					dupe = i
					break
				}
			}
			if dupe < 0 {
				out = append(out, c)
				vis = append(vis, kept{visual: vh, hashed: true})
				continue
			}
		} else {
			key, err := hashFile(c.Path)
			if err != nil {
				out = append(out, c)
				vis = append(vis, kept{})
				continue
			}
			for i := range vis {
				if vis[i].bytes != "" && vis[i].bytes == key {
					dupe = i
					break
				}
			}
			if dupe < 0 {
				out = append(out, c)
				vis = append(vis, kept{bytes: key})
				continue
			}
		}

		if len(c.Label) > len(out[dupe].Label) && out[dupe].Label == "" {
			out[dupe].Label = c.Label
		}
		if out[dupe].Source == SourcePlaceholder && c.Source != SourcePlaceholder {
			out[dupe] = c
		}
	}
	return out
}

func hashFile(path string) (string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:]), nil
}

func Hash(path string) (string, error) { return hashFile(path) }

func CachePath(dir, url string) string {
	sum := sha256.Sum256([]byte(url))
	ext := strings.ToLower(filepath.Ext(url))
	if ext == "" || len(ext) > 5 {
		ext = ".jpg"
	}
	return filepath.Join(dir, hex.EncodeToString(sum[:16])+ext)
}

func humanBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for v := n / unit; v >= unit; v /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(n)/float64(div), "KMGT"[exp])
}