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.Client, language string) *Engine { return &Engine{ db: db, cron: cron.New(), tasks: make(map[int]*Task, 0), language: language, } } type scheduler struct { cron string f func() error } type Engine struct { db *db.Client cron *cron.Cron tasks map[int]*Task language string schedulers utils.Map[string, scheduler] } 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) reloadUsingBuildinDownloader(h *ent.History) error{ cl, err := buildin.NewDownloader(c.db.GetDownloadDir()) 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") } c.tasks[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[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[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[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[t.ID] = &Task{Torrent: to} } } } log.Infof("------ task reloading done ------") } func (c *Engine) buildInDownloader() (pkg.Downloader, error) { dir := c.db.GetDownloadDir() return buildin.NewDownloader(dir) } 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", d.URL) 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 := 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 *Engine) GetTasks() map[int]*Task { return c.tasks }