zurg/internal/torrent/manager.go

package torrent

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

	"github.com/debridmediamanager/zurg/internal/config"
	"github.com/debridmediamanager/zurg/internal/fs"
	"github.com/debridmediamanager/zurg/pkg/logutil"
	"github.com/debridmediamanager/zurg/pkg/realdebrid"
	mapset "github.com/deckarep/golang-set/v2"
	cmap "github.com/orcaman/concurrent-map/v2"
	"github.com/panjf2000/ants/v2"
	"gopkg.in/vansante/go-ffprobe.v2"
)

const (
	INT_ALL = "int__all__"
)

type TorrentManager struct {
	requiredVersion string

	Config     config.ConfigInterface
	api        *realdebrid.RealDebrid
	workerPool *ants.Pool
	log        *logutil.Logger
	repairLog  *logutil.Logger

	DirectoryMap  cmap.ConcurrentMap[string, cmap.ConcurrentMap[string, *Torrent]] // directory -> accessKey -> Torrent
	DownloadMap   cmap.ConcurrentMap[string, *realdebrid.Download]
	UnrestrictMap cmap.ConcurrentMap[string, *realdebrid.Download]

	RootNode *fs.FileNode

	RefreshKillSwitch chan struct{}
	RepairKillSwitch  chan struct{}
	RemountTrigger    chan struct{}

	latestState *LibraryState

	repairTrigger   chan *Torrent
	repairQueue     mapset.Set[*Torrent]
	repairRunning   bool
	repairRunningMu sync.Mutex

	ImmediateBin mapset.Set[string]
	OnceDoneBin  mapset.Set[string]
}

// NewTorrentManager creates a new torrent manager
// it will fetch all torrents and their info in the background
// and store them in-memory and cached in files
func NewTorrentManager(cfg config.ConfigInterface, api *realdebrid.RealDebrid, workerPool *ants.Pool, log, repairLog *logutil.Logger) *TorrentManager {
	t := &TorrentManager{
		requiredVersion: "0.10.0",

		Config:     cfg,
		api:        api,
		workerPool: workerPool,
		log:        log,
		repairLog:  repairLog,

		DirectoryMap:  cmap.New[cmap.ConcurrentMap[string, *Torrent]](),
		DownloadMap:   cmap.New[*realdebrid.Download](),
		UnrestrictMap: cmap.New[*realdebrid.Download](),

		RootNode: fs.NewFileNode("root", true),

		RefreshKillSwitch: make(chan struct{}, 1),
		RepairKillSwitch:  make(chan struct{}, 1),
		RemountTrigger:    make(chan struct{}, 1),

		latestState: &LibraryState{log: log},
	}

	t.initializeBins()
	t.initializeDirectoryMaps()
	var wg sync.WaitGroup
	wg.Add(2)
	t.workerPool.Submit(func() {
		t.refreshTorrents()
		wg.Done()
		t.setNewLatestState(t.getCurrentState())
		t.StartRefreshJob()
		t.StartRepairJob()
		t.TriggerRepair(nil)
	})
	t.workerPool.Submit(func() {
		t.mountNewDownloads()
		wg.Done()
		t.StartDownloadsJob()
	})
	t.workerPool.Submit(func() {
		wg.Wait()
		t.log.Info("Applying media info details to all torrents")
		allTorrents, _ := t.DirectoryMap.Get(INT_ALL)
		allTorrents.IterCb(func(_ string, torrent *Torrent) {
			t.applyMediaInfoDetails(torrent)
			t.writeTorrentToFile(torrent)
			t.log.Debugf("Applied media info details to torrent %s", t.GetKey(torrent))
		})
	})

	return t
}

// proxy function
func (t *TorrentManager) UnrestrictLinkUntilOk(link string, checkFirstByte bool) *realdebrid.Download {
	isRealDebrid := strings.HasPrefix(link, "https://real-debrid.com/d/")
	if isRealDebrid && t.UnrestrictMap.Has(link[0:39]) {
		ret, _ := t.UnrestrictMap.Get(link[0:39])
		return ret
	} else if !isRealDebrid && t.UnrestrictMap.Has(link) {
		ret, _ := t.UnrestrictMap.Get(link)
		return ret
	}
	ret, err := t.api.UnrestrictLink(link, checkFirstByte)
	if err != nil {
		t.log.Warnf("Cannot unrestrict link %s: %v", link, err)
		return nil
	}
	if isRealDebrid {
		t.UnrestrictMap.Set(ret.Link[0:39], ret)
	} else {
		t.UnrestrictMap.Set(ret.Link, ret)
	}
	return ret
}

func (t *TorrentManager) UnrestrictFileUntilOk(file *File, checkFirstByte bool) *realdebrid.Download {
	if !file.State.Is("ok_file") {
		return nil
	}
	return t.UnrestrictLinkUntilOk(file.Link, checkFirstByte)
}

func (t *TorrentManager) GetKey(torrent *Torrent) string {
	if !t.Config.ShouldIgnoreRenames() && torrent.Rename != "" {
		return torrent.Rename
	}
	if t.Config.EnableRetainRDTorrentName() {
		return torrent.Name
	}
	// drop the extension from the name
	if t.Config.EnableRetainFolderNameExtension() && strings.Contains(torrent.Name, torrent.OriginalName) {
		return torrent.Name
	} else {
		ret := strings.TrimSuffix(torrent.OriginalName, ".mp4")
		ret = strings.TrimSuffix(ret, ".mkv")
		return ret
	}
}

func (t *TorrentManager) GetPath(file *File) string {
	if !t.Config.ShouldIgnoreRenames() && file.Rename != "" {
		return file.Rename
	}
	filename := filepath.Base(file.Path)
	return filename
}

/// torrent functions

func (t *TorrentManager) getTorrentFiles(parentDir string) mapset.Set[string] {
	files, err := filepath.Glob(parentDir + "/*.zurgtorrent")
	if err != nil {
		t.log.Warnf("Cannot get files in %s directory: %v", parentDir, err)
		return nil
	}
	return mapset.NewSet[string](files...)
}

func (t *TorrentManager) writeTorrentToFile(torrent *Torrent) {
	filePath := "data/" + t.GetKey(torrent) + ".zurgtorrent"
	file, err := os.Create(filePath)
	if err != nil {
		t.log.Warnf("Cannot create file %s: %v", filePath, err)
		return
	}
	defer file.Close()

	torrent.Version = t.requiredVersion

	jsonData, err := json.Marshal(torrent)
	if err != nil {
		t.log.Warnf("Cannot marshal torrent: %v", err)
		return
	}

	if _, err := file.Write(jsonData); err != nil {
		t.log.Warnf("Cannot write to file %s: %v", filePath, err)
		return
	}

	t.log.Debugf("Saved torrent %s to file", t.GetKey(torrent), torrent.Hash)
}

func (t *TorrentManager) applyMediaInfoDetails(torrent *Torrent) {
	torrent.SelectedFiles.IterCb(func(_ string, file *File) {
		if file.MediaInfo != nil {
			return
		}
		unrestrict := t.UnrestrictFileUntilOk(file, false)
		if unrestrict == nil {
			file.State.Event(context.Background(), "break_file")
			return
		}
		ctx, cancelFn := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancelFn()
		data, err := ffprobe.ProbeURL(ctx, unrestrict.Download)
		if err != nil {
			t.log.Warnf("Cannot probe file %s: %v", file.Path, err)
			file.State.Event(context.Background(), "break_file")
			return
		}
		file.MediaInfo = data
	})
}

func (t *TorrentManager) readTorrentFromFile(filePath string) *Torrent {
	file, err := os.Open(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return nil
	}
	defer file.Close()
	jsonData, err := io.ReadAll(file)
	if err != nil {
		return nil
	}
	var torrent *Torrent
	if err := json.Unmarshal(jsonData, &torrent); err != nil {
		return nil
	}
	if torrent.Version != t.requiredVersion {
		return nil
	}
	return torrent
}

func (t *TorrentManager) deleteTorrentFile(hash string) {
	filePath := "data/" + hash + ".zurgtorrent"
	_ = os.Remove(filePath)
}

/// end torrent functions

/// info functions

func (t *TorrentManager) getInfoFiles() mapset.Set[string] {
	files, err := filepath.Glob("data/*.zurginfo")
	if err != nil {
		t.log.Warnf("Cannot get files in data directory: %v", err)
		return nil
	}
	return mapset.NewSet[string](files...)
}

func (t *TorrentManager) writeInfoToFile(info *realdebrid.TorrentInfo) {
	filePath := "data/" + info.ID + ".zurginfo"
	file, err := os.Create(filePath)
	if err != nil {
		t.log.Warnf("Cannot create info file %s: %v", filePath, err)
		return
	}
	defer file.Close()

	jsonData, err := json.Marshal(info)
	if err != nil {
		t.log.Warnf("Cannot marshal torrent info: %v", err)
		return
	}

	if _, err := file.Write(jsonData); err != nil {
		t.log.Warnf("Cannot write to info file %s: %v", filePath, err)
		return
	}

	// t.log.Debugf("Saved torrent %s to info file", info.ID)
}

func (t *TorrentManager) readInfoFromFile(torrentID string) *realdebrid.TorrentInfo {
	filePath := "data/" + torrentID + ".zurginfo"
	file, err := os.Open(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return nil
	}
	defer file.Close()
	jsonData, err := io.ReadAll(file)
	if err != nil {
		return nil
	}
	var info *realdebrid.TorrentInfo
	if err := json.Unmarshal(jsonData, &info); err != nil {
		return nil
	}
	return info
}

func (t *TorrentManager) deleteInfoFile(torrentID string) {
	filePath := "data/" + torrentID + ".zurginfo"
	_ = os.Remove(filePath)
}

/// end info functions

func (t *TorrentManager) mountNewDownloads() {
	downloads := t.api.GetDownloads()
	for _, download := range downloads {
		isRealDebrid := strings.HasPrefix(download.Link, "https://real-debrid.com/d/")
		if isRealDebrid && !t.UnrestrictMap.Has(download.Link[0:39]) {
			t.UnrestrictMap.Set(download.Link[0:39], &download)
		} else if !isRealDebrid {
			if !t.UnrestrictMap.Has(download.Link) {
				t.UnrestrictMap.Set(download.Link, &download)
			}
			filename := filepath.Base(download.Filename)
			t.DownloadMap.Set(filename, &download)
		}
	}
}

func (t *TorrentManager) StartDownloadsJob() {
	_ = t.workerPool.Submit(func() {
		remountTicker := time.NewTicker(time.Duration(t.Config.GetDownloadsEveryMins()) * time.Minute)
		defer remountTicker.Stop()

		for {
			select {
			case <-remountTicker.C:
				t.mountNewDownloads()
			case <-t.RemountTrigger:
				t.DownloadMap.Clear()
				t.mountNewDownloads()
			}
		}
	})
}

func (t *TorrentManager) initializeDirectoryMaps() {
	// create internal directories
	t.DirectoryMap.Set(INT_ALL, cmap.New[*Torrent]()) // key is GetAccessKey()
	// create directory maps
	for _, directory := range t.Config.GetDirectories() {
		t.DirectoryMap.Set(directory, cmap.New[*Torrent]())
		// t.RootNode.AddChild(fs.NewFileNode(directory, true))
	}
}