polaris/server/core/client.go

package core

import (
	"encoding/json"
	"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() {
	c.reloadTasks()
	c.addSysCron()
}

func (c *Client) reloadTasks() {
	allTasks := c.db.GetRunningHistories()
	for _, t := range allTasks {
		var torrent pkg.Torrent
		if tt, err := c.reloadTransmiision(t.Saved); err == nil {
			torrent = tt
			log.Infof("load transmission task: %v", t.Saved)
		} else if tt, err := c.reloadQbit(t.Saved); err == nil {
			torrent = tt
			log.Infof("load qbit task: %v", t.Saved)
		} else {
			log.Warnf("load task fail: %v", t.Saved)
			continue
		}
		c.tasks[t.ID] = &Task{Torrent: torrent}
	}
}

func (c *Client) reloadTransmiision(s string) (pkg.Torrent, error) {
	var t transmission.Torrent
	if err := json.Unmarshal([]byte(s), &t); err != nil {
		return nil, err
	}
	if err := t.Reload(); err != nil {
		return nil, err
	}
	return &t, nil
}

func (c *Client) reloadQbit(s string) (pkg.Torrent, error) {
	var t qbittorrent.Torrent
	if err := json.Unmarshal([]byte(s), &t); err != nil {
		return nil, err
	}
	if err := t.Reload(); err != nil {
		return nil, err
	}
	return &t, nil
}


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
}