stash/pkg/stashbox/scene.go

package stashbox

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net/http"

	"github.com/stashapp/stash/pkg/logger"
	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/scraper"
	"github.com/stashapp/stash/pkg/sliceutil"
	"github.com/stashapp/stash/pkg/stashbox/graphql"
	"github.com/stashapp/stash/pkg/utils"
)

// QueryScene queries stash-box for scenes using a query string.
func (c Client) QueryScene(ctx context.Context, queryStr string) ([]*models.ScrapedScene, error) {
	scenes, err := c.client.SearchScene(ctx, queryStr)
	if err != nil {
		return nil, err
	}

	sceneFragments := scenes.SearchScene

	var ret []*models.ScrapedScene
	var ignoredTags []string
	for _, s := range sceneFragments {
		ss, err := c.sceneFragmentToScrapedScene(ctx, s)
		if err != nil {
			return nil, err
		}

		var thisIgnoredTags []string
		ss.Tags, thisIgnoredTags = scraper.FilterTags(c.excludeTagRE, ss.Tags)
		ignoredTags = sliceutil.AppendUniques(ignoredTags, thisIgnoredTags)

		ret = append(ret, ss)
	}

	scraper.LogIgnoredTags(ignoredTags)

	return ret, nil
}

// FindStashBoxScenesByFingerprints queries stash-box for a scene using the
// scene's MD5/OSHASH checksum, or PHash.
func (c Client) FindSceneByFingerprints(ctx context.Context, fps models.Fingerprints) ([]*models.ScrapedScene, error) {
	res, err := c.FindScenesByFingerprints(ctx, []models.Fingerprints{fps})
	if len(res) > 0 {
		return res[0], err
	}
	return nil, err
}

// FindScenesByFingerprints queries stash-box for scenes using every
// scene's MD5/OSHASH checksum, or PHash, and returns results in the same order
// as the input slice.
func (c Client) FindScenesByFingerprints(ctx context.Context, fps []models.Fingerprints) ([][]*models.ScrapedScene, error) {
	var fingerprints [][]*graphql.FingerprintQueryInput

	for _, fp := range fps {
		fingerprints = append(fingerprints, convertFingerprints(fp))
	}

	return c.findScenesByFingerprints(ctx, fingerprints)
}

func convertFingerprints(fps models.Fingerprints) []*graphql.FingerprintQueryInput {
	var ret []*graphql.FingerprintQueryInput

	for _, f := range fps {
		var i = &graphql.FingerprintQueryInput{}
		switch f.Type {
		case models.FingerprintTypeMD5:
			i.Algorithm = graphql.FingerprintAlgorithmMd5
			i.Hash = f.String()
		case models.FingerprintTypeOshash:
			i.Algorithm = graphql.FingerprintAlgorithmOshash
			i.Hash = f.String()
		case models.FingerprintTypePhash:
			i.Algorithm = graphql.FingerprintAlgorithmPhash
			i.Hash = utils.PhashToString(f.Int64())
		default:
			continue
		}

		if !i.Algorithm.IsValid() {
			continue
		}

		ret = append(ret, i)
	}

	return ret
}

func (c Client) findScenesByFingerprints(ctx context.Context, scenes [][]*graphql.FingerprintQueryInput) ([][]*models.ScrapedScene, error) {
	var results [][]*models.ScrapedScene

	// filter out nils
	var validScenes [][]*graphql.FingerprintQueryInput
	for _, s := range scenes {
		if len(s) > 0 {
			validScenes = append(validScenes, s)
		}
	}

	var ignoredTags []string

	for i := 0; i < len(validScenes); i += 40 {
		end := i + 40
		if end > len(validScenes) {
			end = len(validScenes)
		}
		scenes, err := c.client.FindScenesBySceneFingerprints(ctx, validScenes[i:end])

		if err != nil {
			return nil, err
		}

		for _, sceneFragments := range scenes.FindScenesBySceneFingerprints {
			var sceneResults []*models.ScrapedScene
			for _, scene := range sceneFragments {
				ss, err := c.sceneFragmentToScrapedScene(ctx, scene)
				if err != nil {
					return nil, err
				}

				var thisIgnoredTags []string
				ss.Tags, thisIgnoredTags = scraper.FilterTags(c.excludeTagRE, ss.Tags)
				ignoredTags = sliceutil.AppendUniques(ignoredTags, thisIgnoredTags)

				sceneResults = append(sceneResults, ss)
			}
			results = append(results, sceneResults)
		}
	}

	scraper.LogIgnoredTags(ignoredTags)

	// repopulate the results to be the same order as the input
	ret := make([][]*models.ScrapedScene, len(scenes))
	upTo := 0

	for i, v := range scenes {
		if len(v) > 0 {
			ret[i] = results[upTo]
			upTo++
		}
	}

	return ret, nil
}

func (c Client) sceneFragmentToScrapedScene(ctx context.Context, s *graphql.SceneFragment) (*models.ScrapedScene, error) {
	stashID := s.ID

	ss := &models.ScrapedScene{
		Title:        s.Title,
		Code:         s.Code,
		Date:         s.Date,
		Details:      s.Details,
		Director:     s.Director,
		URL:          findURL(s.Urls, "STUDIO"),
		Duration:     s.Duration,
		RemoteSiteID: &stashID,
		Fingerprints: getFingerprints(s),
		// Image
		// stash_id
	}

	for _, u := range s.Urls {
		ss.URLs = append(ss.URLs, u.URL)
	}

	if len(ss.URLs) > 0 {
		ss.URL = &ss.URLs[0]
	}

	if len(s.Images) > 0 {
		// TODO - #454 code sorts images by aspect ratio according to a wanted
		// orientation. I'm just grabbing the first for now
		ss.Image = getFirstImage(ctx, c.httpClient, s.Images)
	}

	ss.URLs = make([]string, len(s.Urls))
	for i, u := range s.Urls {
		ss.URLs[i] = u.URL
	}

	if s.Studio != nil {
		var err error
		ss.Studio, err = c.resolveStudio(ctx, s.Studio)
		if err != nil {
			return nil, err
		}
	}

	for _, p := range s.Performers {
		sp := performerFragmentToScrapedPerformer(*p.Performer)
		ss.Performers = append(ss.Performers, sp)
	}

	for _, t := range s.Tags {
		st := &models.ScrapedTag{
			Name: t.Name,
		}
		ss.Tags = append(ss.Tags, st)
	}

	return ss, nil
}

func getFirstImage(ctx context.Context, client *http.Client, images []*graphql.ImageFragment) *string {
	ret, err := fetchImage(ctx, client, images[0].URL)
	if err != nil && !errors.Is(err, context.Canceled) {
		logger.Warnf("Error fetching image %s: %s", images[0].URL, err.Error())
	}

	return ret
}

func getFingerprints(scene *graphql.SceneFragment) []*models.StashBoxFingerprint {
	fingerprints := []*models.StashBoxFingerprint{}
	for _, fp := range scene.Fingerprints {
		fingerprint := models.StashBoxFingerprint{
			Algorithm: fp.Algorithm.String(),
			Hash:      fp.Hash,
			Duration:  fp.Duration,
		}
		fingerprints = append(fingerprints, &fingerprint)
	}
	return fingerprints
}

type SceneDraft struct {
	// Files, URLs, StashIDs must be loaded
	Scene *models.Scene
	// StashIDs must be loaded
	Performers []*models.Performer
	// StashIDs must be loaded
	Studio *models.Studio
	Tags   []*models.Tag
	Cover  []byte
}

func (c Client) SubmitSceneDraft(ctx context.Context, d SceneDraft) (*string, error) {
	draft := newSceneDraftInput(d, c.box.Endpoint)
	var image io.Reader

	if len(d.Cover) > 0 {
		image = bytes.NewReader(d.Cover)
	}

	var id *string
	var ret graphql.SubmitSceneDraft
	err := c.submitDraft(ctx, graphql.SubmitSceneDraftDocument, draft, image, &ret)
	id = ret.SubmitSceneDraft.ID

	return id, err

	// ret, err := c.client.SubmitSceneDraft(ctx, draft, uploadImage(image))
	// if err != nil {
	// 	return nil, err
	// }

	// id := ret.SubmitSceneDraft.ID
	// return id, nil
}

func newSceneDraftInput(d SceneDraft, endpoint string) graphql.SceneDraftInput {
	scene := d.Scene

	draft := graphql.SceneDraftInput{}

	if scene.Title != "" {
		draft.Title = &scene.Title
	}
	if scene.Code != "" {
		draft.Code = &scene.Code
	}
	if scene.Details != "" {
		draft.Details = &scene.Details
	}
	if scene.Director != "" {
		draft.Director = &scene.Director
	}
	draft.Urls = scene.URLs.List()

	if scene.Date != nil {
		v := scene.Date.String()
		draft.Date = &v
	}

	if d.Studio != nil {
		studio := d.Studio

		studioDraft := graphql.DraftEntityInput{
			Name: studio.Name,
		}

		stashIDs := studio.StashIDs.List()
		for _, stashID := range stashIDs {
			c := stashID
			if stashID.Endpoint == endpoint {
				studioDraft.ID = &c.StashID
				break
			}
		}
		draft.Studio = &studioDraft
	}

	fingerprints := []*graphql.FingerprintInput{}

	for _, f := range scene.Files.List() {
		duration := f.Duration

		if duration != 0 {
			fingerprints = appendFingerprintsUnique(fingerprints, fileFingerprintsToInputGraphQL(f.Fingerprints, int(duration))...)
		}
	}
	draft.Fingerprints = fingerprints

	scenePerformers := d.Performers

	inputPerformers := []*graphql.DraftEntityInput{}
	for _, p := range scenePerformers {
		performerDraft := graphql.DraftEntityInput{
			Name: p.Name,
		}

		stashIDs := p.StashIDs.List()
		for _, stashID := range stashIDs {
			c := stashID
			if stashID.Endpoint == endpoint {
				performerDraft.ID = &c.StashID
				break
			}
		}

		inputPerformers = append(inputPerformers, &performerDraft)
	}
	draft.Performers = inputPerformers

	var tags []*graphql.DraftEntityInput
	sceneTags := d.Tags
	for _, tag := range sceneTags {
		tags = append(tags, &graphql.DraftEntityInput{Name: tag.Name})
	}
	draft.Tags = tags

	stashIDs := scene.StashIDs.List()
	var stashID *string
	for _, v := range stashIDs {
		if v.Endpoint == endpoint {
			vv := v.StashID
			stashID = &vv
			break
		}
	}
	draft.ID = stashID

	return draft
}

func fileFingerprintsToInputGraphQL(fps models.Fingerprints, duration int) []*graphql.FingerprintInput {
	var ret []*graphql.FingerprintInput

	for _, f := range fps {
		var i = &graphql.FingerprintInput{
			Duration: duration,
		}
		switch f.Type {
		case models.FingerprintTypeMD5:
			i.Algorithm = graphql.FingerprintAlgorithmMd5
			i.Hash = f.String()
		case models.FingerprintTypeOshash:
			i.Algorithm = graphql.FingerprintAlgorithmOshash
			i.Hash = f.String()
		case models.FingerprintTypePhash:
			i.Algorithm = graphql.FingerprintAlgorithmPhash
			i.Hash = utils.PhashToString(f.Int64())
		default:
			continue
		}

		if !i.Algorithm.IsValid() {
			continue
		}

		ret = appendFingerprintUnique(ret, i)
	}

	return ret
}

func (c Client) SubmitFingerprints(ctx context.Context, scenes []*models.Scene) (bool, error) {
	endpoint := c.box.Endpoint

	var fingerprints []graphql.FingerprintSubmission

	for _, scene := range scenes {
		stashIDs := scene.StashIDs.List()
		sceneStashID := ""
		for _, stashID := range stashIDs {
			if stashID.Endpoint == endpoint {
				sceneStashID = stashID.StashID
			}
		}

		if sceneStashID == "" {
			continue
		}

		for _, f := range scene.Files.List() {
			duration := f.Duration

			if duration == 0 {
				continue
			}

			fps := fileFingerprintsToInputGraphQL(f.Fingerprints, int(duration))
			for _, fp := range fps {
				fingerprints = append(fingerprints, graphql.FingerprintSubmission{
					SceneID:     sceneStashID,
					Fingerprint: fp,
				})
			}
		}
	}

	return c.submitFingerprints(ctx, fingerprints)
}

func (c Client) submitFingerprints(ctx context.Context, fingerprints []graphql.FingerprintSubmission) (bool, error) {
	for _, fingerprint := range fingerprints {
		_, err := c.client.SubmitFingerprint(ctx, fingerprint)
		if err != nil {
			return false, err
		}
	}

	return true, nil
}

func appendFingerprintUnique(v []*graphql.FingerprintInput, toAdd *graphql.FingerprintInput) []*graphql.FingerprintInput {
	for _, vv := range v {
		if vv.Algorithm == toAdd.Algorithm && vv.Hash == toAdd.Hash {
			return v
		}
	}

	return append(v, toAdd)
}

func appendFingerprintsUnique(v []*graphql.FingerprintInput, toAdd ...*graphql.FingerprintInput) []*graphql.FingerprintInput {
	for _, a := range toAdd {
		v = appendFingerprintUnique(v, a)
	}

	return v
}