comic-hasher/ch/hashing.go

package ch

import (
	"cmp"
	_ "embed"
	"fmt"
	"image"
	"math/bits"
	"os"
	"slices"
	"strings"
	"sync"
	"time"
	"unicode"
	"unsafe"

	"gitea.narnian.us/lordwelch/comic-hasher/log"

	"gitea.narnian.us/lordwelch/goimagehash"
)

//go:embed hashes.gz
var Hashes []byte

const (
	H0 uint64 = 0b11111111 << (8 * iota)
	H1
	H2
	H3
	H4
	H5
	H6
	H7
)

const (
	Shift0 = (8 * iota)
	Shift1
	Shift2
	Shift3
	Shift4
	Shift5
	Shift6
	Shift7
)

const (
	ComicVine Source = "comicvine.gamespot.com"
	Metron    Source = "metron.cloud"
)

var sources *sync.Map = newSourceMap()

type Source string

type Match struct {
	Distance int
	Hash     uint64
}

type ID struct {
	Domain *Source
	ID     string
}

type Result struct {
	Hash          Hash
	ID            ID
	Distance      int
	EquivalentIDs []ID
}
type Im struct {
	Im      image.Image
	Format  string
	ID      ID
	NewOnly bool
}

type ImageHash struct {
	Hashes []Hash
	ID     ID
}

type Hash struct {
	Hash uint64
	Kind goimagehash.Kind
}

func (id *ID) Compare(target ID) int {
	return cmp.Or(
		strings.Compare(string(*id.Domain), string(*target.Domain)),
		strings.Compare(id.ID, target.ID),
	)
}

func (id ID) String() string {
	return fmt.Sprintf("%s: %v", *id.Domain, id.ID)
}

func (id ID) GoString() string {
	return fmt.Sprintf("ch.Source{Domain: %#v, ID: %#v}", *id.Domain, id.ID)
}

func newSourceMap() *sync.Map {
	m := &sync.Map{}
	for _, s := range []Source{ComicVine} {
		m.Store(s, &s)
	}
	_, ok := m.Load(ComicVine)
	if !ok {
		log.Log(log.Fatal, "sync.Map really doesn't work, %#v", ComicVine)
		time.Sleep(11 * time.Millisecond) // for logging to complete
		os.Exit(1)
	}
	_, ok = m.Load(Source(strings.ToLower(string(ComicVine))))
	if !ok {
		log.Println(log.Fatal, "sync.Map doesn't work")
		time.Sleep(11 * time.Millisecond) // for logging to complete
		os.Exit(1)
	}
	return m
}

func IsLower(s string) bool {
	for _, r := range s {
		if !unicode.IsLower(r) && unicode.IsLetter(r) {
			return false
		}
	}
	return true
}

func Lower[E string | Source](s E) Source {
	if !IsLower(string(s)) {
		return Source(strings.ToLower(string(s)))
	}
	return Source(s)
}

func NewSource[E string | Source](s E) *Source {
	if sp, ok := sources.Load(Source(s)); ok {
		return sp.(*Source)
	}
	s1 := Lower(s)
	if sp, ok := sources.Load(s1); ok {
		return sp.(*Source)
	}
	s2 := Source(Clone(string(s1)))
	sp, _ := sources.LoadOrStore(s2, &s2)
	return sp.(*Source)
}

func Clone(s string) string {
	if len(s) == 0 {
		return ""
	}
	b := make([]byte, len(s))
	copy(b, s)
	return unsafe.String(&b[0], len(b))
}

type HashStorage interface {
	GetMatches(hashes []Hash, max int, exactOnly bool) ([]Result, error)
	MapHashes(ImageHash) Metrics
	DecodeHashes(hashes *SavedHashesV3) (Metrics, error)
	EncodeHashes() (*SavedHashesV3, error)
	AssociateIDs(OldID, NewID ID) error
	GetIDs(id ID) []ID
	RLock()
	RUnlock()
}

func Atleast(maxDistance int, searchHash uint64, hashes []uint64) []Match {
	matchingHashes := make([]Match, 0, 20) // hope that we don't need all of them
	var storedHash uint64
	for _, storedHash = range hashes {
		distance := bits.OnesCount64(searchHash ^ storedHash)
		if distance <= maxDistance {
			matchingHashes = append(matchingHashes, Match{distance, storedHash})
		}
	}
	return matchingHashes
}

func InsertIdx[S ~[]E, E cmp.Ordered](slice S, item E) (S, int) {
	index, itemFound := slices.BinarySearch(slice, item)
	if itemFound {
		return slice, index
	}
	return slices.Insert(slice, index, item), index
}

func Insert[S ~[]E, E cmp.Ordered](slice S, item E) S {
	slice, _ = InsertIdx(slice, item)
	return slice
}

func HashImage(i Im) ImageHash {
	if i.Format == "webp" {
		i.Im = goimagehash.FancyUpscale(i.Im.(*image.YCbCr))
	}

	var err error

	ahash, err := goimagehash.AverageHash(i.Im)
	if err != nil {
		log.Log(log.Fatal, "Failed to ahash Image: %s", err)
		return ImageHash{}
	}
	dhash, err := goimagehash.DifferenceHash(i.Im)
	if err != nil {
		log.Log(log.Fatal, "Failed to dhash Image: %s", err)
		return ImageHash{}
	}
	phash, err := goimagehash.PerceptionHash(i.Im)
	if err != nil {
		log.Log(log.Fatal, "Failed to phash Image: %s", err)
		return ImageHash{}
	}
	return ImageHash{
		Hashes: []Hash{{ahash.GetHash(), ahash.GetKind()}, {dhash.GetHash(), dhash.GetKind()}, {phash.GetHash(), phash.GetKind()}},
		ID:     i.ID,
	}
}

func SplitHash(hash uint64) [8]uint8 {
	return [8]uint8{
		uint8((hash & H7) >> Shift7),
		uint8((hash & H6) >> Shift6),
		uint8((hash & H5) >> Shift5),
		uint8((hash & H4) >> Shift4),
		uint8((hash & H3) >> Shift3),
		uint8((hash & H2) >> Shift2),
		uint8((hash & H1) >> Shift1),
		uint8((hash & H0) >> Shift0),
	}
}