decypharr/pkg/webdav/cache.go

package webdav

import (
	"bufio"
	"context"
	"fmt"
	"github.com/dgraph-io/badger/v4"
	"github.com/goccy/go-json"
	"github.com/rs/zerolog"
	"github.com/sirrobot01/debrid-blackhole/internal/logger"
	"github.com/sirrobot01/debrid-blackhole/pkg/debrid/debrid"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"github.com/sirrobot01/debrid-blackhole/internal/config"
	"github.com/sirrobot01/debrid-blackhole/pkg/debrid/torrent"
)

type DownloadLinkCache struct {
	Link string `json:"download_link"`
}

type propfindResponse struct {
	data []byte
	ts   time.Time
}

type CachedTorrent struct {
	*torrent.Torrent
	LastRead   time.Time `json:"last_read"`
	IsComplete bool      `json:"is_complete"`
}

type Cache struct {
	dir    string
	client debrid.Client
	db     *badger.DB
	logger zerolog.Logger

	torrents      map[string]*CachedTorrent // key: torrent.Id, value: *CachedTorrent
	torrentsNames map[string]*CachedTorrent // key: torrent.Name, value: torrent
	listings      atomic.Value
	downloadLinks map[string]string // key: file.Link, value: download link
	propfindResp  sync.Map

	workers int

	LastUpdated time.Time `json:"last_updated"`

	// refresh mutex
	listingRefreshMu       sync.Mutex // for refreshing torrents
	downloadLinksRefreshMu sync.Mutex // for refreshing download links
	torrentsRefreshMu      sync.Mutex // for refreshing torrents

	// Data Mutexes
	torrentsMutex      sync.RWMutex // for torrents and torrentsNames
	downloadLinksMutex sync.Mutex   // for downloadLinks
}

func (c *Cache) setTorrent(t *CachedTorrent) {
	c.torrentsMutex.Lock()
	c.torrents[t.Id] = t
	c.torrentsNames[t.Name] = t
	c.torrentsMutex.Unlock()

	go c.refreshListings() // This is concurrent safe

	go func() {
		if err := c.SaveTorrent(t); err != nil {
			c.logger.Debug().Err(err).Msgf("Failed to save torrent %s", t.Id)
		}
	}()
}

func (c *Cache) refreshListings() {
	// Copy the current torrents to avoid concurrent issues
	c.torrentsMutex.RLock()
	torrents := make([]string, 0, len(c.torrents))
	for _, t := range c.torrents {
		if t != nil && t.Torrent != nil {
			torrents = append(torrents, t.Name)
		}
	}
	c.torrentsMutex.RUnlock()

	sort.Slice(torrents, func(i, j int) bool {
		return torrents[i] < torrents[j]
	})

	files := make([]os.FileInfo, 0, len(torrents))
	now := time.Now()
	for _, t := range torrents {
		files = append(files, &FileInfo{
			name:    t,
			size:    0,
			mode:    0755 | os.ModeDir,
			modTime: now,
			isDir:   true,
		})
	}
	// Atomic store of the complete ready-to-use slice
	c.listings.Store(files)
}

func (c *Cache) GetListing() []os.FileInfo {
	return c.listings.Load().([]os.FileInfo)
}

func (c *Cache) setTorrents(torrents map[string]*CachedTorrent) {
	c.torrentsMutex.Lock()
	for _, t := range torrents {
		c.torrents[t.Id] = t
		c.torrentsNames[t.Name] = t
	}

	c.torrentsMutex.Unlock()

	go c.refreshListings() // This is concurrent safe

	go func() {
		if err := c.SaveTorrents(); err != nil {
			c.logger.Debug().Err(err).Msgf("Failed to save torrents")
		}
	}()
}

func (c *Cache) GetTorrents() map[string]*CachedTorrent {
	c.torrentsMutex.RLock()
	defer c.torrentsMutex.RUnlock()
	result := make(map[string]*CachedTorrent, len(c.torrents))
	for k, v := range c.torrents {
		result[k] = v
	}
	return result
}

func (c *Cache) GetTorrentNames() map[string]*CachedTorrent {
	c.torrentsMutex.RLock()
	defer c.torrentsMutex.RUnlock()
	return c.torrentsNames
}

type Manager struct {
	caches map[string]*Cache
}

func NewCacheManager(clients []debrid.Client) *Manager {
	m := &Manager{
		caches: make(map[string]*Cache),
	}

	for _, client := range clients {
		m.caches[client.GetName()] = NewCache(client)
	}

	return m
}

func (m *Manager) GetCaches() map[string]*Cache {
	return m.caches
}

func (m *Manager) GetCache(debridName string) *Cache {
	return m.caches[debridName]
}

func NewCache(client debrid.Client) *Cache {
	cfg := config.GetConfig()
	dbPath := filepath.Join(cfg.Path, "cache", client.GetName())
	return &Cache{
		dir:           dbPath,
		torrents:      make(map[string]*CachedTorrent),
		torrentsNames: make(map[string]*CachedTorrent),
		client:        client,
		logger:        logger.NewLogger(fmt.Sprintf("%s-cache", client.GetName())),
		workers:       200,
		downloadLinks: make(map[string]string),
	}
}

func (c *Cache) Start() error {
	if err := os.MkdirAll(c.dir, 0755); err != nil {
		return fmt.Errorf("failed to create cache directory: %w", err)
	}

	if err := c.Sync(); err != nil {
		return fmt.Errorf("failed to sync cache: %w", err)
	}

	// initial download links
	go func() {
		// lock download refresh mutex
		c.downloadLinksRefreshMu.Lock()
		defer c.downloadLinksRefreshMu.Unlock()
		// This prevents the download links from being refreshed twice
		c.refreshDownloadLinks()
	}()

	go func() {
		err := c.Refresh()
		if err != nil {
			c.logger.Error().Err(err).Msg("Failed to start cache refresh worker")
		}
	}()

	return nil
}

func (c *Cache) Close() error {
	if c.db != nil {
		return c.db.Close()
	}
	return nil
}

func (c *Cache) load() (map[string]*CachedTorrent, error) {
	torrents := make(map[string]*CachedTorrent)
	if err := os.MkdirAll(c.dir, 0755); err != nil {
		return torrents, fmt.Errorf("failed to create cache directory: %w", err)
	}

	files, err := os.ReadDir(c.dir)
	if err != nil {
		return torrents, fmt.Errorf("failed to read cache directory: %w", err)
	}

	for _, file := range files {
		if file.IsDir() || filepath.Ext(file.Name()) != ".json" {
			continue
		}

		filePath := filepath.Join(c.dir, file.Name())
		data, err := os.ReadFile(filePath)
		if err != nil {
			c.logger.Debug().Err(err).Msgf("Failed to read file: %s", filePath)
			continue
		}

		var ct CachedTorrent
		if err := json.Unmarshal(data, &ct); err != nil {
			c.logger.Debug().Err(err).Msgf("Failed to unmarshal file: %s", filePath)
			continue
		}
		if len(ct.Files) != 0 {
			// We can assume the torrent is complete
			ct.IsComplete = true
			torrents[ct.Id] = &ct
		}
	}

	return torrents, nil
}

func (c *Cache) GetTorrent(id string) *CachedTorrent {
	c.torrentsMutex.RLock()
	defer c.torrentsMutex.RUnlock()
	if t, ok := c.torrents[id]; ok {
		return t
	}
	return nil
}

func (c *Cache) GetTorrentByName(name string) *CachedTorrent {
	if t, ok := c.GetTorrentNames()[name]; ok {
		return t
	}
	return nil
}

func (c *Cache) SaveTorrents() error {
	for _, ct := range c.GetTorrents() {
		if err := c.SaveTorrent(ct); err != nil {
			return err
		}
	}
	return nil
}

func (c *Cache) SaveTorrent(ct *CachedTorrent) error {
	data, err := json.MarshalIndent(ct, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal torrent: %w", err)
	}

	fileName := ct.Torrent.Id + ".json"
	filePath := filepath.Join(c.dir, fileName)
	tmpFile := filePath + ".tmp"

	f, err := os.Create(tmpFile)
	if err != nil {
		return fmt.Errorf("failed to create temp file: %w", err)
	}
	defer f.Close()

	w := bufio.NewWriter(f)
	if _, err := w.Write(data); err != nil {
		return fmt.Errorf("failed to write data: %w", err)
	}

	if err := w.Flush(); err != nil {
		return fmt.Errorf("failed to flush data: %w", err)
	}

	return os.Rename(tmpFile, filePath)
}

func (c *Cache) Sync() error {
	cachedTorrents, err := c.load()
	if err != nil {
		c.logger.Debug().Err(err).Msg("Failed to load cache")
	}

	torrents, err := c.client.GetTorrents()
	if err != nil {
		return fmt.Errorf("failed to sync torrents: %v", err)
	}

	c.logger.Info().Msgf("Got %d torrents from %s", len(torrents), c.client.GetName())

	newTorrents := make([]*torrent.Torrent, 0)
	idStore := make(map[string]bool, len(torrents))
	for _, t := range torrents {
		idStore[t.Id] = true
		if _, ok := cachedTorrents[t.Id]; !ok {
			newTorrents = append(newTorrents, t)
		}
	}

	// Check for deleted torrents
	deletedTorrents := make([]string, 0)
	for _, t := range cachedTorrents {
		if _, ok := idStore[t.Id]; !ok {
			deletedTorrents = append(deletedTorrents, t.Id)
		}
	}

	if len(deletedTorrents) > 0 {
		c.logger.Info().Msgf("Found %d deleted torrents", len(deletedTorrents))
		for _, id := range deletedTorrents {
			if _, ok := cachedTorrents[id]; ok {
				delete(cachedTorrents, id)
				c.removeFromDB(id)
			}
		}
	}

	// Write these torrents to the cache
	c.setTorrents(cachedTorrents)
	c.logger.Info().Msgf("Loaded %d torrents from cache", len(cachedTorrents))

	if len(newTorrents) > 0 {
		c.logger.Info().Msgf("Found %d new torrents", len(newTorrents))
		if err := c.sync(newTorrents); err != nil {
			return fmt.Errorf("failed to sync torrents: %v", err)
		}
	}

	return nil
}

func (c *Cache) sync(torrents []*torrent.Torrent) error {
	// Calculate optimal workers - balance between CPU and IO
	workers := runtime.NumCPU() * 50 // A more balanced multiplier for BadgerDB

	// Create channels with appropriate buffering
	workChan := make(chan *torrent.Torrent, workers*2)

	// Use an atomic counter for progress tracking
	var processed int64
	var errorCount int64

	// Create a context with cancellation in case of critical errors
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create a wait group for workers
	var wg sync.WaitGroup

	// Start workers
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case t, ok := <-workChan:
					if !ok {
						return // Channel closed, exit goroutine
					}

					if err := c.processTorrent(t); err != nil {
						c.logger.Error().Err(err).Str("torrent", t.Name).Msg("sync error")
						atomic.AddInt64(&errorCount, 1)
					}

					count := atomic.AddInt64(&processed, 1)
					if count%1000 == 0 {
						c.logger.Info().Msgf("Progress: %d/%d torrents processed", count, len(torrents))
					}

				case <-ctx.Done():
					return // Context cancelled, exit goroutine
				}
			}
		}()
	}

	// Feed work to workers
	for _, t := range torrents {
		select {
		case workChan <- t:
			// Work sent successfully
		case <-ctx.Done():
			break // Context cancelled
		}
	}

	// Signal workers that no more work is coming
	close(workChan)

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

	c.logger.Info().Msgf("Sync complete: %d torrents processed, %d errors", len(torrents), errorCount)
	return nil
}

func (c *Cache) processTorrent(t *torrent.Torrent) error {
	var err error
	err = c.client.UpdateTorrent(t)
	if err != nil {
		return fmt.Errorf("failed to get torrent files: %v", err)
	}

	ct := &CachedTorrent{
		Torrent:    t,
		LastRead:   time.Now(),
		IsComplete: len(t.Files) > 0,
	}
	c.setTorrent(ct)
	return nil
}

func (c *Cache) GetDownloadLink(torrentId, filename, fileLink string) string {

	// Check link cache
	if dl := c.checkDownloadLink(fileLink); dl != "" {
		return dl
	}

	ct := c.GetTorrent(torrentId)
	if ct == nil {
		return ""
	}
	file := ct.Files[filename]

	if file.Link == "" {
		// file link is empty, refresh the torrent to get restricted links
		if ct.IsComplete {
			return ""
		}
		ct = c.refreshTorrent(ct) // Refresh the torrent from the debrid service
		if ct == nil {
			return ""
		} else {
			file = ct.Files[filename]
		}
	}

	c.logger.Debug().Msgf("Getting download link for %s", ct.Name)
	f := c.client.GetDownloadLink(ct.Torrent, &file)
	if f == nil {
		return ""
	}
	file.DownloadLink = f.DownloadLink
	ct.Files[filename] = file

	go c.updateDownloadLink(f)
	go c.setTorrent(ct)
	return f.DownloadLink
}

func (c *Cache) updateDownloadLink(file *torrent.File) {
	c.downloadLinksMutex.Lock()
	defer c.downloadLinksMutex.Unlock()
	c.downloadLinks[file.Link] = file.DownloadLink
}

func (c *Cache) checkDownloadLink(link string) string {
	if dl, ok := c.downloadLinks[link]; ok {
		return dl
	}
	return ""
}

func (c *Cache) refreshTorrent(t *CachedTorrent) *CachedTorrent {
	_torrent := t.Torrent
	err := c.client.UpdateTorrent(_torrent)
	if err != nil {
		c.logger.Debug().Msgf("Failed to get torrent files for %s: %v", t.Id, err)
		return nil
	}
	if len(t.Files) == 0 {
		return nil
	}

	ct := &CachedTorrent{
		Torrent:    _torrent,
		LastRead:   time.Now(),
		IsComplete: len(t.Files) > 0,
	}
	c.setTorrent(ct)

	return ct
}

func (c *Cache) refreshDownloadLinks() map[string]string {
	c.downloadLinksMutex.Lock()
	defer c.downloadLinksMutex.Unlock()

	downloadLinks, err := c.client.GetDownloads()
	if err != nil {
		c.logger.Debug().Err(err).Msg("Failed to get download links")
		return nil
	}
	for k, v := range downloadLinks {
		c.downloadLinks[k] = v.DownloadLink
	}
	return c.downloadLinks
}

func (c *Cache) GetClient() debrid.Client {
	return c.client
}

func (c *Cache) refreshTorrents() {
	c.torrentsMutex.RLock()
	currentTorrents := c.torrents //
	// Create a copy of the current torrents to avoid concurrent issues
	torrents := make(map[string]string, len(currentTorrents)) // a mpa of id and name
	for _, v := range currentTorrents {
		torrents[v.Id] = v.Name
	}
	c.torrentsMutex.RUnlock()

	// Get new torrents from the debrid service
	debTorrents, err := c.client.GetTorrents()
	if err != nil {
		c.logger.Debug().Err(err).Msg("Failed to get torrents")
		return
	}

	if len(debTorrents) == 0 {
		// Maybe an error occurred
		return
	}

	// Get the newly added torrents only
	newTorrents := make([]*torrent.Torrent, 0)
	idStore := make(map[string]bool, len(debTorrents))
	for _, t := range debTorrents {
		idStore[t.Id] = true
		if _, ok := torrents[t.Id]; !ok {
			newTorrents = append(newTorrents, t)
		}
	}

	// Check for deleted torrents
	deletedTorrents := make([]string, 0)
	for id, _ := range torrents {
		if _, ok := idStore[id]; !ok {
			deletedTorrents = append(deletedTorrents, id)
		}
	}

	if len(deletedTorrents) > 0 {
		c.DeleteTorrent(deletedTorrents)
	}

	if len(newTorrents) == 0 {
		return
	}
	c.logger.Info().Msgf("Found %d new torrents", len(newTorrents))

	// No need for a complex sync process, just add the new torrents
	wg := sync.WaitGroup{}
	wg.Add(len(newTorrents))
	for _, t := range newTorrents {
		// processTorrent is concurrent safe
		go func() {
			defer wg.Done()
			if err := c.processTorrent(t); err != nil {
				c.logger.Info().Err(err).Msg("Failed to process torrent")
			}

		}()
	}
	wg.Wait()
}

func (c *Cache) DeleteTorrent(ids []string) {
	c.logger.Info().Msgf("Deleting %d torrents", len(ids))
	c.torrentsMutex.Lock()
	defer c.torrentsMutex.Unlock()
	for _, id := range ids {
		if t, ok := c.torrents[id]; ok {
			delete(c.torrents, id)
			delete(c.torrentsNames, t.Name)
			c.removeFromDB(id)
		}
	}
}

func (c *Cache) removeFromDB(torrentId string) {
	filePath := filepath.Join(c.dir, torrentId+".json")
	if err := os.Remove(filePath); err != nil {
		c.logger.Debug().Err(err).Msgf("Failed to remove file: %s", filePath)
	}
}

func (c *Cache) resetPropfindResponse() {
	// Right now, parents are hardcoded
	parents := []string{"__all__", "torrents"}
	// Reset only the parent directories
	// Convert the parents to a keys
	// This is a bit hacky, but it works
	// Instead of deleting all the keys, we only delete the parent keys, e.g __all__/ or torrents/
	keys := make([]string, 0, len(parents))
	for _, p := range parents {
		// Construct the key
		// construct url
		url := filepath.Join("/webdav/%s/%s", c.client.GetName(), p)
		key0 := fmt.Sprintf("propfind:%s:0", url)
		key1 := fmt.Sprintf("propfind:%s:1", url)
		keys = append(keys, key0, key1)
	}

	// Delete the keys
	for _, k := range keys {
		c.propfindResp.Delete(k)
	}
}