decypharr/pkg/debrid/realdebrid/realdebrid.go

package realdebrid

import (
	"fmt"
	"github.com/goccy/go-json"
	"github.com/rs/zerolog"
	"github.com/sirrobot01/debrid-blackhole/internal/config"
	"github.com/sirrobot01/debrid-blackhole/internal/logger"
	"github.com/sirrobot01/debrid-blackhole/internal/request"
	"github.com/sirrobot01/debrid-blackhole/internal/utils"
	"github.com/sirrobot01/debrid-blackhole/pkg/debrid/torrent"
	"io"
	"net/http"
	gourl "net/url"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
)

type RealDebrid struct {
	Name             string
	Host             string `json:"host"`
	APIKey           string
	DownloadUncached bool
	client           *request.Client

	MountPath   string
	logger      zerolog.Logger
	CheckCached bool
}

func (r *RealDebrid) GetName() string {
	return r.Name
}

func (r *RealDebrid) GetLogger() zerolog.Logger {
	return r.logger
}

// getTorrentFiles returns a list of torrent files from the torrent info
// validate is used to determine if the files should be validated
// if validate is false, selected files will be returned
func getTorrentFiles(t *torrent.Torrent, data TorrentInfo, validate bool) map[string]torrent.File {
	files := make(map[string]torrent.File)
	cfg := config.GetConfig()
	idx := 0
	for _, f := range data.Files {

		name := filepath.Base(f.Path)

		if validate {
			if utils.RegexMatch(utils.SAMPLEMATCH, name) {
				// Skip sample files
				continue
			}
			if !cfg.IsAllowedFile(name) {
				continue
			}
			if !cfg.IsSizeAllowed(f.Bytes) {
				continue
			}
		} else {
			if f.Selected == 0 {
				continue
			}
		}

		fileId := f.ID
		_link := ""
		if len(data.Links) > idx {
			_link = data.Links[idx]
		}

		if a, ok := t.Files[name]; ok {
			a.Link = _link
			files[name] = a
			continue
		}

		file := torrent.File{
			Name: name,
			Path: name,
			Size: f.Bytes,
			Id:   strconv.Itoa(fileId),
			Link: _link,
		}
		files[name] = file
		idx++
	}
	return files
}

func (r *RealDebrid) IsAvailable(hashes []string) map[string]bool {
	// Check if the infohashes are available in the local cache
	result := make(map[string]bool)

	// Divide hashes into groups of 100
	for i := 0; i < len(hashes); i += 200 {
		end := i + 200
		if end > len(hashes) {
			end = len(hashes)
		}

		// Filter out empty strings
		validHashes := make([]string, 0, end-i)
		for _, hash := range hashes[i:end] {
			if hash != "" {
				validHashes = append(validHashes, hash)
			}
		}

		// If no valid hashes in this batch, continue to the next batch
		if len(validHashes) == 0 {
			continue
		}

		hashStr := strings.Join(validHashes, "/")
		url := fmt.Sprintf("%s/torrents/instantAvailability/%s", r.Host, hashStr)
		req, _ := http.NewRequest(http.MethodGet, url, nil)
		resp, err := r.client.MakeRequest(req)
		if err != nil {
			r.logger.Info().Msgf("Error checking availability: %v", err)
			return result
		}
		var data AvailabilityResponse
		err = json.Unmarshal(resp, &data)
		if err != nil {
			r.logger.Info().Msgf("Error marshalling availability: %v", err)
			return result
		}
		for _, h := range hashes[i:end] {
			hosters, exists := data[strings.ToLower(h)]
			if exists && len(hosters.Rd) > 0 {
				result[h] = true
			}
		}
	}
	return result
}

func (r *RealDebrid) SubmitMagnet(t *torrent.Torrent) (*torrent.Torrent, error) {
	url := fmt.Sprintf("%s/torrents/addMagnet", r.Host)
	payload := gourl.Values{
		"magnet": {t.Magnet.Link},
	}
	var data AddMagnetSchema
	req, _ := http.NewRequest(http.MethodPost, url, strings.NewReader(payload.Encode()))
	resp, err := r.client.MakeRequest(req)
	if err != nil {
		return nil, err
	}
	if err = json.Unmarshal(resp, &data); err != nil {
		return nil, err
	}
	t.Id = data.Id
	t.Debrid = r.Name
	t.MountPath = r.MountPath
	return t, nil
}

func (r *RealDebrid) UpdateTorrent(t *torrent.Torrent) error {
	url := fmt.Sprintf("%s/torrents/info/%s", r.Host, t.Id)
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	resp, err := r.client.MakeRequest(req)
	if err != nil {
		return err
	}
	var data TorrentInfo
	err = json.Unmarshal(resp, &data)
	if err != nil {
		return err
	}
	name := utils.RemoveInvalidChars(data.OriginalFilename)
	t.Name = name
	t.Bytes = data.Bytes
	t.Folder = name
	t.Progress = data.Progress
	t.Status = data.Status
	t.Speed = data.Speed
	t.Seeders = data.Seeders
	t.Filename = data.Filename
	t.OriginalFilename = data.OriginalFilename
	t.Links = data.Links
	t.MountPath = r.MountPath
	t.Debrid = r.Name
	t.Files = getTorrentFiles(t, data, false) // Get selected files
	return nil
}

func (r *RealDebrid) CheckStatus(t *torrent.Torrent, isSymlink bool) (*torrent.Torrent, error) {
	url := fmt.Sprintf("%s/torrents/info/%s", r.Host, t.Id)
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	for {
		resp, err := r.client.MakeRequest(req)
		if err != nil {
			r.logger.Info().Msgf("ERROR Checking file: %v", err)
			return t, err
		}
		var data TorrentInfo
		if err = json.Unmarshal(resp, &data); err != nil {
			return t, err
		}
		status := data.Status
		name := utils.RemoveInvalidChars(data.OriginalFilename)
		t.Name = name // Important because some magnet changes the name
		t.Folder = name
		t.Filename = data.Filename
		t.OriginalFilename = data.OriginalFilename
		t.Bytes = data.Bytes
		t.Progress = data.Progress
		t.Speed = data.Speed
		t.Seeders = data.Seeders
		t.Links = data.Links
		t.Status = status
		t.Debrid = r.Name
		t.MountPath = r.MountPath
		if status == "waiting_files_selection" {
			t.Files = getTorrentFiles(t, data, true)
			if len(t.Files) == 0 {
				return t, fmt.Errorf("no video files found")
			}
			filesId := make([]string, 0)
			for _, f := range t.Files {
				filesId = append(filesId, f.Id)
			}
			p := gourl.Values{
				"files": {strings.Join(filesId, ",")},
			}
			payload := strings.NewReader(p.Encode())
			req, _ := http.NewRequest(http.MethodPost, fmt.Sprintf("%s/torrents/selectFiles/%s", r.Host, t.Id), payload)
			_, err = r.client.MakeRequest(req)
			if err != nil {
				return t, err
			}
		} else if status == "downloaded" {
			t.Files = getTorrentFiles(t, data, false) // Get selected files
			r.logger.Info().Msgf("Torrent: %s downloaded to RD", t.Name)
			if !isSymlink {
				err = r.GenerateDownloadLinks(t)
				if err != nil {
					return t, err
				}
			}
			break
		} else if slices.Contains(r.GetDownloadingStatus(), status) {
			if !t.DownloadUncached {
				return t, fmt.Errorf("torrent: %s not cached", t.Name)
			}
		} else {
			return t, fmt.Errorf("torrent: %s has error: %s", t.Name, status)
		}

	}
	return t, nil
}

func (r *RealDebrid) DeleteTorrent(torrent *torrent.Torrent) {
	url := fmt.Sprintf("%s/torrents/delete/%s", r.Host, torrent.Id)
	req, _ := http.NewRequest(http.MethodDelete, url, nil)
	_, err := r.client.MakeRequest(req)
	if err == nil {
		r.logger.Info().Msgf("Torrent: %s deleted", torrent.Name)
	} else {
		r.logger.Info().Msgf("Error deleting torrent: %s", err)
	}
}

func (r *RealDebrid) GenerateDownloadLinks(t *torrent.Torrent) error {
	url := fmt.Sprintf("%s/unrestrict/link/", r.Host)
	for _, f := range t.Files {
		if f.DownloadLink != "" {
			// Or check the generated link
			continue
		}
		payload := gourl.Values{
			"link": {f.Link},
		}
		req, _ := http.NewRequest(http.MethodPost, url, strings.NewReader(payload.Encode()))
		resp, err := r.client.MakeRequest(req)
		if err != nil {
			return err
		}
		var data UnrestrictResponse
		if err = json.Unmarshal(resp, &data); err != nil {
			return err
		}
		f.DownloadLink = data.Download
		f.Generated = time.Now()
		t.Files[f.Name] = f
	}
	return nil
}

func (r *RealDebrid) ConvertLinksToFiles(links []string) []torrent.File {
	files := make([]torrent.File, 0)
	for _, l := range links {
		url := fmt.Sprintf("%s/unrestrict/link/", r.Host)
		payload := gourl.Values{
			"link": {l},
		}
		req, _ := http.NewRequest(http.MethodPost, url, strings.NewReader(payload.Encode()))
		resp, err := r.client.MakeRequest(req)
		if err != nil {
			continue
		}
		var data UnrestrictResponse
		if err = json.Unmarshal(resp, &data); err != nil {
			continue
		}
		files = append(files, torrent.File{
			Name:         data.Filename,
			Size:         data.Filesize,
			Link:         l,
			DownloadLink: data.Download,
			Generated:    time.Now(),
		})
	}
	return files
}

func (r *RealDebrid) GetDownloadLink(t *torrent.Torrent, file *torrent.File) *torrent.File {
	url := fmt.Sprintf("%s/unrestrict/link/", r.Host)
	payload := gourl.Values{
		"link": {file.Link},
	}
	req, _ := http.NewRequest(http.MethodPost, url, strings.NewReader(payload.Encode()))
	resp, err := r.client.MakeRequest(req)
	if err != nil {
		return nil
	}
	var data UnrestrictResponse
	if err = json.Unmarshal(resp, &data); err != nil {
		return nil
	}
	file.DownloadLink = data.Download
	file.Generated = time.Now()
	return file
}

func (r *RealDebrid) GetCheckCached() bool {
	return r.CheckCached
}

func (r *RealDebrid) getTorrents(offset int, limit int) (int, []*torrent.Torrent, error) {
	url := fmt.Sprintf("%s/torrents?limit=%d", r.Host, limit)
	torrents := make([]*torrent.Torrent, 0)
	if offset > 0 {
		url = fmt.Sprintf("%s&offset=%d", url, offset)
	}
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	resp, err := r.client.Do(req)

	if err != nil {
		return 0, torrents, err
	}

	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return 0, torrents, fmt.Errorf("realdebrid API error: %d", resp.StatusCode)
	}

	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return 0, torrents, err
	}
	totalItems, _ := strconv.Atoi(resp.Header.Get("X-Total-Count"))
	var data []TorrentsResponse
	if err = json.Unmarshal(body, &data); err != nil {
		return 0, nil, err
	}
	filenames := map[string]bool{}
	for _, t := range data {
		if _, exists := filenames[t.Filename]; exists {
			continue
		}
		torrents = append(torrents, &torrent.Torrent{
			Id:               t.Id,
			Name:             t.Filename,
			Bytes:            t.Bytes,
			Progress:         t.Progress,
			Status:           t.Status,
			Filename:         t.Filename,
			OriginalFilename: t.Filename,
			Links:            t.Links,
			Files:            make(map[string]torrent.File),
			InfoHash:         t.Hash,
			Debrid:           r.Name,
			MountPath:        r.MountPath,
		})
	}
	return totalItems, torrents, nil
}

func (r *RealDebrid) GetTorrents() ([]*torrent.Torrent, error) {
	limit := 5000

	// Get first batch and total count
	totalItems, firstBatch, err := r.getTorrents(0, limit)
	if err != nil {
		return nil, err
	}

	allTorrents := firstBatch

	// Calculate remaining requests
	remaining := totalItems - len(firstBatch)
	if remaining <= 0 {
		return allTorrents, nil
	}

	// Prepare for concurrent fetching
	var wg sync.WaitGroup
	var mu sync.Mutex
	var fetchError error

	// Calculate how many more requests we need
	batchCount := (remaining + limit - 1) / limit // ceiling division

	for i := 1; i <= batchCount; i++ {
		wg.Add(1)
		go func(batchOffset int) {
			defer wg.Done()

			_, batch, err := r.getTorrents(batchOffset, limit)
			if err != nil {
				mu.Lock()
				fetchError = err
				mu.Unlock()
				return
			}

			mu.Lock()
			allTorrents = append(allTorrents, batch...)
			mu.Unlock()
		}(i * limit)
	}

	// Wait for all fetches to complete
	wg.Wait()

	if fetchError != nil {
		return nil, fetchError
	}

	return allTorrents, nil
}

func (r *RealDebrid) GetDownloads() (map[string]torrent.DownloadLinks, error) {
	links := make(map[string]torrent.DownloadLinks)
	offset := 0
	limit := 5000
	for {
		dl, err := r._getDownloads(offset, limit)
		if err != nil {
			break
		}
		if len(dl) == 0 {
			break
		}

		for _, d := range dl {
			if _, exists := links[d.Link]; exists {
				// This is ordered by date, so we can skip the rest
				continue
			}
			links[d.Link] = d
		}

		offset += len(dl)
	}
	return links, nil
}

func (r *RealDebrid) _getDownloads(offset int, limit int) ([]torrent.DownloadLinks, error) {
	url := fmt.Sprintf("%s/downloads?limit=%d", r.Host, limit)
	if offset > 0 {
		url = fmt.Sprintf("%s&offset=%d", url, offset)
	}
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	resp, err := r.client.MakeRequest(req)
	if err != nil {
		return nil, err
	}
	var data []DownloadsResponse
	if err = json.Unmarshal(resp, &data); err != nil {
		return nil, err
	}
	links := make([]torrent.DownloadLinks, 0)
	for _, d := range data {
		links = append(links, torrent.DownloadLinks{
			Filename:     d.Filename,
			Size:         d.Filesize,
			Link:         d.Link,
			DownloadLink: d.Download,
			Generated:    d.Generated,
			Id:           d.Id,
		})

	}
	return links, nil
}

func (r *RealDebrid) GetDownloadingStatus() []string {
	return []string{"downloading", "magnet_conversion", "queued", "compressing", "uploading"}
}

func (r *RealDebrid) GetDownloadUncached() bool {
	return r.DownloadUncached
}

func New(dc config.Debrid) *RealDebrid {
	rl := request.ParseRateLimit(dc.RateLimit)
	headers := map[string]string{
		"Authorization": fmt.Sprintf("Bearer %s", dc.APIKey),
	}
	_log := logger.NewLogger(dc.Name)
	client := request.New().
		WithHeaders(headers).
		WithRateLimiter(rl).WithLogger(_log)
	return &RealDebrid{
		Name:             "realdebrid",
		Host:             dc.Host,
		APIKey:           dc.APIKey,
		DownloadUncached: dc.DownloadUncached,
		client:           client,
		MountPath:        dc.Folder,
		logger:           logger.NewLogger(dc.Name),
		CheckCached:      dc.CheckCached,
	}
}