From 9dacad70a1e33a8b81ff66295d53a1324b51af97 Mon Sep 17 00:00:00 2001
From: bnkai <48220860+bnkai@users.noreply.github.com>
Date: Thu, 19 Mar 2020 03:36:00 +0200
Subject: [PATCH] Autoassociate galleries to scenes when scanning (#405)

---
 pkg/manager/manager_tasks.go | 37 +++++++++++++++++++++++-
 pkg/manager/task_scan.go     | 54 +++++++++++++++++++++++++++++++++++-
 2 files changed, 89 insertions(+), 2 deletions(-)

diff --git a/pkg/manager/manager_tasks.go b/pkg/manager/manager_tasks.go
index 169372036..34286eda0 100644
--- a/pkg/manager/manager_tasks.go
+++ b/pkg/manager/manager_tasks.go
@@ -12,6 +12,32 @@ import (
 	"time"
 )
 
+var extensionsToScan = []string{"zip", "m4v", "mp4", "mov", "wmv", "avi", "mpg", "mpeg", "rmvb", "rm", "flv", "asf", "mkv", "webm"}
+var extensionsGallery = []string{"zip"}
+
+func constructGlob() string { // create a sequence for glob doublestar from our extensions
+	extLen := len(extensionsToScan)
+	glb := "{"
+	for i := 0; i < extLen-1; i++ { // append extensions and commas
+		glb += extensionsToScan[i] + ","
+	}
+	if extLen >= 1 { // append last extension without comma
+		glb += extensionsToScan[extLen-1]
+	}
+	glb += "}"
+	return glb
+
+}
+
+func isGallery(pathname string) bool {
+	for _, ext := range extensionsGallery {
+		if filepath.Ext(pathname) == "."+ext {
+			return true
+		}
+	}
+	return false
+}
+
 type TaskStatus struct {
 	Status     JobStatus
 	Progress   float64
@@ -67,7 +93,7 @@ func (s *singleton) Scan(useFileMetadata bool) {
 
 		var results []string
 		for _, path := range config.GetStashPaths() {
-			globPath := filepath.Join(path, "**/*.{zip,m4v,mp4,mov,wmv,avi,mpg,mpeg,rmvb,rm,flv,asf,mkv,webm}") // TODO: Make this configurable
+			globPath := filepath.Join(path, "**/*."+constructGlob())
 			globResults, _ := doublestar.Glob(globPath)
 			results = append(results, globResults...)
 		}
@@ -96,6 +122,15 @@ func (s *singleton) Scan(useFileMetadata bool) {
 		}
 
 		logger.Info("Finished scan")
+		for _, path := range results {
+			if isGallery(path) {
+				wg.Add(1)
+				task := ScanTask{FilePath: path, UseFileMetadata: false}
+				go task.associateGallery(&wg)
+				wg.Wait()
+			}
+		}
+		logger.Info("Finished gallery association")
 	}()
 }
 
diff --git a/pkg/manager/task_scan.go b/pkg/manager/task_scan.go
index ecdb995ee..98fde4652 100644
--- a/pkg/manager/task_scan.go
+++ b/pkg/manager/task_scan.go
@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"path/filepath"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -21,7 +22,7 @@ type ScanTask struct {
 }
 
 func (t *ScanTask) Start(wg *sync.WaitGroup) {
-	if filepath.Ext(t.FilePath) == ".zip" {
+	if isGallery(t.FilePath) {
 		t.scanGallery()
 	} else {
 		t.scanScene()
@@ -60,6 +61,7 @@ func (t *ScanTask) scanGallery() {
 	} else {
 		logger.Infof("%s doesn't exist.  Creating new item...", t.FilePath)
 		currentTime := time.Now()
+
 		newGallery := models.Gallery{
 			Checksum:  checksum,
 			Path:      t.FilePath,
@@ -77,6 +79,56 @@ func (t *ScanTask) scanGallery() {
 	}
 }
 
+// associates a gallery to a scene with the same basename
+func (t *ScanTask) associateGallery(wg *sync.WaitGroup) {
+	qb := models.NewGalleryQueryBuilder()
+	gallery, _ := qb.FindByPath(t.FilePath)
+	if gallery == nil {
+		// shouldn't happen , associate is run after scan is finished
+		logger.Errorf("associate: gallery %s not found in DB", t.FilePath)
+		wg.Done()
+		return
+	}
+
+	if !gallery.SceneID.Valid { // gallery has no SceneID
+		basename := strings.TrimSuffix(t.FilePath, filepath.Ext(t.FilePath))
+		var relatedFiles []string
+		for _, ext := range extensionsToScan { // make a list of media files that can be related to the gallery
+			related := basename + "." + ext
+			if !isGallery(related) { //exclude gallery extensions from the related files
+				relatedFiles = append(relatedFiles, related)
+			}
+		}
+		for _, scenePath := range relatedFiles {
+			qbScene := models.NewSceneQueryBuilder()
+			scene, _ := qbScene.FindByPath(scenePath)
+			if scene != nil { // found related Scene
+				logger.Infof("associate: Gallery %s is related to scene: %d", t.FilePath, scene.ID)
+
+				gallery.SceneID.Int64 = int64(scene.ID)
+				gallery.SceneID.Valid = true
+
+				ctx := context.TODO()
+				tx := database.DB.MustBeginTx(ctx, nil)
+
+				_, err := qb.Update(*gallery, tx)
+				if err != nil {
+					logger.Errorf("associate: Error updating gallery sceneId %s", err)
+					_ = tx.Rollback()
+				} else if err := tx.Commit(); err != nil {
+					logger.Error(err.Error())
+				}
+
+				break // since a gallery can have only one related scene
+				// only first found is associated
+			}
+
+		}
+
+	}
+	wg.Done()
+}
+
 func (t *ScanTask) scanScene() {
 	qb := models.NewSceneQueryBuilder()
 	scene, _ := qb.FindByPath(t.FilePath)