refactor: size limiter

db/const.go

@@ -1,5 +1,7 @@
 package db
 
+import "polaris/ent/media"
+
 var Version = "undefined"
 
 const (
@@ -17,9 +19,9 @@ const (
 	SettingTvNamingFormat         = "tv_naming_format"
 	SettingMovieNamingFormat      = "movie_naming_format"
 	SettingProwlarrInfo           = "prowlarr_info"
-	Setting720pSizeLimiter        = "720p_size_limiter"
-	Setting1080ppSizeLimiter      = "1080p_size_limiter"
-	Setting2160ppSizeLimiter      = "2160p_size_limiter"
+
+	SettingTvSizeLimiter          = "tv_size_limiter"
+	SettingMovieSizeLimiter       = "movie_size_limiter"
 )
 
 const (
@@ -53,6 +55,23 @@ type ResolutionType string
 
 const JwtSerectKey = "jwt_secrect_key"
 
+type MediaSizeLimiter struct {
+	P720p SizeLimiter `json:"720p"`
+	P1080 SizeLimiter `json:"1080p"`
+	P2160 SizeLimiter `json:"2160p"`
+}
+
+func (m *MediaSizeLimiter) GetLimiter(r media.Resolution) SizeLimiter {
+	if r == media.Resolution1080p {
+		return m.P1080
+	} else if r == media.Resolution720p {
+		return m.P720p
+	} else if r == media.Resolution2160p {
+		return m.P2160
+	}
+	return SizeLimiter{}
+}
+
 type SizeLimiter struct {
 	MaxSIze    int64 `json:"max_size"`
 	MinSize    int64 `json:"min_size"`

db/db.go

@@ -441,7 +441,7 @@ type Storage struct {
 }
 
 func (s *Storage) ToWebDavSetting() WebdavSetting {
-	if s.Implementation != storage.ImplementationWebdav && s.Implementation != storage.ImplementationAlist{
+	if s.Implementation != storage.ImplementationWebdav && s.Implementation != storage.ImplementationAlist {
 		panic("not webdav storage")
 	}
 	var webdavSetting WebdavSetting
@@ -631,19 +631,17 @@ func (c *Client) DeleteImportlist(id int) error {
 	return c.ent.ImportList.DeleteOneID(id).Exec(context.TODO())
 }
 
-func (c *Client) GetSizeLimiter(r media.Resolution) (*SizeLimiter, error) {
+func (c *Client) GetSizeLimiter(mediaType string) (*MediaSizeLimiter, error) {
 	var v string
-	if r == media.Resolution720p {
-		v = c.GetSetting(Setting720pSizeLimiter)
-	} else if r == media.Resolution1080p {
-		v = c.GetSetting(Setting1080ppSizeLimiter)
-	} else if r == media.Resolution2160p {
-		v = c.GetSetting(Setting2160ppSizeLimiter)
+	if mediaType == "tv" {
+		v = c.GetSetting(SettingTvSizeLimiter)
+	} else if mediaType == "movie" {
+		v = c.GetSetting(SettingMovieSizeLimiter)
 	} else {
-		return nil, errors.Errorf("resolution not supported: %v", r)
+		return nil, errors.Errorf("media type not supported: %v", mediaType)
 	}
 
-	var limiter SizeLimiter
+	var limiter MediaSizeLimiter
 	if v == "" {
 		return &limiter, nil
 	}
@@ -652,20 +650,19 @@ func (c *Client) GetSizeLimiter(r media.Resolution) (*SizeLimiter, error) {
 	return &limiter, err
 }
 
-func (c *Client) SetSizeLimiter(r media.Resolution, limiter *SizeLimiter) error {
+func (c *Client) SetSizeLimiter(mediaType string, limiter *MediaSizeLimiter) error {
 	data, err := json.Marshal(limiter)
 	if err != nil {
 		return err
 	}
-	if r == media.Resolution720p {
-		return c.SetSetting(Setting720pSizeLimiter, string(data))
-	} else if r == media.Resolution1080p {
-		return c.SetSetting(Setting1080ppSizeLimiter, string(data))
-	} else if r == media.Resolution2160p {
-		return c.SetSetting(Setting2160ppSizeLimiter, string(data))
+	if mediaType == "tv" {
+		return c.SetSetting(SettingTvSizeLimiter, string(data))
+	} else if mediaType == "movie" {
+		return c.SetSetting(SettingMovieSizeLimiter, string(data))
 	} else {
-		return errors.Errorf("resolution not supported: %v", r)
+		return errors.Errorf("media type not supported: %v", mediaType)
 	}
+
 }
 
 func (c *Client) GetTvNamingFormat() string {

ent/schema/media.go

@@ -42,8 +42,9 @@ func (Media) Edges() []ent.Edge {
 }
 
 type MediaLimiter struct {
-	SizeMin int `json:"size_min"` //in B
-	SizeMax int `json:"size_max"` //in B
+	SizeMin    int64 `json:"size_min"` //in B
+	SizeMax    int64 `json:"size_max"` //in B
+	PreferSize int64 `json:"prefer_max"`
 }
 
 type MediaExtras struct {
@@ -51,7 +52,7 @@ type MediaExtras struct {
 	JavId        string `json:"javid"`
 	//OriginCountry    []string `json:"origin_country"`
 	OriginalLanguage string `json:"original_language"`
-	Genres []struct {
+	Genres           []struct {
 		ID   int64  `json:"id"`
 		Name string `json:"name"`
 	} `json:"genres"`

pkg/torznab/torznab.go

@@ -96,9 +96,9 @@ func (r *Response) ToResults(indexer *db.TorznabInfo) []Result {
 			Description:          item.Description,
 			Link:                 item.Link,
 			Size:                 mustAtoI(item.Size),
-			Seeders:              mustAtoI(item.GetAttr("seeders")),
-			Peers:                mustAtoI(item.GetAttr("peers")),
-			Category:             mustAtoI(item.GetAttr("category")),
+			Seeders:              int(mustAtoI(item.GetAttr("seeders"))),
+			Peers:                int(mustAtoI(item.GetAttr("peers"))),
+			Category:             int(mustAtoI(item.GetAttr("category"))),
 			ImdbId:               imdb,
 			DownloadVolumeFactor: tryParseFloat(item.GetAttr("downloadvolumefactor")),
 			UploadVolumeFactor:   tryParseFloat(item.GetAttr("uploadvolumefactor")),
@@ -112,8 +112,8 @@ func (r *Response) ToResults(indexer *db.TorznabInfo) []Result {
 	return res
 }
 
-func mustAtoI(key string) int {
-	i, err := strconv.Atoi(key)
+func mustAtoI(key string) int64 {
+	i, err := strconv.ParseInt(key, 10, 64)
 	if err != nil {
 		log.Errorf("must atoi error: %v", err)
 		panic(err)
@@ -183,7 +183,7 @@ type Result struct {
 	Name                 string  `json:"name"`
 	Description          string  `json:"description"`
 	Link                 string  `json:"link"`
-	Size                 int     `json:"size"`
+	Size                 int64     `json:"size"`
 	Seeders              int     `json:"seeders"`
 	Peers                int     `json:"peers"`
 	Category             int     `json:"category"`

server/core/importlist.go

@@ -114,8 +114,9 @@ type AddWatchlistIn struct {
 	Resolution              string `json:"resolution" binding:"required"`
 	Folder                  string `json:"folder" binding:"required"`
 	DownloadHistoryEpisodes bool   `json:"download_history_episodes"` //for tv
-	SizeMin                 int    `json:"size_min"`
-	SizeMax                 int    `json:"size_max"`
+	SizeMin                 int64  `json:"size_min"`
+	SizeMax                 int64  `json:"size_max"`
+	PreferSize              int64  `json:"prefer_size"`
 }
 
 func (c *Client) AddTv2Watchlist(in AddWatchlistIn) (interface{}, error) {
@@ -139,7 +140,7 @@ func (c *Client) AddTv2Watchlist(in AddWatchlistIn) (interface{}, error) {
 	}
 	log.Infof("find detail for tv id %d: %+v", in.TmdbID, detail)
 
-	 lastSeason := 0
+	lastSeason := 0
 	for _, season := range detail.Seasons {
 		if season.SeasonNumber > lastSeason && season.EpisodeCount > 0 { //如果最新一季已经有剧集信息，则以最新一季为准
 			lastSeason = season.SeasonNumber
@@ -327,8 +328,8 @@ func (c *Client) checkMovieFolder(m *ent.Media) error {
 		return err
 	}
 
-	for _,f := range files {
-		if f.IsDir() || f.Size() < 100 * 1000 * 1000 /* 100M */{ //忽略路径和小于100M的文件
+	for _, f := range files {
+		if f.IsDir() || f.Size() < 100*1000*1000 /* 100M */ { //忽略路径和小于100M的文件
 			continue
 		}
 		meta := metadata.ParseMovie(f.Name())

server/core/resources.go

@@ -68,7 +68,7 @@ func (c *Client) DownloadEpisodeTorrent(r1 torznab.Result, seriesId, seasonNum i
 		SourceTitle: r1.Name,
 		TargetDir:   dir,
 		Status:      history.StatusRunning,
-		Size:        r1.Size,
+		Size:        int(r1.Size),
 		//Saved:            torrent.Save(),
 		Link:             magnet,
 		DownloadClientID: dlc.ID,
@@ -192,7 +192,7 @@ lo:
 	return torrentNames, nil
 }
 
-func (c *Client) DownloadMovie(m *ent.Media, link, name string, size int, indexerID int) (*string, error) {
+func (c *Client) DownloadMovie(m *ent.Media, link, name string, size int64, indexerID int) (*string, error) {
 	trc, dlc, err := c.GetDownloadClient()
 	if err != nil {
 		return nil, errors.Wrap(err, "connect transmission")
@@ -219,7 +219,7 @@ func (c *Client) DownloadMovie(m *ent.Media, link, name string, size int, indexe
 			SourceTitle: name,
 			TargetDir:   m.TargetDir,
 			Status:      history.StatusRunning,
-			Size:        size,
+			Size:        int(size),
 			//Saved:            torrent.Save(),
 			Link:             magnet,
 			DownloadClientID: dlc.ID,

server/core/scheduler.go

@@ -434,7 +434,7 @@ func (c *Client) downloadMovieSingleEpisode(ep *ent.Episode, targetDir string) (
 		SourceTitle: r1.Name,
 		TargetDir:   targetDir,
 		Status:      history.StatusRunning,
-		Size:        r1.Size,
+		Size:        int(r1.Size),
 		//Saved:            torrent.Save(),
 		Link:             magnet,
 		DownloadClientID: dlc.ID,

server/core/torrent.go

@@ -114,8 +114,8 @@ func imdbIDMatchExact(id1, id2 string) bool {
 	return id1 == id2
 }
 
-func torrentSizeOk(detail *db.MediaDetails, torrentSize int, torrentEpisodeNum int, param *SearchParam) bool {
-	defaultMinSize := 80 * 1000 * 1000 //tv, 80M min
+func torrentSizeOk(detail *db.MediaDetails, torrentSize int64, torrentEpisodeNum int, param *SearchParam) bool {
+	defaultMinSize := int64(80 * 1000 * 1000) //tv, 80M min
 	if detail.MediaType == media.MediaTypeMovie {
 		defaultMinSize = 200 * 1000 * 1000 // movie, 200M min
 	}
@@ -134,20 +134,20 @@ func torrentSizeOk(detail *db.MediaDetails, torrentSize int, torrentEpisodeNum i
 
 	if param.CheckFileSize { //check file size when trigger automatic download
 		if detail.Limiter.SizeMin > 0 { //min size
-			sizeMin := detail.Limiter.SizeMin * multiplier
+			sizeMin := detail.Limiter.SizeMin * int64(multiplier)
 			if torrentSize < sizeMin { //比最小要求的大小还要小, min size not qualify
 				return false
 			}
 		}
 
 		if detail.Limiter.SizeMax > 0 { //max size
-			sizeMax := detail.Limiter.SizeMax * multiplier
+			sizeMax := detail.Limiter.SizeMax * int64(multiplier)
 			if torrentSize > sizeMax { //larger than max size wanted, max size not qualify
 				return false
 			}
 		}
 	}
-	return torrentSize > defaultMinSize*multiplier
+	return torrentSize > defaultMinSize*int64(multiplier)
 }
 
 func seasonEpisodeCount(detail *db.MediaDetails, seasonNum int) int {

server/setting.go

@@ -7,7 +7,6 @@ import (
 	"polaris/db"
 	"polaris/ent"
 	"polaris/ent/downloadclients"
-	"polaris/ent/media"
 	"polaris/log"
 	"polaris/pkg/prowlarr"
 	"polaris/pkg/qbittorrent"
@@ -331,29 +330,22 @@ func (s *Server) SaveProwlarrSetting(c *gin.Context) (interface{}, error) {
 }
 
 type ResolutionSizeLimiter struct {
-	P720  db.SizeLimiter `json:"720p"`
-	P1080 db.SizeLimiter `json:"1080p"`
-	P2160 db.SizeLimiter `json:"2160p"`
+	TvLimiter    *db.MediaSizeLimiter `json:"tv_limiter"`
+	MovieLimiter *db.MediaSizeLimiter `json:"movie_limiter"`
 }
 
 func (s *Server) GetSizeLimiter(c *gin.Context) (interface{}, error) {
-	p720, err := s.db.GetSizeLimiter(media.Resolution720p)
+	tv, err := s.db.GetSizeLimiter("tv")
 	if err != nil {
 		return nil, errors.Wrap(err, "db")
 	}
-	p1080, err := s.db.GetSizeLimiter(media.Resolution1080p)
-	if err != nil {
-		return nil, errors.Wrap(err, "db")
-	}
-
-	p2160, err := s.db.GetSizeLimiter(media.Resolution2160p)
+	movie, err := s.db.GetSizeLimiter("movie")
 	if err != nil {
 		return nil, errors.Wrap(err, "db")
 	}
 	r := ResolutionSizeLimiter{
-		P720: *p720,
-		P1080: *p1080,
-		P2160: *p2160,
+		TvLimiter:    tv,
+		MovieLimiter: movie,
 	}
 	return r, nil
 }
@@ -363,15 +355,11 @@ func (s *Server) SetSizeLimiter(c *gin.Context) (interface{}, error) {
 	if err := c.ShouldBindJSON(&in); err != nil {
 		return nil, err
 	}
-	if err := s.db.SetSizeLimiter(media.Resolution720p, &in.P720); err != nil {
+	if err := s.db.SetSizeLimiter("tv", in.TvLimiter); err != nil {
 		return nil, errors.Wrap(err, "db")
 	}
-	if err := s.db.SetSizeLimiter(media.Resolution1080p, &in.P1080); err != nil {
-		return nil, errors.Wrap(err, "db")
-	}
-
-	if err := s.db.SetSizeLimiter(media.Resolution2160p, &in.P2160); err != nil {
+	if err := s.db.SetSizeLimiter("movie", in.MovieLimiter); err != nil {
 		return nil, errors.Wrap(err, "db")
 	}
 	return "success", nil
-}
+}