polaris/db/db.go

package db

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"polaris/ent"
	"polaris/ent/downloadclients"
	"polaris/ent/episode"
	"polaris/ent/history"
	"polaris/ent/importlist"
	"polaris/ent/indexers"
	"polaris/ent/media"
	"polaris/ent/schema"
	"polaris/ent/settings"
	"polaris/ent/storage"
	"polaris/log"
	"polaris/pkg/utils"
	"strings"
	"time"

	"entgo.io/ent/dialect"
	"entgo.io/ent/dialect/sql"
	_ "github.com/mattn/go-sqlite3"
	"github.com/pkg/errors"
)

type Client struct {
	ent *ent.Client
}

func Open() (*Client, error) {
	os.Mkdir(DataPath, 0777)
	client, err := ent.Open(dialect.SQLite, fmt.Sprintf("file:%v/polaris.db?cache=shared&_fk=1", DataPath))
	if err != nil {
		return nil, errors.Wrap(err, "failed opening connection to sqlite")
	}
	//defer client.Close()
	// Run the auto migration tool.
	if err := client.Schema.Create(context.Background()); err != nil {
		return nil, errors.Wrap(err, "failed creating schema resources")
	}
	c := &Client{
		ent: client,
	}
	c.init()

	return c, nil
}

func (c *Client) init() {
	c.generateJwtSerectIfNotExist()
	if err := c.generateDefaultLocalStorage(); err != nil {
		log.Errorf("generate default storage: %v", err)
	}

	downloadDir := c.GetSetting(SettingDownloadDir)
	if downloadDir == "" {
		log.Infof("set default download dir")
		c.SetSetting(SettingDownloadDir, "/downloads")
	}
	logLevel := c.GetSetting(SettingLogLevel)
	if logLevel == "" {
		log.Infof("set default log level")
		c.SetSetting(SettingLogLevel, "info")
	}
	if tr := c.GetTransmission(); tr == nil {
		log.Warnf("no download client, set default download client")
		c.SaveTransmission("transmission", "http://transmission:9091", "", "")
	}
}

func (c *Client) generateJwtSerectIfNotExist() {
	v := c.GetSetting(JwtSerectKey)
	if v == "" {
		log.Infof("generate jwt serect")
		key := utils.RandString(32)
		c.SetSetting(JwtSerectKey, key)
	}
}

func (c *Client) generateDefaultLocalStorage() error {
	n, _ := c.ent.Storage.Query().Count(context.TODO())
	if n != 0 {
		return nil
	}
	log.Infof("add default storage")
	return c.AddStorage(&StorageInfo{
		Name:           "local",
		Implementation: "local",
		TvPath:         "/data/tv/",
		MoviePath:      "/data/movies/",
		Default:        true,
	})
}

func (c *Client) GetSetting(key string) string {
	v, err := c.ent.Settings.Query().Where(settings.Key(key)).Only(context.TODO())
	if err != nil {
		log.Debugf("get setting by key: %s error: %v", key, err)
		return ""
	}
	return v.Value
}

func (c *Client) SetSetting(key, value string) error {
	v, err := c.ent.Settings.Query().Where(settings.Key(key)).Only(context.TODO())
	if err != nil {
		log.Infof("create new setting")
		_, err := c.ent.Settings.Create().SetKey(key).SetValue(value).Save(context.TODO())
		return err
	}
	_, err = c.ent.Settings.UpdateOneID(v.ID).SetValue(value).Save(context.TODO())
	return err
}

func (c *Client) GetLanguage() string {
	lang := c.GetSetting(SettingLanguage)
	log.Infof("get application language: %s", lang)
	if lang == "" {
		return "zh-CN"
	}
	return lang
}

func (c *Client) AddMediaWatchlist(m *ent.Media, episodes []int) (*ent.Media, error) {
	count := c.ent.Media.Query().Where(media.TmdbID(m.TmdbID)).CountX(context.Background())
	if count > 0 {
		return nil, fmt.Errorf("tv series %s already in watchlist", m.NameEn)
	}

	if m.StorageID == 0 {
		r, err := c.ent.Storage.Query().Where(storage.And(storage.Default(true), storage.Deleted(false))).First(context.TODO())
		if err == nil {
			log.Infof("use default storage: %v", r.Name)
			m.StorageID = r.ID
		}
	}
	r, err := c.ent.Media.Create().
		SetTmdbID(m.TmdbID).
		SetImdbID(m.ImdbID).
		SetStorageID(m.StorageID).
		SetOverview(m.Overview).
		SetNameCn(m.NameCn).
		SetNameEn(m.NameEn).
		SetOriginalName(m.OriginalName).
		SetMediaType(m.MediaType).
		SetAirDate(m.AirDate).
		SetResolution(m.Resolution).
		SetTargetDir(m.TargetDir).
		SetDownloadHistoryEpisodes(m.DownloadHistoryEpisodes).
		SetLimiter(m.Limiter).
		SetExtras(m.Extras).
		AddEpisodeIDs(episodes...).
		Save(context.TODO())
	return r, err

}

func (c *Client) GetMediaWatchlist(mediaType media.MediaType) []*ent.Media {
	list, err := c.ent.Media.Query().Where(media.MediaTypeEQ(mediaType)).All(context.TODO())
	if err != nil {
		log.Infof("query wtach list error: %v", err)
		return nil
	}
	return list
}

func (c *Client) GetEpisode(seriesId, seasonNum, episodeNum int) (*ent.Episode, error) {
	return c.ent.Episode.Query().Where(episode.MediaID(seriesId), episode.SeasonNumber(seasonNum),
		episode.EpisodeNumber(episodeNum)).First(context.TODO())
}
func (c *Client) GetEpisodeByID(epID int) (*ent.Episode, error) {
	return c.ent.Episode.Query().Where(episode.ID(epID)).First(context.TODO())
}

func (c *Client) UpdateEpiode(episodeId int, name, overview string) error {
	return c.ent.Episode.Update().Where(episode.ID(episodeId)).SetTitle(name).SetOverview(overview).Exec(context.TODO())
}

func (c *Client) UpdateEpiode2(episodeId int, name, overview, airdate string) error {
	return c.ent.Episode.Update().Where(episode.ID(episodeId)).SetTitle(name).SetOverview(overview).SetAirDate(airdate).Exec(context.TODO())
}

type MediaDetails struct {
	*ent.Media
	Episodes []*ent.Episode `json:"episodes"`
}

func (c *Client) GetMediaDetails(id int) *MediaDetails {
	se, err := c.ent.Media.Query().Where(media.ID(id)).First(context.TODO())
	if err != nil {
		log.Errorf("get series %d: %v", id, err)
		return nil
	}
	var md = &MediaDetails{
		Media: se,
	}

	ep, err := se.QueryEpisodes().All(context.Background())
	if err != nil {
		log.Errorf("get episodes %d: %v", id, err)
		return nil
	}
	md.Episodes = ep

	return md
}

func (c *Client) GetMedia(id int) (*ent.Media, error) {
	return c.ent.Media.Query().Where(media.ID(id)).First(context.TODO())
}

func (c *Client) DeleteMedia(id int) error {
	_, err := c.ent.Episode.Delete().Where(episode.MediaID(id)).Exec(context.TODO())
	if err != nil {
		return err
	}
	_, err = c.ent.Media.Delete().Where(media.ID(id)).Exec(context.TODO())
	return err
}

func (c *Client) SaveEposideDetail(d *ent.Episode) (int, error) {
	ep, err := c.ent.Episode.Create().
		SetAirDate(d.AirDate).
		SetSeasonNumber(d.SeasonNumber).
		SetEpisodeNumber(d.EpisodeNumber).
		SetOverview(d.Overview).
		SetMonitored(d.Monitored).
		SetTitle(d.Title).Save(context.TODO())
	if err != nil {
		return 0, errors.Wrap(err, "save episode")
	}
	return ep.ID, nil
}

func (c *Client) SaveEposideDetail2(d *ent.Episode) (int, error) {
	ep, err := c.ent.Episode.Create().
		SetAirDate(d.AirDate).
		SetSeasonNumber(d.SeasonNumber).
		SetEpisodeNumber(d.EpisodeNumber).
		SetMediaID(d.MediaID).
		SetStatus(d.Status).
		SetOverview(d.Overview).
		SetTitle(d.Title).Save(context.TODO())

	return ep.ID, err
}

type TorznabSetting struct {
	URL    string `json:"url"`
	ApiKey string `json:"api_key"`
}

func (c *Client) SaveIndexer(in *ent.Indexers) error {

	if in.ID != 0 {
		//update setting
		return c.ent.Indexers.Update().Where(indexers.ID(in.ID)).SetName(in.Name).SetImplementation(in.Implementation).
			SetPriority(in.Priority).SetSettings(in.Settings).SetSeedRatio(in.SeedRatio).SetDisabled(in.Disabled).Exec(context.Background())
	}
	//create new one
	count := c.ent.Indexers.Query().Where(indexers.Name(in.Name)).CountX(context.TODO())
	if count > 0 {
		return fmt.Errorf("name already esxits: %v", in.Name)
	}

	_, err := c.ent.Indexers.Create().
		SetName(in.Name).SetImplementation(in.Implementation).SetPriority(in.Priority).SetSettings(in.Settings).SetSeedRatio(in.SeedRatio).
		SetDisabled(in.Disabled).Save(context.TODO())
	if err != nil {
		return errors.Wrap(err, "save db")
	}

	return nil
}

func (c *Client) DeleteTorznab(id int) {
	c.ent.Indexers.Delete().Where(indexers.ID(id)).Exec(context.TODO())
}

func (c *Client) GetIndexer(id int) (*TorznabInfo, error) {
	res, err := c.ent.Indexers.Query().Where(indexers.ID(id)).First(context.TODO())
	if err != nil {
		return nil, err
	}
	var ss TorznabSetting
	err = json.Unmarshal([]byte(res.Settings), &ss)
	if err != nil {

		return nil, fmt.Errorf("unmarshal torznab %s error: %v", res.Name, err)
	}
	return &TorznabInfo{Indexers: res, TorznabSetting: ss}, nil
}

type TorznabInfo struct {
	*ent.Indexers
	TorznabSetting
}

func (c *Client) GetAllTorznabInfo() []*TorznabInfo {
	res := c.ent.Indexers.Query().Where(indexers.Implementation(IndexerTorznabImpl)).AllX(context.TODO())

	var l = make([]*TorznabInfo, 0, len(res))
	for _, r := range res {
		var ss TorznabSetting
		err := json.Unmarshal([]byte(r.Settings), &ss)
		if err != nil {
			log.Errorf("unmarshal torznab %s error: %v", r.Name, err)
			continue
		}
		l = append(l, &TorznabInfo{
			Indexers:       r,
			TorznabSetting: ss,
		})
	}
	return l
}

func (c *Client) SaveTransmission(name, url, user, password string) error {
	count := c.ent.DownloadClients.Query().Where(downloadclients.Name(name)).CountX(context.TODO())
	if count != 0 {
		err := c.ent.DownloadClients.Update().Where(downloadclients.Name(name)).
			SetURL(url).SetUser(user).SetPassword(password).Exec(context.TODO())
		return err
	}

	_, err := c.ent.DownloadClients.Create().SetEnable(true).SetImplementation("transmission").
		SetName(name).SetURL(url).SetUser(user).SetPassword(password).Save(context.TODO())
	return err
}

func (c *Client) GetTransmission() *ent.DownloadClients {
	dc, err := c.ent.DownloadClients.Query().Where(downloadclients.Implementation("transmission")).First(context.TODO())
	if err != nil {
		log.Errorf("no transmission client found: %v", err)
		return nil
	}
	return dc
}

func (c *Client) GetAllDonloadClients() []*ent.DownloadClients {
	cc, err := c.ent.DownloadClients.Query().All(context.TODO())
	if err != nil {
		log.Errorf("no download client")
		return nil
	}
	return cc
}

func (c *Client) DeleteDownloadCLient(id int) {
	c.ent.DownloadClients.Delete().Where(downloadclients.ID(id)).Exec(context.TODO())
}

// Storage is the model entity for the Storage schema.
type StorageInfo struct {
	Name           string            `json:"name" binding:"required"`
	Implementation string            `json:"implementation" binding:"required"`
	Settings       map[string]string `json:"settings" binding:"required"`
	TvPath         string            `json:"tv_path" binding:"required"`
	MoviePath      string            `json:"movie_path" binding:"required"`
	Default        bool              `json:"default"`
}

func (s *StorageInfo) ToWebDavSetting() WebdavSetting {
	if s.Implementation != storage.ImplementationWebdav.String() {
		panic("not webdav storage")
	}
	return WebdavSetting{
		URL:            s.Settings["url"],
		User:           s.Settings["user"],
		Password:       s.Settings["password"],
		ChangeFileHash: s.Settings["change_file_hash"],
	}
}

type WebdavSetting struct {
	URL            string `json:"url"`
	User           string `json:"user"`
	Password       string `json:"password"`
	ChangeFileHash string `json:"change_file_hash"`
}

func (c *Client) AddStorage(st *StorageInfo) error {
	if !strings.HasSuffix(st.TvPath, "/") {
		st.TvPath += "/"
	}
	if !strings.HasSuffix(st.MoviePath, "/") {
		st.MoviePath += "/"
	}
	if st.Settings == nil {
		st.Settings = map[string]string{}
	}

	data, err := json.Marshal(st.Settings)
	if err != nil {
		return errors.Wrap(err, "json marshal")
	}

	count := c.ent.Storage.Query().Where(storage.Name(st.Name)).CountX(context.TODO())
	if count > 0 {
		//storage already exist, edit exist one
		return c.ent.Storage.Update().Where(storage.Name(st.Name)).
			SetImplementation(storage.Implementation(st.Implementation)).SetTvPath(st.TvPath).SetMoviePath(st.MoviePath).
			SetSettings(string(data)).Exec(context.TODO())
	}
	countAll := c.ent.Storage.Query().Where(storage.Deleted(false)).CountX(context.TODO())
	if countAll == 0 {
		log.Infof("first storage, make it default: %s", st.Name)
		st.Default = true
	}
	_, err = c.ent.Storage.Create().SetName(st.Name).
		SetImplementation(storage.Implementation(st.Implementation)).SetTvPath(st.TvPath).SetMoviePath(st.MoviePath).
		SetSettings(string(data)).SetDefault(st.Default).Save(context.TODO())
	if err != nil {
		return err
	}
	return nil
}

func (c *Client) GetAllStorage() []*ent.Storage {
	data, err := c.ent.Storage.Query().Where(storage.Deleted(false)).All(context.TODO())
	if err != nil {
		log.Errorf("get storage: %v", err)
		return nil
	}
	return data
}

type Storage struct {
	ent.Storage
}

func (s *Storage) ToWebDavSetting() WebdavSetting {
	if s.Implementation != storage.ImplementationWebdav {
		panic("not webdav storage")
	}
	var webdavSetting WebdavSetting
	json.Unmarshal([]byte(s.Settings), &webdavSetting)
	return webdavSetting
}

func (c *Client) GetStorage(id int) *Storage {
	r, err := c.ent.Storage.Query().Where(storage.ID(id)).First(context.TODO())
	if err != nil {
		//use default storage
		r := c.ent.Storage.Query().Where(storage.Default(true)).FirstX(context.TODO())
		return &Storage{*r}
	}
	return &Storage{*r}
}

func (c *Client) DeleteStorage(id int) error {
	return c.ent.Storage.Update().Where(storage.ID(id)).SetDeleted(true).Exec(context.TODO())
}

func (c *Client) SetDefaultStorage(id int) error {
	err := c.ent.Storage.Update().Where(storage.ID(id)).SetDefault(true).Exec(context.TODO())
	if err != nil {
		return err
	}
	err = c.ent.Storage.Update().Where(storage.Or(storage.ID(id))).SetDefault(false).Exec(context.TODO())
	return err
}

func (c *Client) SetDefaultStorageByName(name string) error {
	err := c.ent.Storage.Update().Where(storage.Name(name)).SetDefault(true).Exec(context.TODO())
	if err != nil {
		return err
	}
	err = c.ent.Storage.Update().Where(storage.Or(storage.Name(name))).SetDefault(false).Exec(context.TODO())
	return err
}

func (c *Client) SaveHistoryRecord(h ent.History) (*ent.History, error) {
	return c.ent.History.Create().SetMediaID(h.MediaID).SetEpisodeID(h.EpisodeID).SetDate(time.Now()).
		SetStatus(h.Status).SetTargetDir(h.TargetDir).SetSourceTitle(h.SourceTitle).SetIndexerID(h.IndexerID).
		SetDownloadClientID(h.DownloadClientID).SetSize(h.Size).SetSaved(h.Saved).Save(context.TODO())
}

func (c *Client) SetHistoryStatus(id int, status history.Status) error {
	return c.ent.History.Update().Where(history.ID(id)).SetStatus(status).Exec(context.TODO())
}

func (c *Client) GetHistories() ent.Histories {
	h, err := c.ent.History.Query().Order(history.ByID(sql.OrderDesc())).All(context.TODO())
	if err != nil {
		return nil
	}
	return h
}

func (c *Client) GetRunningHistories() ent.Histories {
	h, err := c.ent.History.Query().Where(history.Or(history.StatusEQ(history.StatusRunning),
		history.StatusEQ(history.StatusUploading), history.StatusEQ(history.StatusSeeding))).All(context.TODO())
	if err != nil {
		return nil
	}
	return h
}

func (c *Client) GetHistory(id int) *ent.History {
	return c.ent.History.Query().Where(history.ID(id)).FirstX(context.TODO())
}

func (c *Client) DeleteHistory(id int) error {
	_, err := c.ent.History.Delete().Where(history.ID(id)).Exec(context.Background())
	return err
}

func (c *Client) GetDownloadDir() string {
	r, err := c.ent.Settings.Query().Where(settings.Key(SettingDownloadDir)).First(context.TODO())
	if err != nil {
		return "/downloads"
	}
	return r.Value
}

func (c *Client) UpdateEpisodeStatus(mediaID int, seasonNum, episodeNum int) error {
	ep, err := c.ent.Episode.Query().Where(episode.MediaID(mediaID)).Where(episode.EpisodeNumber(episodeNum)).
		Where(episode.SeasonNumber(seasonNum)).First(context.TODO())
	if err != nil {
		return errors.Wrap(err, "finding episode")
	}
	return ep.Update().SetStatus(episode.StatusDownloaded).Exec(context.TODO())
}

func (c *Client) SetEpisodeStatus(id int, status episode.Status) error {
	return c.ent.Episode.Update().Where(episode.ID(id)).SetStatus(status).Exec(context.TODO())
}

func (c *Client) SetSeasonAllEpisodeStatus(mediaID, seasonNum int, status episode.Status) error {
	return c.ent.Episode.Update().Where(episode.MediaID(mediaID), episode.SeasonNumber(seasonNum)).SetStatus(status).Exec(context.TODO())
}

func (c *Client) TmdbIdInWatchlist(tmdb_id int) bool {
	return c.ent.Media.Query().Where(media.TmdbID(tmdb_id)).CountX(context.TODO()) > 0
}

func (c *Client) GetDownloadHistory(mediaID int) ([]*ent.History, error) {
	return c.ent.History.Query().Where(history.MediaID(mediaID)).All(context.TODO())
}

func (c *Client) GetMovieDummyEpisode(movieId int) (*ent.Episode, error) {
	_, err := c.ent.Media.Query().Where(media.ID(movieId), media.MediaTypeEQ(media.MediaTypeMovie)).First(context.TODO())
	if err != nil {
		return nil, errors.Wrap(err, "get movie")
	}
	ep, err := c.ent.Episode.Query().Where(episode.MediaID(movieId)).First(context.TODO())
	if err != nil {
		return nil, errors.Wrap(err, "query episode")
	}
	return ep, nil
}

func (c *Client) GetDownloadClient(id int) (*ent.DownloadClients, error) {
	return c.ent.DownloadClients.Query().Where(downloadclients.ID(id)).First(context.Background())
}

func (c *Client) SetEpisodeMonitoring(id int, b bool) error {
	return c.ent.Episode.Update().Where(episode.ID(id)).SetMonitored(b).Exec(context.Background())
}

type EditMediaData struct {
	ID         int                 `json:"id"`
	Resolution media.Resolution    `json:"resolution"`
	TargetDir  string              `json:"target_dir"`
	Limiter    schema.MediaLimiter `json:"limiter"`
}

func (c *Client) EditMediaMetadata(in EditMediaData) error {
	return c.ent.Media.Update().Where(media.ID(in.ID)).SetResolution(in.Resolution).SetTargetDir(in.TargetDir).SetLimiter(in.Limiter).
		Exec(context.Background())
}

func (c *Client) UpdateEpisodeTargetFile(id int, filename string) error {
	return c.ent.Episode.Update().Where(episode.ID(id)).SetTargetFile(filename).Exec(context.Background())
}

func (c *Client) GetSeasonEpisodes(mediaId, seasonNum int) ([]*ent.Episode, error) {
	return c.ent.Episode.Query().Where(episode.MediaID(mediaId), episode.SeasonNumber(seasonNum)).All(context.Background())
}

func (c *Client) GetAllImportLists() ([]*ent.ImportList, error) {
	return c.ent.ImportList.Query().All(context.Background())
}

func (c *Client) AddImportlist(il *ent.ImportList) error {
	count, err := c.ent.ImportList.Query().Where(importlist.Name(il.Name)).Count(context.Background())
	if err != nil {
		return err
	}
	if count > 0 {
		//edit exist record
		return c.ent.ImportList.Update().Where(importlist.Name(il.Name)).
			SetURL(il.URL).SetQulity(il.Qulity).SetType(il.Type).SetStorageID(il.StorageID).Exec(context.Background())
	}
	return c.ent.ImportList.Create().SetName(il.Name).SetURL(il.URL).SetQulity(il.Qulity).SetStorageID(il.StorageID).
		SetType(il.Type).Exec(context.Background())
}

func (c *Client) DeleteImportlist(id int) error {
	return c.ent.ImportList.DeleteOneID(id).Exec(context.TODO())
}