stash/pkg/scene/import.go

package scene

import (
	"context"
	"fmt"
	"slices"
	"strings"
	"time"

	"github.com/stashapp/stash/pkg/models"
	"github.com/stashapp/stash/pkg/models/json"
	"github.com/stashapp/stash/pkg/models/jsonschema"
	"github.com/stashapp/stash/pkg/sliceutil"
	"github.com/stashapp/stash/pkg/utils"
)

type ImporterReaderWriter interface {
	models.SceneCreatorUpdater
	models.ViewHistoryWriter
	models.OHistoryWriter
	FindByFileID(ctx context.Context, fileID models.FileID) ([]*models.Scene, error)
}

type Importer struct {
	ReaderWriter        ImporterReaderWriter
	FileFinder          models.FileFinder
	StudioWriter        models.StudioFinderCreator
	GalleryFinder       models.GalleryFinder
	PerformerWriter     models.PerformerFinderCreator
	GroupWriter         models.GroupFinderCreator
	TagWriter           models.TagFinderCreator
	Input               jsonschema.Scene
	MissingRefBehaviour models.ImportMissingRefEnum
	FileNamingAlgorithm models.HashAlgorithm

	ID             int
	scene          models.Scene
	coverImageData []byte
	viewHistory    []time.Time
	oHistory       []time.Time
}

func (i *Importer) PreImport(ctx context.Context) error {
	i.scene = i.sceneJSONToScene(i.Input)

	if err := i.populateFiles(ctx); err != nil {
		return err
	}

	if err := i.populateStudio(ctx); err != nil {
		return err
	}

	if err := i.populateGalleries(ctx); err != nil {
		return err
	}

	if err := i.populatePerformers(ctx); err != nil {
		return err
	}

	if err := i.populateTags(ctx); err != nil {
		return err
	}

	if err := i.populateGroups(ctx); err != nil {
		return err
	}

	var err error
	if len(i.Input.Cover) > 0 {
		i.coverImageData, err = utils.ProcessBase64Image(i.Input.Cover)
		if err != nil {
			return fmt.Errorf("invalid cover image: %v", err)
		}
	}

	i.populateViewHistory()
	i.populateOHistory()

	return nil
}

func (i *Importer) sceneJSONToScene(sceneJSON jsonschema.Scene) models.Scene {
	newScene := models.Scene{
		Title:        sceneJSON.Title,
		Code:         sceneJSON.Code,
		Details:      sceneJSON.Details,
		Director:     sceneJSON.Director,
		PerformerIDs: models.NewRelatedIDs([]int{}),
		TagIDs:       models.NewRelatedIDs([]int{}),
		GalleryIDs:   models.NewRelatedIDs([]int{}),
		Groups:       models.NewRelatedGroups([]models.GroupsScenes{}),
		StashIDs:     models.NewRelatedStashIDs(sceneJSON.StashIDs),
	}

	if len(sceneJSON.URLs) > 0 {
		newScene.URLs = models.NewRelatedStrings(sceneJSON.URLs)
	} else if sceneJSON.URL != "" {
		newScene.URLs = models.NewRelatedStrings([]string{sceneJSON.URL})
	}

	if sceneJSON.Date != "" {
		d, err := models.ParseDate(sceneJSON.Date)
		if err == nil {
			newScene.Date = &d
		}
	}
	if sceneJSON.Rating != 0 {
		newScene.Rating = &sceneJSON.Rating
	}

	newScene.Organized = sceneJSON.Organized
	newScene.CreatedAt = sceneJSON.CreatedAt.GetTime()
	newScene.UpdatedAt = sceneJSON.UpdatedAt.GetTime()
	newScene.ResumeTime = sceneJSON.ResumeTime
	newScene.PlayDuration = sceneJSON.PlayDuration

	return newScene
}

func getHistory(historyJSON []json.JSONTime, count int, last json.JSONTime, createdAt json.JSONTime) []time.Time {
	var ret []time.Time

	if len(historyJSON) > 0 {
		for _, d := range historyJSON {
			ret = append(ret, d.GetTime())
		}
	} else if count > 0 {
		createdAt := createdAt.GetTime()
		for j := 0; j < count; j++ {
			t := createdAt
			if j+1 == count && !last.IsZero() {
				// last one, use last play date
				t = last.GetTime()
			}
			ret = append(ret, t)
		}
	}

	return ret
}

func (i *Importer) populateViewHistory() {
	i.viewHistory = getHistory(
		i.Input.PlayHistory,
		i.Input.PlayCount,
		i.Input.LastPlayedAt,
		i.Input.CreatedAt,
	)
}

func (i *Importer) populateOHistory() {
	i.oHistory = getHistory(
		i.Input.OHistory,
		i.Input.OCounter,
		i.Input.CreatedAt, // no last o count date
		i.Input.CreatedAt,
	)
}

func (i *Importer) populateFiles(ctx context.Context) error {
	files := make([]*models.VideoFile, 0)

	for _, ref := range i.Input.Files {
		path := ref
		f, err := i.FileFinder.FindByPath(ctx, path, true)
		if err != nil {
			return fmt.Errorf("error finding file: %w", err)
		}

		if f == nil {
			return fmt.Errorf("scene file '%s' not found", path)
		} else {
			files = append(files, f.(*models.VideoFile))
		}
	}

	i.scene.Files = models.NewRelatedVideoFiles(files)

	return nil
}

func (i *Importer) populateStudio(ctx context.Context) error {
	if i.Input.Studio != "" {
		studio, err := i.StudioWriter.FindByName(ctx, i.Input.Studio, false)
		if err != nil {
			return fmt.Errorf("error finding studio by name: %v", err)
		}

		if studio == nil {
			if i.MissingRefBehaviour == models.ImportMissingRefEnumFail {
				return fmt.Errorf("scene studio '%s' not found", i.Input.Studio)
			}

			if i.MissingRefBehaviour == models.ImportMissingRefEnumIgnore {
				return nil
			}

			if i.MissingRefBehaviour == models.ImportMissingRefEnumCreate {
				studioID, err := i.createStudio(ctx, i.Input.Studio)
				if err != nil {
					return err
				}
				i.scene.StudioID = &studioID
			}
		} else {
			i.scene.StudioID = &studio.ID
		}
	}

	return nil
}

func (i *Importer) createStudio(ctx context.Context, name string) (int, error) {
	newStudio := models.NewStudio()
	newStudio.Name = name

	err := i.StudioWriter.Create(ctx, &newStudio)
	if err != nil {
		return 0, err
	}

	return newStudio.ID, nil
}

func (i *Importer) locateGallery(ctx context.Context, ref jsonschema.GalleryRef) (*models.Gallery, error) {
	var galleries []*models.Gallery
	var err error
	switch {
	case ref.FolderPath != "":
		galleries, err = i.GalleryFinder.FindByPath(ctx, ref.FolderPath)
	case len(ref.ZipFiles) > 0:
		for _, p := range ref.ZipFiles {
			galleries, err = i.GalleryFinder.FindByPath(ctx, p)
			if err != nil {
				break
			}

			if len(galleries) > 0 {
				break
			}
		}
	case ref.Title != "":
		galleries, err = i.GalleryFinder.FindUserGalleryByTitle(ctx, ref.Title)
	}

	var ret *models.Gallery
	if len(galleries) > 0 {
		ret = galleries[0]
	}

	return ret, err
}

func (i *Importer) populateGalleries(ctx context.Context) error {
	for _, ref := range i.Input.Galleries {
		gallery, err := i.locateGallery(ctx, ref)
		if err != nil {
			return err
		}

		if gallery == nil {
			if i.MissingRefBehaviour == models.ImportMissingRefEnumFail {
				return fmt.Errorf("scene gallery '%s' not found", ref.String())
			}

			// we don't create galleries - just ignore
		} else {
			i.scene.GalleryIDs.Add(gallery.ID)
		}
	}

	return nil
}

func (i *Importer) populatePerformers(ctx context.Context) error {
	if len(i.Input.Performers) > 0 {
		names := i.Input.Performers
		performers, err := i.PerformerWriter.FindByNames(ctx, names, false)
		if err != nil {
			return err
		}

		var pluckedNames []string
		for _, performer := range performers {
			if performer.Name == "" {
				continue
			}
			pluckedNames = append(pluckedNames, performer.Name)
		}

		missingPerformers := sliceutil.Filter(names, func(name string) bool {
			return !slices.Contains(pluckedNames, name)
		})

		if len(missingPerformers) > 0 {
			if i.MissingRefBehaviour == models.ImportMissingRefEnumFail {
				return fmt.Errorf("scene performers [%s] not found", strings.Join(missingPerformers, ", "))
			}

			if i.MissingRefBehaviour == models.ImportMissingRefEnumCreate {
				createdPerformers, err := i.createPerformers(ctx, missingPerformers)
				if err != nil {
					return fmt.Errorf("error creating scene performers: %v", err)
				}

				performers = append(performers, createdPerformers...)
			}

			// ignore if MissingRefBehaviour set to Ignore
		}

		for _, p := range performers {
			i.scene.PerformerIDs.Add(p.ID)
		}
	}

	return nil
}

func (i *Importer) createPerformers(ctx context.Context, names []string) ([]*models.Performer, error) {
	var ret []*models.Performer
	for _, name := range names {
		newPerformer := models.NewPerformer()
		newPerformer.Name = name

		err := i.PerformerWriter.Create(ctx, &models.CreatePerformerInput{
			Performer: &newPerformer,
		})
		if err != nil {
			return nil, err
		}

		ret = append(ret, &newPerformer)
	}

	return ret, nil
}

func (i *Importer) populateGroups(ctx context.Context) error {
	if len(i.Input.Groups) > 0 {
		for _, inputGroup := range i.Input.Groups {
			group, err := i.GroupWriter.FindByName(ctx, inputGroup.GroupName, false)
			if err != nil {
				return fmt.Errorf("error finding scene group: %v", err)
			}

			var groupID int
			if group == nil {
				if i.MissingRefBehaviour == models.ImportMissingRefEnumFail {
					return fmt.Errorf("scene group [%s] not found", inputGroup.GroupName)
				}

				if i.MissingRefBehaviour == models.ImportMissingRefEnumCreate {
					groupID, err = i.createGroup(ctx, inputGroup.GroupName)
					if err != nil {
						return fmt.Errorf("error creating scene group: %v", err)
					}
				}

				// ignore if MissingRefBehaviour set to Ignore
				if i.MissingRefBehaviour == models.ImportMissingRefEnumIgnore {
					continue
				}
			} else {
				groupID = group.ID
			}

			toAdd := models.GroupsScenes{
				GroupID: groupID,
			}

			if inputGroup.SceneIndex != 0 {
				index := inputGroup.SceneIndex
				toAdd.SceneIndex = &index
			}

			i.scene.Groups.Add(toAdd)
		}
	}

	return nil
}

func (i *Importer) createGroup(ctx context.Context, name string) (int, error) {
	newGroup := models.NewGroup()
	newGroup.Name = name

	err := i.GroupWriter.Create(ctx, &newGroup)
	if err != nil {
		return 0, err
	}

	return newGroup.ID, nil
}

func (i *Importer) populateTags(ctx context.Context) error {
	if len(i.Input.Tags) > 0 {

		tags, err := importTags(ctx, i.TagWriter, i.Input.Tags, i.MissingRefBehaviour)
		if err != nil {
			return err
		}

		for _, p := range tags {
			i.scene.TagIDs.Add(p.ID)
		}
	}

	return nil
}

func (i *Importer) addViewHistory(ctx context.Context) error {
	if len(i.viewHistory) > 0 {
		_, err := i.ReaderWriter.AddViews(ctx, i.ID, i.viewHistory)
		if err != nil {
			return fmt.Errorf("error adding view date: %v", err)
		}
	}

	return nil
}

func (i *Importer) addOHistory(ctx context.Context) error {
	if len(i.oHistory) > 0 {
		_, err := i.ReaderWriter.AddO(ctx, i.ID, i.oHistory)
		if err != nil {
			return fmt.Errorf("error adding o date: %v", err)
		}
	}

	return nil
}

func (i *Importer) PostImport(ctx context.Context, id int) error {
	if len(i.coverImageData) > 0 {
		if err := i.ReaderWriter.UpdateCover(ctx, id, i.coverImageData); err != nil {
			return fmt.Errorf("error setting scene images: %v", err)
		}
	}

	// add histories
	if err := i.addViewHistory(ctx); err != nil {
		return err
	}

	if err := i.addOHistory(ctx); err != nil {
		return err
	}

	return nil
}

func (i *Importer) Name() string {
	if i.Input.Title != "" {
		return i.Input.Title
	}

	if len(i.Input.Files) > 0 {
		return i.Input.Files[0]
	}

	return ""
}

func (i *Importer) FindExistingID(ctx context.Context) (*int, error) {
	var existing []*models.Scene
	var err error

	for _, f := range i.scene.Files.List() {
		existing, err = i.ReaderWriter.FindByFileID(ctx, f.ID)
		if err != nil {
			return nil, err
		}

		if len(existing) > 0 {
			id := existing[0].ID
			return &id, nil
		}
	}

	return nil, nil
}

func (i *Importer) Create(ctx context.Context) (*int, error) {
	var fileIDs []models.FileID
	for _, f := range i.scene.Files.List() {
		fileIDs = append(fileIDs, f.Base().ID)
	}
	if err := i.ReaderWriter.Create(ctx, &i.scene, fileIDs); err != nil {
		return nil, fmt.Errorf("error creating scene: %v", err)
	}

	id := i.scene.ID
	i.ID = id
	return &id, nil
}

func (i *Importer) Update(ctx context.Context, id int) error {
	scene := i.scene
	scene.ID = id
	i.ID = id
	if err := i.ReaderWriter.Update(ctx, &scene); err != nil {
		return fmt.Errorf("error updating existing scene: %v", err)
	}

	return nil
}

func importTags(ctx context.Context, tagWriter models.TagFinderCreator, names []string, missingRefBehaviour models.ImportMissingRefEnum) ([]*models.Tag, error) {
	tags, err := tagWriter.FindByNames(ctx, names, false)
	if err != nil {
		return nil, err
	}

	var pluckedNames []string
	for _, tag := range tags {
		pluckedNames = append(pluckedNames, tag.Name)
	}

	missingTags := sliceutil.Filter(names, func(name string) bool {
		return !slices.Contains(pluckedNames, name)
	})

	if len(missingTags) > 0 {
		if missingRefBehaviour == models.ImportMissingRefEnumFail {
			return nil, fmt.Errorf("tags [%s] not found", strings.Join(missingTags, ", "))
		}

		if missingRefBehaviour == models.ImportMissingRefEnumCreate {
			createdTags, err := createTags(ctx, tagWriter, missingTags)
			if err != nil {
				return nil, fmt.Errorf("error creating tags: %v", err)
			}

			tags = append(tags, createdTags...)
		}

		// ignore if MissingRefBehaviour set to Ignore
	}

	return tags, nil
}

func createTags(ctx context.Context, tagWriter models.TagCreator, names []string) ([]*models.Tag, error) {
	var ret []*models.Tag
	for _, name := range names {
		newTag := models.NewTag()
		newTag.Name = name

		err := tagWriter.Create(ctx, &newTag)
		if err != nil {
			return nil, err
		}

		ret = append(ret, &newTag)
	}

	return ret, nil
}