Studio Tagger (#3510)

* Studio image and parent studio support in scene tagger
* Refactor studio backend and add studio tagger
---------
Co-authored-by: WithoutPants <53250216+WithoutPants@users.noreply.github.com>

internal/identify/studio.go

@@ -2,64 +2,95 @@ package identify
 
 import (
 	"context"
-	"fmt"
-	"time"
+	"strconv"
 
+	"github.com/stashapp/stash/pkg/logger"
 	"github.com/stashapp/stash/pkg/models"
-	"github.com/stashapp/stash/pkg/utils"
+	"github.com/stashapp/stash/pkg/studio"
 )
 
-type StudioCreator interface {
-	Create(ctx context.Context, newStudio *models.Studio) error
-	UpdateStashIDs(ctx context.Context, studioID int, stashIDs []models.StashID) error
-	UpdateImage(ctx context.Context, studioID int, image []byte) error
-}
+func createMissingStudio(ctx context.Context, endpoint string, w models.StudioReaderWriter, s *models.ScrapedStudio) (*int, error) {
+	var err error
 
-func createMissingStudio(ctx context.Context, endpoint string, w StudioCreator, studio *models.ScrapedStudio) (*int, error) {
-	studioInput := scrapedToStudioInput(studio)
-	err := w.Create(ctx, &studioInput)
+	if s.Parent != nil {
+		if s.Parent.StoredID == nil {
+			// The parent needs to be created
+			newParentStudio := s.Parent.ToStudio(endpoint, nil)
+			parentImage, err := s.Parent.GetImage(ctx, nil)
+			if err != nil {
+				logger.Errorf("Failed to make parent studio from scraped studio %s: %s", s.Parent.Name, err.Error())
+				return nil, err
+			}
+
+			// Create the studio
+			err = w.Create(ctx, newParentStudio)
+			if err != nil {
+				return nil, err
+			}
+
+			// Update image table
+			if len(parentImage) > 0 {
+				if err := w.UpdateImage(ctx, newParentStudio.ID, parentImage); err != nil {
+					return nil, err
+				}
+			}
+
+			storedId := strconv.Itoa(newParentStudio.ID)
+			s.Parent.StoredID = &storedId
+		} else {
+			// The parent studio matched an existing one and the user has chosen in the UI to link and/or update it
+			existingStashIDs := getStashIDsForStudio(ctx, *s.Parent.StoredID, w)
+			studioPartial := s.Parent.ToPartial(s.Parent.StoredID, endpoint, nil, existingStashIDs)
+			parentImage, err := s.Parent.GetImage(ctx, nil)
+			if err != nil {
+				return nil, err
+			}
+
+			if err := studio.ValidateModify(ctx, *studioPartial, w); err != nil {
+				return nil, err
+			}
+
+			_, err = w.UpdatePartial(ctx, *studioPartial)
+			if err != nil {
+				return nil, err
+			}
+
+			if len(parentImage) > 0 {
+				if err := w.UpdateImage(ctx, studioPartial.ID, parentImage); err != nil {
+					return nil, err
+				}
+			}
+		}
+	}
+
+	newStudio := s.ToStudio(endpoint, nil)
+	studioImage, err := s.GetImage(ctx, nil)
 	if err != nil {
-		return nil, fmt.Errorf("error creating studio: %w", err)
+		return nil, err
 	}
 
-	// update image table
-	if studio.Image != nil && len(*studio.Image) > 0 {
-		imageData, err := utils.ReadImageFromURL(ctx, *studio.Image)
-		if err != nil {
-			return nil, err
-		}
+	err = w.Create(ctx, newStudio)
+	if err != nil {
+		return nil, err
+	}
 
-		err = w.UpdateImage(ctx, studioInput.ID, imageData)
-		if err != nil {
+	// Update image table
+	if len(studioImage) > 0 {
+		if err := w.UpdateImage(ctx, newStudio.ID, studioImage); err != nil {
 			return nil, err
 		}
 	}
 
-	if endpoint != "" && studio.RemoteSiteID != nil {
-		if err := w.UpdateStashIDs(ctx, studioInput.ID, []models.StashID{
-			{
-				Endpoint: endpoint,
-				StashID:  *studio.RemoteSiteID,
-			},
-		}); err != nil {
-			return nil, fmt.Errorf("error setting studio stash id: %w", err)
-		}
-	}
-
-	return &studioInput.ID, nil
+	return &newStudio.ID, nil
 }
 
-func scrapedToStudioInput(studio *models.ScrapedStudio) models.Studio {
-	currentTime := time.Now()
-	ret := models.Studio{
-		Name:      studio.Name,
-		CreatedAt: currentTime,
-		UpdatedAt: currentTime,
-	}
+func getStashIDsForStudio(ctx context.Context, studioID string, w models.StudioReaderWriter) []models.StashID {
+	id, _ := strconv.Atoi(studioID)
+	tempStudio := &models.Studio{ID: id}
 
-	if studio.URL != nil {
-		ret.URL = *studio.URL
+	err := tempStudio.LoadStashIDs(ctx, w)
+	if err != nil {
+		return nil
 	}
-
-	return ret
+	return tempStudio.StashIDs.List()
 }