polaris/server/core/client.go

package core

import (
	"polaris/db"
	"polaris/ent"
	"polaris/ent/downloadclients"
	"polaris/log"
	"polaris/pkg"
	"polaris/pkg/qbittorrent"
	"polaris/pkg/tmdb"
	"polaris/pkg/transmission"
	"polaris/pkg/utils"

	"github.com/pkg/errors"
	"github.com/robfig/cron"
)

func NewClient(db *db.Client, language string) *Client {
	return &Client{
		db:       db,
		cron:     cron.New(),
		tasks:    make(map[int]*Task, 0),
		language: language,
	}
}

type scheduler struct {
	cron string
	f    func() error
}
type Client struct {
	db         *db.Client
	cron       *cron.Cron
	tasks      map[int]*Task
	language   string
	schedulers utils.Map[string, scheduler]
}

func (c *Client) registerCronJob(name string, cron string, f func() error) {
	c.schedulers.Store(name, scheduler{
		cron: cron,
		f:    f,
	})
}

func (c *Client) Init() {
	go c.reloadTasks()
	c.addSysCron()
	go c.checkW500PosterOnStartup()
}

func (c *Client) reloadTasks() {
	allTasks := c.db.GetRunningHistories()
	for _, t := range allTasks {
		dl, err := c.db.GetDownloadClient(t.DownloadClientID)
		if err != nil {
			log.Warnf("no download client related: %v", t.SourceTitle)
			continue
		}

		if dl.Implementation == downloadclients.ImplementationTransmission {
			to, err := transmission.NewTorrent(transmission.Config{
				URL:      dl.URL,
				User:     dl.User,
				Password: dl.Password,
			}, t.Link)
			if err != nil {
				log.Warnf("get task error: %v", err)
				continue
			}
			c.tasks[t.ID] = &Task{Torrent: to}
		} else if dl.Implementation == downloadclients.ImplementationQbittorrent {
			to, err := qbittorrent.NewTorrent(qbittorrent.Info{
				URL:      dl.URL,
				User:     dl.User,
				Password: dl.Password,
			}, t.Link)
			if err != nil {
				log.Warnf("get task error: %v", err)
				continue
			}
			c.tasks[t.ID] = &Task{Torrent: to}
		}

	}
	log.Infof("------ task reloading done ------")
}

func (c *Client) GetDownloadClient() (pkg.Downloader, *ent.DownloadClients, error) {
	downloaders := c.db.GetAllDonloadClients()
	for _, d := range downloaders {
		if !d.Enable {
			continue
		}
		if d.Implementation == downloadclients.ImplementationTransmission {
			trc, err := transmission.NewClient(transmission.Config{
				URL:      d.URL,
				User:     d.User,
				Password: d.Password,
			})
			if err != nil {
				log.Warnf("connect to download client error: %v", d.URL)
				continue
			}
			return trc, d, nil

		} else if d.Implementation == downloadclients.ImplementationQbittorrent {
			qbt, err := qbittorrent.NewClient(d.URL, d.User, d.Password)
			if err != nil {
				log.Warnf("connect to download client error: %v", d.URL)
				continue
			}
			return qbt, d, nil
		}
	}
	return nil, nil, errors.Errorf("no available download client")
}

func (c *Client) TMDB() (*tmdb.Client, error) {
	api := c.db.GetSetting(db.SettingTmdbApiKey)
	if api == "" {
		return nil, errors.New("TMDB apiKey not set")
	}
	proxy := c.db.GetSetting(db.SettingProxy)
	adult := c.db.GetSetting(db.SettingEnableTmdbAdultContent)
	return tmdb.NewClient(api, proxy, adult == "true")
}

func (c *Client) MustTMDB() *tmdb.Client {
	t, err := c.TMDB()
	if err != nil {
		log.Panicf("get tmdb: %v", err)
	}
	return t
}

func (c *Client) RemoveTaskAndTorrent(id int) error {
	torrent := c.tasks[id]
	if torrent != nil {
		if err := torrent.Remove(); err != nil {
			return errors.Wrap(err, "remove torrent")
		}
		delete(c.tasks, id)
	}
	return nil
}

func (c *Client) GetTasks() map[int]*Task {
	return c.tasks
}