zurg/internal/torrent/manager.go

package torrent

import (
	"encoding/gob"
	"log"
	"os"
	"strings"
	"sync"

	"github.com/debridmediamanager.com/zurg/internal/config"
	"github.com/debridmediamanager.com/zurg/pkg/realdebrid"
)

type TorrentManager struct {
	token      string
	torrents   []Torrent
	workerPool chan bool
	config     config.ConfigInterface
}

// NewTorrentManager creates a new torrent manager
// it will fetch all torrents and their info in the background
// and store them in-memory
func NewTorrentManager(token string, config config.ConfigInterface) *TorrentManager {
	handler := &TorrentManager{
		token:      token,
		workerPool: make(chan bool, 10),
		config:     config,
	}

	handler.torrents = handler.getAll()

	for _, torrent := range handler.torrents {
		go func(id string) {
			handler.workerPool <- true
			handler.getInfo(id)
			// sleep for 1 second to avoid rate limiting
			<-handler.workerPool
		}(torrent.ID)
	}

	return handler
}

func (t *TorrentManager) getAll() []Torrent {
	log.Println("Getting all torrents")
	torrents, err := realdebrid.GetTorrents(t.token)
	if err != nil {
		log.Printf("Cannot get torrents: %v\n", err.Error())
		return nil
	}
	var torrentsV2 []Torrent
	for _, torrent := range torrents {
		torrentV2 := Torrent{
			Torrent:       torrent,
			SelectedFiles: nil,
		}
		torrentsV2 = append(torrentsV2, torrentV2)
	}
	log.Printf("Fetched %d torrents", len(torrentsV2))

	version := t.config.GetVersion()
	if version == "v1" {
		configV1 := t.config.(*config.ZurgConfigV1)
		groupMap := configV1.GetGroupMap()
		for group, directories := range groupMap {
			log.Printf("Processing group %s\n", group)
			for i := range torrents {
				for _, directory := range directories {
					if configV1.MeetsConditions(directory, torrentsV2[i].ID, torrentsV2[i].Name) {
						torrentsV2[i].Directories = append(torrentsV2[i].Directories, directory)
						break
					}
				}
			}
		}
	}
	log.Println("Finished mapping to groups")

	return torrentsV2
}

func (t *TorrentManager) GetByDirectory(directory string) []Torrent {
	var torrents []Torrent
	for i := range t.torrents {
		for _, dir := range t.torrents[i].Directories {
			if dir == directory {
				torrents = append(torrents, t.torrents[i])
			}
		}
	}
	return torrents
}

func (t *TorrentManager) getInfo(torrentID string) *Torrent {
	torrentFromFile := t.readFromFile(torrentID)
	if torrentFromFile != nil {
		torrent := t.getByID(torrentID)
		if torrent != nil {
			torrent.SelectedFiles = torrentFromFile.SelectedFiles
		}
		return torrent
	}
	info, err := realdebrid.GetTorrentInfo(t.token, torrentID)
	if err != nil {
		log.Printf("Cannot get info: %v\n", err.Error())
		return nil
	}
	var selectedFiles []File
	for _, file := range info.Files {
		if file.Selected == 0 {
			continue
		}
		selectedFiles = append(selectedFiles, File{
			File: file,
			Link: "",
		})
	}
	if len(selectedFiles) != len(info.Links) {
		type Result struct {
			Filename string
			Link     string
		}

		resultsChan := make(chan Result, len(info.Links))
		var wg sync.WaitGroup

		// Limit concurrency
		sem := make(chan struct{}, 10) // e.g., 10 concurrent requests

		for _, link := range info.Links {
			wg.Add(1)
			sem <- struct{}{} // Acquire semaphore
			go func(lnk string) {
				defer wg.Done()
				defer func() { <-sem }() // Release semaphore

				unrestrictFn := func() (*realdebrid.UnrestrictResponse, error) {
					return realdebrid.UnrestrictLink(t.token, lnk)
				}
				resp := realdebrid.RetryUntilOk(unrestrictFn)
				if resp != nil {
					resultsChan <- Result{Filename: resp.Filename, Link: resp.Link}
				}
			}(link)
		}

		go func() {
			wg.Wait()
			close(resultsChan)
		}()

		for result := range resultsChan {
			for i := range selectedFiles {
				if strings.HasSuffix(selectedFiles[i].Path, result.Filename) {
					selectedFiles[i].Link = result.Link
				}
			}
		}
	} else {
		for i, link := range info.Links {
			selectedFiles[i].Link = link
		}
	}
	torrent := t.getByID(torrentID)
	if torrent != nil {
		torrent.SelectedFiles = selectedFiles
	}
	if len(torrent.SelectedFiles) > 0 {
		t.writeToFile(torrentID, torrent)
	}
	return torrent
}

func (t *TorrentManager) GetInfo(torrentID string) *Torrent {
	for i := range t.torrents {
		if t.torrents[i].ID == torrentID {
			return &t.torrents[i]
		}
	}
	return t.getInfo(torrentID)
}

func (t *TorrentManager) getByID(torrentID string) *Torrent {
	for i := range t.torrents {
		if t.torrents[i].ID == torrentID {
			return &t.torrents[i]
		}
	}
	return nil
}

func (t *TorrentManager) writeToFile(torrentID string, torrent *Torrent) {
	filePath := "data/" + torrentID + ".bin"
	file, err := os.Create(filePath)
	if err != nil {
		log.Fatalf("Failed creating file: %s", err)
		return
	}
	defer file.Close()

	dataEncoder := gob.NewEncoder(file)
	dataEncoder.Encode(torrent)
}

func (t *TorrentManager) readFromFile(torrentID string) *Torrent {
	filePath := "data/" + torrentID + ".bin"
	file, err := os.Open(filePath)
	if err != nil {
		return nil
	}
	defer file.Close()

	var torrent Torrent
	dataDecoder := gob.NewDecoder(file)
	err = dataDecoder.Decode(&torrent)
	if err != nil {
		log.Fatalf("Failed decoding file: %s", err)
		return nil
	}

	return &torrent
}