From 42dce61529e5b023f60e945ffb4b97df0bd6ffb8 Mon Sep 17 00:00:00 2001
From: Ben Sarmiento <ben.sarmiento@forto.com>
Date: Thu, 7 Dec 2023 12:49:59 +0100
Subject: [PATCH] Fix bugs

---
 internal/torrent/delete.go  |  4 +++-
 internal/torrent/manager.go |  4 ++--
 internal/torrent/refresh.go | 19 +++++++++----------
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/internal/torrent/delete.go b/internal/torrent/delete.go
index 98e4150..374e647 100644
--- a/internal/torrent/delete.go
+++ b/internal/torrent/delete.go
@@ -21,8 +21,8 @@ func (t *TorrentManager) CheckDeletedState(torrent *Torrent) bool {
 }
 
 func (t *TorrentManager) Delete(accessKey string, deleteInRD bool) {
+	allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
 	if deleteInRD {
-		allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
 		infoCache, _ := t.DirectoryMap.Get(INT_INFO_CACHE)
 		if torrent, ok := allTorrents.Get(accessKey); ok {
 			torrent.DownloadedIDs.Each(func(id string) bool {
@@ -34,8 +34,10 @@ func (t *TorrentManager) Delete(accessKey string, deleteInRD bool) {
 			})
 		}
 	}
+	t.allAccessKeys.Remove(accessKey)
 	t.log.Infof("Removing torrent %s from zurg database", accessKey)
 	t.DirectoryMap.IterCb(func(directory string, torrents cmap.ConcurrentMap[string, *Torrent]) {
 		torrents.Remove(accessKey)
 	})
+	allTorrents.Remove(accessKey)
 }
diff --git a/internal/torrent/manager.go b/internal/torrent/manager.go
index 80991fc..f2abe70 100644
--- a/internal/torrent/manager.go
+++ b/internal/torrent/manager.go
@@ -25,7 +25,7 @@ type TorrentManager struct {
 	Api             *realdebrid.RealDebrid
 	DirectoryMap    cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
 	DownloadCache   cmap.ConcurrentMap[string, *realdebrid.Download]
-	accessKeySet    *strset.Set
+	allAccessKeys   *strset.Set
 	latestState     *LibraryState
 	requiredVersion string
 	workerPool      *ants.Pool
@@ -42,7 +42,7 @@ func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, p
 	t := &TorrentManager{
 		Config:          cfg,
 		Api:             api,
-		accessKeySet:    set.NewStringSet(),
+		allAccessKeys:   set.NewStringSet(),
 		latestState:     &initialSate,
 		requiredVersion: "06.12.2023",
 		workerPool:      p,
diff --git a/internal/torrent/refresh.go b/internal/torrent/refresh.go
index 04937f4..50dde2b 100644
--- a/internal/torrent/refresh.go
+++ b/internal/torrent/refresh.go
@@ -43,6 +43,7 @@ func (t *TorrentManager) RefreshTorrents() []string {
 		}
 		freshKeys.Add(info.AccessKey)
 		if torrent, exists := allTorrents.Get(info.AccessKey); !exists {
+			t.allAccessKeys.Add(info.AccessKey)
 			allTorrents.Set(info.AccessKey, info)
 		} else if !strset.Difference(info.DownloadedIDs, torrent.DownloadedIDs).IsEmpty() {
 			mainTorrent := t.mergeToMain(torrent, info)
@@ -50,17 +51,10 @@ func (t *TorrentManager) RefreshTorrents() []string {
 		}
 	}
 	t.log.Infof("Compiled into %d torrents, %d were missing info", allTorrents.Count(), noInfoCount)
-
-	// removed
-	strset.Difference(t.accessKeySet, freshKeys).Each(func(accessKey string) bool {
-		t.accessKeySet.Remove(accessKey)
-		t.Delete(accessKey, false)
-		return true
-	})
-	// new
 	var updatedPaths []string
-	strset.Difference(freshKeys, t.accessKeySet).Each(func(accessKey string) bool {
-		t.accessKeySet.Add(accessKey)
+	// torrents yet to be assigned in a directory
+	strset.Difference(freshKeys, t.allAccessKeys).Each(func(accessKey string) bool {
+		// assign to directories
 		tor, _ := allTorrents.Get(accessKey)
 		t.assignedDirectoryCb(tor, func(directory string) {
 			if strings.HasPrefix(directory, "int__") {
@@ -72,6 +66,11 @@ func (t *TorrentManager) RefreshTorrents() []string {
 		})
 		return true
 	})
+	// removed torrents
+	strset.Difference(t.allAccessKeys, freshKeys).Each(func(accessKey string) bool {
+		t.Delete(accessKey, false)
+		return true
+	})
 
 	t.SetNewLatestState(t.getCurrentState())