zurg/internal/torrent/refresh.go

package torrent

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/debridmediamanager/zurg/internal/config"
	"github.com/debridmediamanager/zurg/pkg/realdebrid"
	"github.com/debridmediamanager/zurg/pkg/utils"
	mapset "github.com/deckarep/golang-set/v2"
	cmap "github.com/orcaman/concurrent-map/v2"
)

func (t *TorrentManager) refreshTorrents() []string {
	instances, _, err := t.api.GetTorrents(false)
	if err != nil {
		t.log.Warnf("Cannot get torrents: %v", err)
		return nil
	}

	var wg sync.WaitGroup
	var mergeChan = make(chan *Torrent, len(instances))

	updatedPaths := mapset.NewSet[string]()

	allTorrents, _ := t.DirectoryMap.Get(INT_ALL)

	freshIDs := mapset.NewSet[string]()
	freshAccessKeys := mapset.NewSet[string]()

	t.getTorrentFiles("dump").Each(func(filePath string) bool {
		torrent := t.readTorrentFromFile(filePath)
		if torrent != nil {
			accessKey := t.GetKey(torrent)
			t.log.Debugf("Adding dumped torrent %s", accessKey)
			allTorrents.Set(accessKey, torrent)
			t.assignDirectory(torrent, func(directory string) {
				listing, _ := t.DirectoryMap.Get(directory)
				listing.Set(accessKey, torrent)
				// note that we're not adding it to updatedPaths
			})
			freshAccessKeys.Add(accessKey) // to prevent being deleted
		}
		return false
	})

	existingHashes := mapset.NewSet[string]()
	t.getTorrentFiles("data").Each(func(path string) bool {
		path = filepath.Base(path)
		hash := strings.TrimSuffix(path, ".zurgtorrent")
		existingHashes.Add(hash)
		return false
	})

	for i := range instances {
		wg.Add(1)
		idx := i
		_ = t.workerPool.Submit(func() {
			defer wg.Done()
			if t.binImmediately(instances[idx].ID) {
				// t.log.Debugf("Skipping trashed torrent %s (id=%s)", instances[idx].Name, instances[idx].ID)
				mergeChan <- nil
				return
			}

			if instances[idx].Progress != 100 {
				t.log.Debugf("Skipping incomplete torrent %s (id=%s)", instances[idx].Name, instances[idx].ID)
				mergeChan <- nil
				return
			}

			freshIDs.Add(instances[idx].ID)

			tInfo := t.getMoreInfo(instances[idx])
			torrent := t.convertToTorrent(tInfo)
			accessKey := t.GetKey(torrent)
			freshAccessKeys.Add(accessKey)

			var forMerging *Torrent
			mainTorrent, exists := allTorrents.Get(accessKey)
			if !exists {
				allTorrents.Set(accessKey, torrent)

				t.assignDirectory(torrent, func(directory string) {
					listing, _ := t.DirectoryMap.Get(directory)
					listing.Set(accessKey, torrent)

					updatedPaths.Add(fmt.Sprintf("%s/%s", directory, accessKey))
				})
			} else if !mainTorrent.DownloadedIDs.Contains(tInfo.ID) {
				forMerging = torrent
			}

			// write to file if it is a new torrent
			if forMerging == nil && !existingHashes.Contains(tInfo.Hash) {
				t.writeTorrentToFile(torrent)
			}

			mergeChan <- forMerging
		})
	}

	wg.Wait()
	close(mergeChan)

	for torrent := range mergeChan {
		if torrent == nil {
			continue
		}

		accessKey := t.GetKey(torrent)
		existing, ok := allTorrents.Get(accessKey)
		if !ok {
			t.log.Warnf("Cannot merge %s", accessKey)
			continue
		}
		mainTorrent := t.mergeTorrents(existing, torrent)
		allTorrents.Set(accessKey, mainTorrent)
		t.writeTorrentToFile(mainTorrent)

		t.DirectoryMap.IterCb(func(directory string, torrents cmap.ConcurrentMap[string, *Torrent]) {
			if strings.HasPrefix(directory, "int__") {
				return
			}
			torrents.Remove(accessKey)
		})

		t.assignDirectory(mainTorrent, func(directory string) {
			listing, _ := t.DirectoryMap.Get(directory)
			listing.Set(accessKey, mainTorrent)

			updatedPaths.Add(fmt.Sprintf("%s/%s", directory, accessKey))
		})
	}

	noInfoCount := 0

	// removed torrents
	oldPlusNewKeys := mapset.NewSet[string](allTorrents.Keys()...)
	oldPlusNewKeys.Difference(freshAccessKeys).Each(func(accessKey string) bool {
		t.Delete(accessKey, false)
		return false
	})

	t.log.Infof("Compiled into %d torrents, %d were missing info", allTorrents.Count(), noInfoCount)

	existingIDs := mapset.NewSet[string]()
	t.getInfoFiles().Each(func(path string) bool {
		path = filepath.Base(path)
		torrentID := strings.TrimSuffix(path, ".info_zurg")
		if !t.binOnceDone(torrentID) {
			existingIDs.Add(torrentID)
		}
		return false
	})
	existingIDs.Difference(freshIDs).Each(func(id string) bool {
		t.log.Infof("Deleting stale info file %s", id)
		t.deleteInfoFile(id)
		return false
	})

	t.cleanupBins(freshIDs)

	return updatedPaths.ToSlice()
}

// StartRefreshJob periodically refreshes the torrents
func (t *TorrentManager) StartRefreshJob() {
	go func() {
		t.log.Debug("Starting periodic refresh job")
		refreshTicker := time.NewTicker(time.Duration(t.Config.GetRefreshEverySecs()) * time.Second)
		defer refreshTicker.Stop()

		for {
			select {
			case <-refreshTicker.C:
				checksum := t.getCurrentState()
				if t.latestState.Eq(checksum) {
					continue
				}
				t.setNewLatestState(checksum)

				updatedPaths := t.refreshTorrents()
				t.log.Info("Finished refreshing torrents")

				t.TriggerHookOnLibraryUpdate(updatedPaths)
			case <-t.RefreshKillSwitch:
				t.log.Info("Stopping periodic refresh job")
				return
			}
		}
	}()
}

func (t *TorrentManager) getMoreInfo(rdTorrent realdebrid.Torrent) *realdebrid.TorrentInfo {
	info := t.readInfoFromFile(rdTorrent.ID)
	if info == nil {
		var err error
		info, err = t.api.GetTorrentInfo(rdTorrent.ID)
		if err != nil {
			t.log.Warnf("Cannot get info for id=%s: %v", rdTorrent.ID, err)
			return nil
		}
		t.writeInfoToFile(info)
	}
	return info
}

func (t *TorrentManager) convertToTorrent(info *realdebrid.TorrentInfo) *Torrent {
	torrent := t.readTorrentFromFile("data/" + info.Hash + ".zurgtorrent")
	if torrent != nil && torrent.DownloadedIDs.Contains(info.ID) {
		return torrent
	}

	torrent = &Torrent{
		Name:          info.Name,
		OriginalName:  info.OriginalName,
		Added:         info.Added,
		Hash:          info.Hash,
		State:         NewTorrentState("broken_torrent"),
		DownloadedIDs: mapset.NewSet[string](),
	}

	// SelectedFiles is a subset of Files with only the selected ones
	// it also has a Link field, which can be empty
	// if it is empty, it means the file is no longer available
	// Files+Links together are the same as SelectedFiles
	allFilenames := mapset.NewSet[string]()
	dupeFilenames := mapset.NewSet[string]()
	var selectedFiles []*File
	for _, file := range info.Files {
		filename := filepath.Base(file.Path)
		if allFilenames.Contains(filename) {
			dupeFilenames.Add(filename)
		} else {
			allFilenames.Add(filename)
		}
		if file.Selected == 0 {
			continue
		}
		selectedFiles = append(selectedFiles, &File{
			File:  file,
			Ended: info.Ended,
			Link:  "", // no link yet, consider it broken
			State: NewFileState("broken_file"),
		})
	}
	if len(selectedFiles) == len(info.Links) {
		// all links are still intact! good!
		for i, file := range selectedFiles {
			file.Link = info.Links[i]
			if strings.HasPrefix(file.Link, "https://real-debrid.com/d/") {
				file.Link = file.Link[0:39]
			}
			if err := file.State.Event(context.Background(), "repair_file"); err != nil {
				t.log.Errorf("Cannot repair file %s: %v", file.Path, err)
			}
		}
		torrent.UnassignedLinks = mapset.NewSet[string]()
		if err := torrent.State.Event(context.Background(), "mark_as_repaired"); err != nil {
			t.log.Errorf("Cannot repair torrent %s: %v", torrent.Hash, err)
		}
	} else {
		torrent.UnassignedLinks = mapset.NewSet[string]()
		for _, link := range info.Links {
			if strings.HasPrefix(link, "https://real-debrid.com/d/") {
				link = link[0:39]
			}
			torrent.UnassignedLinks.Add(link)
		}
	}
	torrent.SelectedFiles = cmap.New[*File]()
	for _, file := range selectedFiles {
		baseFilename := t.GetPath(file)
		// todo better handling of duplicate filenames
		if dupeFilenames.Contains(baseFilename) {
			extension := filepath.Ext(baseFilename)
			filenameNoExt := strings.TrimSuffix(baseFilename, extension)
			newName := fmt.Sprintf("%s (%d)%s", filenameNoExt, file.ID, extension)
			torrent.SelectedFiles.Set(newName, file)
		} else {
			torrent.SelectedFiles.Set(baseFilename, file)
		}
	}

	torrent.DownloadedIDs.Add(info.ID)

	return torrent
}

func (t *TorrentManager) mergeTorrents(existing, toMerge *Torrent) *Torrent {
	var newer, older *Torrent
	if existing.Added < toMerge.Added {
		newer = toMerge
		older = existing
	} else {
		newer = existing
		older = toMerge
	}

	// base of the merged torrent
	mergedTorrent := &Torrent{
		Name:          older.Name,
		OriginalName:  older.OriginalName,
		Rename:        older.Rename,
		Hash:          older.Hash,
		Added:         older.Added,
		DownloadedIDs: older.DownloadedIDs.Union(newer.DownloadedIDs),

		State: older.State,
	}

	// unrepairable reason
	reasons := mapset.NewSet[string]()
	reasons.Add(older.UnrepairableReason)
	reasons.Add(newer.UnrepairableReason)
	mergedTorrent.UnrepairableReason = strings.Join(reasons.ToSlice(), ", ")

	// selected files
	mergedTorrent.SelectedFiles = cmap.New[*File]()
	newer.SelectedFiles.IterCb(func(key string, file *File) {
		if f, ok := mergedTorrent.SelectedFiles.Get(key); !ok || !f.State.Is("ok_file") {
			mergedTorrent.SelectedFiles.Set(key, file)
		}
	})
	older.SelectedFiles.IterCb(func(key string, file *File) {
		if f, ok := mergedTorrent.SelectedFiles.Get(key); !ok || !f.State.Is("ok_file") {
			mergedTorrent.SelectedFiles.Set(key, file)
		}
	})

	// unassigned links
	mergedTorrent.UnassignedLinks = mapset.NewSet[string]()
	links := newer.UnassignedLinks.Union(older.UnassignedLinks)
	links.Each(func(link string) bool {
		found := false
		mergedTorrent.SelectedFiles.IterCb(func(key string, file *File) {
			if !found && file.Link == link {
				found = true
			}
		})
		if !found {
			mergedTorrent.UnassignedLinks.Add(link)
		}
		return false
	})

	brokenCount := 0
	okCount := 0
	mergedTorrent.SelectedFiles.IterCb(func(key string, file *File) {
		if file.State.Is("broken_file") {
			brokenCount++
		} else if file.State.Is("ok_file") {
			okCount++
		}
	})

	if brokenCount == 0 && okCount > 0 {
		if err := mergedTorrent.State.Event(context.Background(), "mark_as_repaired"); err != nil {
			t.log.Errorf("Cannot repair torrent %s: %v", t, t.GetKey(mergedTorrent), err)
		}
	}

	return mergedTorrent
}

func (t *TorrentManager) assignDirectory(tor *Torrent, cb func(string)) {
	torrentIDs := tor.DownloadedIDs.ToSlice()
	// get filenames needed for directory conditions
	var filenames []string
	var fileSizes []int64
	unplayable := true
	tor.SelectedFiles.IterCb(func(key string, file *File) {
		filenames = append(filenames, filepath.Base(file.Path))
		fileSizes = append(fileSizes, file.Bytes)
		if utils.IsPlayable(file.Path) || t.IsPlayable(file.Path) {
			unplayable = false
		}
	})

	if unplayable {
		t.markAsUnplayable(tor, "no playable files")
		return
	}

	// Map torrents to directories
	switch t.Config.GetVersion() {
	case "v1":
		configV1 := t.Config.(*config.ZurgConfigV1)
		for _, directories := range configV1.GetGroupMap() {
			for _, directory := range directories {
				if t.Config.MeetsConditions(directory, t.GetKey(tor), tor.ComputeTotalSize(), torrentIDs, filenames, fileSizes) {
					cb(directory)
					break
				}
			}
		}
	}
}

func (t *TorrentManager) IsPlayable(filePath string) bool {
	filePath = strings.ToLower(filePath)
	playableExts := t.Config.GetPlayableExtensions()
	for _, ext := range playableExts {
		if strings.HasSuffix(filePath, fmt.Sprintf(".%s", ext)) {
			return true
		}
	}
	return false
}

// initializeBins reads from bins.json and assigns values to t.trashBin and t.repairBin
func (t *TorrentManager) initializeBins() {
	if _, err := os.Stat("data/bins.json"); os.IsNotExist(err) {
		t.log.Warn("data/bins.json does not exist. Initializing empty bins.")
		t.immediateBin = mapset.NewSet[string]()
		t.onceDoneBin = mapset.NewSet[string]()
		return
	}

	fileData, err := os.ReadFile("data/bins.json")
	if err != nil {
		t.log.Errorf("Failed to read bins.json file: %v", err)
		t.immediateBin = mapset.NewSet[string]()
		t.onceDoneBin = mapset.NewSet[string]()
		return
	}

	data := map[string][]string{}

	err = json.Unmarshal(fileData, &data)
	if err != nil {
		t.log.Errorf("Failed to unmarshal bin data: %v", err)
		return
	}

	t.immediateBin = mapset.NewSet[string](data["trash_bin"]...)
	t.onceDoneBin = mapset.NewSet[string](data["repair_bin"]...)

	t.log.Debug("Successfully read bins from bins.json")
	t.log.Debugf("Bin immediately: %v", t.immediateBin.ToSlice())
	t.log.Debugf("Bin once done: %v", t.onceDoneBin.ToSlice())
}
func (t *TorrentManager) setToBinImmediately(torrentId string) {
	t.log.Debugf("Set to delete immediately: %s", torrentId)
	t.immediateBin.Add(torrentId)
	t.persistBins()
}

func (t *TorrentManager) setToBinOnceDone(torrentId string) {
	t.log.Debugf("Set to delete once completed: %s", torrentId)
	t.onceDoneBin.Add(torrentId)
	t.persistBins()
}

func (t *TorrentManager) binImmediately(torrentId string) bool {
	if t.immediateBin.Contains(torrentId) {
		if err := t.api.DeleteTorrent(torrentId); err != nil {
			t.log.Errorf("Failed to delete torrent %s: %v", torrentId, err)
			return false
		}
		t.immediateBin.Remove(torrentId)
		return true
	}
	return false
}

func (t *TorrentManager) binOnceDone(torrentId string) bool {
	if t.onceDoneBin.Contains(torrentId) {
		if err := t.api.DeleteTorrent(torrentId); err != nil {
			t.log.Errorf("Failed to delete torrent %s: %v", torrentId, err)
			return false
		}
		t.onceDoneBin.Remove(torrentId)
		return true
	}
	return false
}

func (t *TorrentManager) persistBins() {
	data := map[string]interface{}{
		"trash_bin":  t.immediateBin.ToSlice(), // Assuming trashBin is a mapset.Set[string]
		"repair_bin": t.onceDoneBin.ToSlice(),  // Assuming repairBin is a mapset.Set[string]
	}

	jsonData, err := json.Marshal(data)
	if err != nil {
		t.log.Errorf("Failed to marshal bin data: %v", err)
		return
	}

	file, err := os.Create("data/bins.json")
	if err != nil {
		t.log.Errorf("Failed to create bins.json file: %v", err)
		return
	}
	defer file.Close()

	_, err = file.Write(jsonData)
	if err != nil {
		t.log.Errorf("Failed to write to bins.json file: %v", err)
	} else {
		t.log.Debug("Successfully saved bins to file")
	}
}

func (t *TorrentManager) cleanupBins(freshIDs mapset.Set[string]) {
	t.immediateBin.Difference(freshIDs).Each(func(id string) bool {
		t.immediateBin.Remove(id)
		return false
	})
	t.onceDoneBin.Difference(freshIDs).Each(func(id string) bool {
		t.onceDoneBin.Remove(id)
		return false
	})
	t.persistBins()
}