tstor/src/sources/torrent/controller.go

package torrent

import (
	"context"
	"log/slog"
	"slices"
	"strings"

	"git.kmsign.ru/royalcat/tstor/pkg/rlog"
	"github.com/anacrolix/torrent"
)

type Controller struct {
	torrentFilePath string
	t               *torrent.Torrent
	rep             *filesMappingsStore
	log             *rlog.Logger
}

func newController(t *torrent.Torrent, rep *filesMappingsStore) *Controller {
	return &Controller{
		t:   t,
		rep: rep,
		log: rlog.Component("torrent/controller").With(slog.String("infohash", t.InfoHash().HexString())),
	}
}

func (s *Controller) TorrentFilePath() string {
	return s.torrentFilePath
}

func (s *Controller) Torrent() *torrent.Torrent {
	return s.t
}

func (c *Controller) Name() string {
	<-c.t.GotInfo()
	if name := c.t.Name(); name != "" {
		return name
	}

	return c.InfoHash()
}

func (s *Controller) InfoHash() string {
	<-s.t.GotInfo()
	return s.t.InfoHash().HexString()
}

func (s *Controller) BytesCompleted() int64 {
	<-s.t.GotInfo()
	return s.t.BytesCompleted()
}

func (s *Controller) BytesMissing() int64 {
	<-s.t.GotInfo()
	return s.t.BytesMissing()
}

func (s *Controller) Length() int64 {
	<-s.t.GotInfo()
	return s.t.Length()
}

func (s *Controller) Files(ctx context.Context) ([]*torrent.File, error) {
	fileMappings, err := s.rep.FileMappings(ctx, s.t.InfoHash())
	if err != nil {
		return nil, err
	}

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-s.t.GotInfo():
	}

	files := s.t.Files()
	files = slices.DeleteFunc(files, func(file *torrent.File) bool {
		if file == nil {
			return true
		}

		p := file.Path()
		if strings.Contains(p, "/.pad/") {
			return true
		}
		if target, ok := fileMappings[p]; ok && target == "" {
			return true
		}
		return false
	})

	return files, nil
}

func Map[T, U any](ts []T, f func(T) U) []U {
	us := make([]U, len(ts))
	for i := range ts {
		us[i] = f(ts[i])
	}
	return us
}

func (s *Controller) ExcludeFile(ctx context.Context, f *torrent.File) error {
	return s.rep.ExcludeFile(ctx, f)
}

func (s *Controller) isFileComplete(startIndex int, endIndex int) bool {
	for i := startIndex; i < endIndex; i++ {
		if !s.t.Piece(i).State().Complete {
			return false
		}
	}
	return true
}

func (s *Controller) ValidateTorrent(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-s.t.GotInfo():
	}

	for i := 0; i < s.t.NumPieces(); i++ {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		s.t.Piece(i).VerifyData()
	}

	return nil
}

func (c *Controller) initializeTorrentPriories(ctx context.Context) error {
	log := c.log.WithComponent("initializeTorrentPriories")

	// files, err := c.Files(ctx)
	// if err != nil {
	// 	return err
	// }

	// for _, file := range files {
	// 	if file == nil {
	// 		continue
	// 	}

	// 	file.SetPriority(torrent.PiecePriorityNormal)
	// }

	log.Info(ctx, "torrent initialization complete", slog.String("infohash", c.InfoHash()), slog.String("torrent_name", c.Name()))

	return nil
}