package engine import ( "polaris/db" "polaris/ent" "polaris/ent/downloadclients" "polaris/log" "polaris/pkg" "polaris/pkg/buildin" "polaris/pkg/qbittorrent" "polaris/pkg/tmdb" "polaris/pkg/transmission" "polaris/pkg/utils" "github.com/pkg/errors" "github.com/robfig/cron" ) func NewEngine(db db.Database, language string) *Engine { return &Engine{ db: db, cron: cron.New(), tasks: utils.Map[int, *Task]{}, schedulers: utils.Map[string, scheduler]{}, language: language, } } type scheduler struct { cron string f func() error } type Engine struct { db db.Database cron *cron.Cron tasks utils.Map[int, *Task] language string schedulers utils.Map[string, scheduler] buildin *buildin.Downloader } func (c *Engine) registerCronJob(name string, cron string, f func() error) { c.schedulers.Store(name, scheduler{ cron: cron, f: f, }) } func (c *Engine) Init() { go c.reloadTasks() c.addSysCron() go c.checkW500PosterOnStartup() } func (c *Engine) GetTask(id int) (*Task, bool) { return c.tasks.Load(id) } func (c *Engine) reloadUsingBuildinDownloader(h *ent.History) error { cl, err := c.buildInDownloader() if err != nil { log.Warnf("buildin downloader error: %v", err) } t, err := cl.Download(h.Link, h.Hash, c.db.GetDownloadDir()) if err != nil { return errors.Wrap(err, "download torrent") } t.Start() c.tasks.Store(h.ID, &Task{Torrent: t}) return nil } func (c *Engine) reloadTasks() { allTasks := c.db.GetRunningHistories() for _, t := range allTasks { if t.DownloadClientID == 0 { log.Warnf("assume buildin downloader: %v", t.SourceTitle) err := c.reloadUsingBuildinDownloader(t) if err != nil { log.Warnf("buildin downloader error: %v", err) } else { log.Infof("success reloading buildin task: %v", t.SourceTitle) } continue } 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 { if t.Hash != "" { //优先使用hash to, err := transmission.NewTorrentHash(transmission.Config{ URL: dl.URL, User: dl.User, Password: dl.Password, }, t.Hash) if err != nil { log.Warnf("get task error: %v", err) continue } c.tasks.Store(t.ID, &Task{Torrent: to}) } else if t.Link != "" { 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.Store(t.ID, &Task{Torrent: to}) } } else if dl.Implementation == downloadclients.ImplementationQbittorrent { if t.Hash != "" { to, err := qbittorrent.NewTorrentHash(qbittorrent.Info{ URL: dl.URL, User: dl.User, Password: dl.Password, }, t.Hash) if err != nil { log.Warnf("get task error: %v", err) continue } c.tasks.Store(t.ID, &Task{Torrent: to}) } else if t.Link != "" { 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.Store(t.ID, &Task{Torrent: to}) } } else if dl.Implementation == downloadclients.ImplementationBuildin { err := c.reloadUsingBuildinDownloader(t) if err != nil { log.Warnf("buildin downloader error: %v", err) } else { log.Infof("success reloading buildin task: %v", t.SourceTitle) } } } log.Infof("------ task reloading done ------") } func (c *Engine) buildInDownloader() (pkg.Downloader, error) { if c.buildin != nil { return c.buildin, nil } dir := c.db.GetDownloadDir() d, err := buildin.NewDownloader(dir) if err != nil { return nil, errors.Wrap(err, "buildin downloader") } c.buildin = d return d, nil } func (c *Engine) 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 } else if d.Implementation == downloadclients.ImplementationBuildin { bin, err := c.buildInDownloader() if err != nil { log.Warnf("connect to download client error: %v", err) continue } return bin, d, nil } } return nil, nil, errors.Errorf("no available download client") } func (c *Engine) TMDB() (*tmdb.Client, error) { api := c.db.GetTmdbApiKey() 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 *Engine) MustTMDB() *tmdb.Client { t, err := c.TMDB() if err != nil { log.Panicf("get tmdb: %v", err) } return t } func (c *Engine) RemoveTaskAndTorrent(id int) error { torrent, ok := c.tasks.Load(id) if ok { if err := torrent.Remove(); err != nil { return errors.Wrap(err, "remove torrent") } c.tasks.Delete(id) } return nil } func (c *Engine) GetTasks() *utils.Map[int, *Task] { return &c.tasks }