stash/internal/identify/scene.go

package identify

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/stashapp/stash/pkg/logger"
	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/scene"
	"github.com/stashapp/stash/pkg/sliceutil"
	"github.com/stashapp/stash/pkg/sliceutil/intslice"
	"github.com/stashapp/stash/pkg/tag"
	"github.com/stashapp/stash/pkg/utils"
)

type SceneReaderUpdater interface {
	GetCover(ctx context.Context, sceneID int) ([]byte, error)
	scene.Updater
	models.PerformerIDLoader
	models.TagIDLoader
	models.StashIDLoader
	models.URLLoader
}

type TagCreatorFinder interface {
	Create(ctx context.Context, newTag *models.Tag) error
	tag.Finder
}

type sceneRelationships struct {
	sceneReader              SceneReaderUpdater
	studioReaderWriter       models.StudioReaderWriter
	performerCreator         PerformerCreator
	tagCreatorFinder         TagCreatorFinder
	scene                    *models.Scene
	result                   *scrapeResult
	fieldOptions             map[string]*FieldOptions
	skipSingleNamePerformers bool
}

func (g sceneRelationships) studio(ctx context.Context) (*int, error) {
	existingID := g.scene.StudioID
	fieldStrategy := g.fieldOptions["studio"]
	createMissing := fieldStrategy != nil && utils.IsTrue(fieldStrategy.CreateMissing)

	scraped := g.result.result.Studio
	endpoint := g.result.source.RemoteSite

	if scraped == nil || !shouldSetSingleValueField(fieldStrategy, existingID != nil) {
		return nil, nil
	}

	if scraped.StoredID != nil {
		// existing studio, just set it
		studioID, err := strconv.Atoi(*scraped.StoredID)
		if err != nil {
			return nil, fmt.Errorf("error converting studio ID %s: %w", *scraped.StoredID, err)
		}

		// only return value if different to current
		if existingID == nil || *existingID != studioID {
			return &studioID, nil
		}
	} else if createMissing {
		return createMissingStudio(ctx, endpoint, g.studioReaderWriter, scraped)
	}

	return nil, nil
}

func (g sceneRelationships) performers(ctx context.Context, ignoreMale bool) ([]int, error) {
	fieldStrategy := g.fieldOptions["performers"]
	scraped := g.result.result.Performers

	// just check if ignored
	if len(scraped) == 0 || !shouldSetSingleValueField(fieldStrategy, false) {
		return nil, nil
	}

	createMissing := fieldStrategy != nil && utils.IsTrue(fieldStrategy.CreateMissing)
	strategy := FieldStrategyMerge
	if fieldStrategy != nil {
		strategy = fieldStrategy.Strategy
	}

	endpoint := g.result.source.RemoteSite

	var performerIDs []int
	originalPerformerIDs := g.scene.PerformerIDs.List()

	if strategy == FieldStrategyMerge {
		// add to existing
		performerIDs = originalPerformerIDs
	}

	singleNamePerformerSkipped := false

	for _, p := range scraped {
		if ignoreMale && p.Gender != nil && strings.EqualFold(*p.Gender, models.GenderEnumMale.String()) {
			continue
		}

		performerID, err := getPerformerID(ctx, endpoint, g.performerCreator, p, createMissing, g.skipSingleNamePerformers)
		if err != nil {
			if errors.Is(err, ErrSkipSingleNamePerformer) {
				singleNamePerformerSkipped = true
				continue
			}
			return nil, err
		}

		if performerID != nil {
			performerIDs = intslice.IntAppendUnique(performerIDs, *performerID)
		}
	}

	// don't return if nothing was added
	if sliceutil.SliceSame(originalPerformerIDs, performerIDs) {
		if singleNamePerformerSkipped {
			return nil, ErrSkipSingleNamePerformer
		}
		return nil, nil
	}

	if singleNamePerformerSkipped {
		return performerIDs, ErrSkipSingleNamePerformer
	}
	return performerIDs, nil
}

func (g sceneRelationships) tags(ctx context.Context) ([]int, error) {
	fieldStrategy := g.fieldOptions["tags"]
	scraped := g.result.result.Tags
	target := g.scene

	// just check if ignored
	if len(scraped) == 0 || !shouldSetSingleValueField(fieldStrategy, false) {
		return nil, nil
	}

	createMissing := fieldStrategy != nil && utils.IsTrue(fieldStrategy.CreateMissing)
	strategy := FieldStrategyMerge
	if fieldStrategy != nil {
		strategy = fieldStrategy.Strategy
	}

	var tagIDs []int
	originalTagIDs := target.TagIDs.List()

	if strategy == FieldStrategyMerge {
		// add to existing
		tagIDs = originalTagIDs
	}

	for _, t := range scraped {
		if t.StoredID != nil {
			// existing tag, just add it
			tagID, err := strconv.ParseInt(*t.StoredID, 10, 64)
			if err != nil {
				return nil, fmt.Errorf("error converting tag ID %s: %w", *t.StoredID, err)
			}

			tagIDs = intslice.IntAppendUnique(tagIDs, int(tagID))
		} else if createMissing {
			now := time.Now()
			newTag := models.Tag{
				Name:      t.Name,
				CreatedAt: now,
				UpdatedAt: now,
			}
			err := g.tagCreatorFinder.Create(ctx, &newTag)
			if err != nil {
				return nil, fmt.Errorf("error creating tag: %w", err)
			}

			tagIDs = append(tagIDs, newTag.ID)
		}
	}

	// don't return if nothing was added
	if sliceutil.SliceSame(originalTagIDs, tagIDs) {
		return nil, nil
	}

	return tagIDs, nil
}

func (g sceneRelationships) stashIDs(ctx context.Context) ([]models.StashID, error) {
	remoteSiteID := g.result.result.RemoteSiteID
	fieldStrategy := g.fieldOptions["stash_ids"]
	target := g.scene

	endpoint := g.result.source.RemoteSite

	// just check if ignored
	if remoteSiteID == nil || endpoint == "" || !shouldSetSingleValueField(fieldStrategy, false) {
		return nil, nil
	}

	strategy := FieldStrategyMerge
	if fieldStrategy != nil {
		strategy = fieldStrategy.Strategy
	}

	var stashIDs []models.StashID
	originalStashIDs := target.StashIDs.List()

	if strategy == FieldStrategyMerge {
		// add to existing
		// make a copy so we don't modify the original
		stashIDs = append(stashIDs, originalStashIDs...)
	}

	for i, stashID := range stashIDs {
		if endpoint == stashID.Endpoint {
			// if stashID is the same, then don't set
			if stashID.StashID == *remoteSiteID {
				return nil, nil
			}

			// replace the stash id and return
			stashID.StashID = *remoteSiteID
			stashIDs[i] = stashID
			return stashIDs, nil
		}
	}

	// not found, create new entry
	stashIDs = append(stashIDs, models.StashID{
		StashID:  *remoteSiteID,
		Endpoint: endpoint,
	})

	if sliceutil.SliceSame(originalStashIDs, stashIDs) {
		return nil, nil
	}

	return stashIDs, nil
}

func (g sceneRelationships) cover(ctx context.Context) ([]byte, error) {
	scraped := g.result.result.Image

	if scraped == nil || *scraped == "" {
		return nil, nil
	}

	// always overwrite if present
	existingCover, err := g.sceneReader.GetCover(ctx, g.scene.ID)
	if err != nil {
		logger.Errorf("Error getting scene cover: %v", err)
	}

	data, err := utils.ProcessImageInput(ctx, *scraped)
	if err != nil {
		return nil, fmt.Errorf("error processing image input: %w", err)
	}

	// only return if different
	if !bytes.Equal(existingCover, data) {
		return data, nil
	}

	return nil, nil
}