aiforge/models/cloudbrain_image.go

package models

import (
	"fmt"
	"strings"

	"xorm.io/builder"

	"code.gitea.io/gitea/modules/timeutil"
)

const RECOMMOND_TYPE = 5
const NORMAL_TYPE = 0

type Image struct {
	ID             int64              `xorm:"pk autoincr"`
	Type           int                `xorm:"INDEX NOT NULL"` //0 normal  5官方推荐,中间值保留为后续扩展
	CloudbrainType int                `xorm:"INDEX NOT NULL"` //0 云脑一   1云脑二
	UID            int64              `xorm:"INDEX NOT NULL"`
	IsPrivate      bool               `xorm:"INDEX NOT NULL"`
	Tag            string             `xorm:"varchar(100) UNIQUE"`
	Description    string             `xorm:"varchar(765)"`
	Topics         []string           `xorm:"TEXT JSON"`
	Place          string             `xorm:"varchar(300)"`
	NumStars       int                `xorm:"NOT NULL DEFAULT 0"`
	CreatedUnix    timeutil.TimeStamp `xorm:"INDEX created"`
	UpdatedUnix    timeutil.TimeStamp `xorm:"INDEX updated"`
}

type ImageStar struct {
	ID          int64              `xorm:"pk autoincr"`
	UID         int64              `xorm:"UNIQUE(s)"`
	ImageID     int64              `xorm:"UNIQUE(s)"`
	CreatedUnix timeutil.TimeStamp `xorm:"created"`
}

type ImageTopic struct {
	ID          int64
	Name        string `xorm:"UNIQUE VARCHAR(105)"`
	ImageCount  int
	CreatedUnix timeutil.TimeStamp `xorm:"INDEX created"`
	UpdatedUnix timeutil.TimeStamp `xorm:"INDEX updated"`
}

type ImageTopicRelation struct {
	ImageID int64 `xorm:"UNIQUE(s)"`
	TopicID int64 `xorm:"UNIQUE(s)"`
}

type SearchImageOptions struct {
	Keyword         string
	UID             int64
	IncludePublic   bool
	IncludeStarByMe bool
	IncludeCustom   bool
	CodeLanguage    string
	Framework       string
	CudaVersion     string
	ListOptions
	SearchOrderBy
}
type ErrorImageTagExist struct {
	Tag string
}

func (err ErrorImageTagExist) Error() string {
	return fmt.Sprintf("Image already exists [tag: %s]", err.Tag)
}

type ErrImageNotExist struct {
	ID int64
}

func (err ErrImageNotExist) Error() string {
	return fmt.Sprintf("Image does not exist [id: %d]", err.ID)
}

func IsErrImageTagExist(err error) bool {
	_, ok := err.(ErrorImageTagExist)
	return ok
}

func IsImageExist(tag string) (bool, error) {
	return x.Exist(&Image{
		Tag: tag,
	})
}

type FindImageTopicOptions struct {
	ListOptions
	ImageID int64
	Keyword string
}

func (opts *FindImageTopicOptions) toConds() builder.Cond {
	var cond = builder.NewCond()
	if opts.ImageID > 0 {
		cond = cond.And(builder.Eq{"image_topic.image_id": opts.ImageID})
	}

	if opts.Keyword != "" {
		cond = cond.And(builder.Like{"image_topic.name", strings.ToLower(opts.Keyword)})
	}

	return cond
}

func GetImageByID(id int64) (*Image, error) {
	rel := new(Image)
	has, err := x.
		ID(id).
		Get(rel)
	if err != nil {
		return nil, err
	} else if !has {
		return nil, ErrImageNotExist{id}
	}

	return rel, nil
}
func FindImageTopics(opts *FindImageTopicOptions) (topics []*ImageTopic, err error) {
	sess := x.Select("image_topic.*").Where(opts.toConds())
	if opts.ImageID > 0 {
		sess.Join("INNER", "image_topic_relation", "image_topic_relation.topic_id = image_topic.id")
	}
	if opts.PageSize != 0 && opts.Page != 0 {
		sess = opts.setSessionPagination(sess)
	}
	return topics, sess.Desc("image_topic.image_count").Find(&topics)
}

func SaveImageTopics(imageID int64, topicNames ...string) error {
	topics, err := FindImageTopics(&FindImageTopicOptions{
		ImageID: imageID,
	})
	if err != nil {
		return err
	}

	sess := x.NewSession()
	defer sess.Close()

	if err := sess.Begin(); err != nil {
		return err
	}

	var addedTopicNames []string
	for _, topicName := range topicNames {
		if strings.TrimSpace(topicName) == "" {
			continue
		}

		var found bool
		for _, t := range topics {
			if strings.EqualFold(topicName, t.Name) {
				found = true
				break
			}
		}
		if !found {
			addedTopicNames = append(addedTopicNames, topicName)
		}
	}

	var removeTopics []*ImageTopic
	for _, t := range topics {
		var found bool
		for _, topicName := range topicNames {
			if strings.EqualFold(topicName, t.Name) {
				found = true
				break
			}
		}
		if !found {
			removeTopics = append(removeTopics, t)
		}
	}

	for _, topicName := range addedTopicNames {
		_, err := addTopicByNameToImage(sess, imageID, topicName)
		if err != nil {
			return err
		}
	}

	for _, topic := range removeTopics {
		err := removeTopicFromImage(sess, imageID, topic)
		if err != nil {
			return err
		}
	}

	topicNames = make([]string, 0, 25)
	if err := sess.Table("image_topic").Cols("name").
		Join("INNER", "image_topic_relation", "image_topic_relation.topic_id = image_topic.id").
		Where("image_topic_relation.image_id = ?", imageID).Desc("image_topic.image_count").Find(&topicNames); err != nil {
		return err
	}

	if _, err := sess.ID(imageID).Cols("topics").Update(&Image{
		Topics: topicNames,
	}); err != nil {
		return err
	}

	return sess.Commit()
}

func addTopicByNameToImage(e Engine, imageID int64, topicName string) (*ImageTopic, error) {
	var topic ImageTopic
	has, err := e.Where("name = ?", topicName).Get(&topic)
	if err != nil {
		return nil, err
	}
	if !has {
		topic.Name = topicName
		topic.ImageCount = 1
		if _, err := e.Insert(&topic); err != nil {
			return nil, err
		}
	} else {
		topic.ImageCount++
		if _, err := e.ID(topic.ID).Cols("image_count").Update(&topic); err != nil {
			return nil, err
		}
	}

	if _, err := e.Insert(&ImageTopicRelation{
		ImageID: imageID,
		TopicID: topic.ID,
	}); err != nil {
		return nil, err
	}

	return &topic, nil
}

func removeTopicFromImage(e Engine, imageId int64, topic *ImageTopic) error {
	topic.ImageCount--
	if _, err := e.ID(topic.ID).Cols("image_count").Update(topic); err != nil {
		return err
	}

	if _, err := e.Delete(&ImageTopicRelation{
		ImageID: imageId,
		TopicID: topic.ID,
	}); err != nil {
		return err
	}

	return nil
}

/*func SearchImage(opts *SearchImageOptions) ([]*Image, int64, error) {
	cond := SearchImageCondition(opts)
	return SearchImageByCondition(opts, cond)
}*/

func SearchImageCondition(opts *SearchImageOptions) builder.Cond {
	var cond = builder.NewCond()

	if len(opts.Keyword) > 0 {
		cond = cond.And(builder.Or(builder.Like{"image.tag", opts.Keyword}, builder.Like{"image.description", opts.Keyword}))
	}

	if len(opts.CudaVersion) > 0 {
		cond = cond.And(builder.Eq{"image.cuda_version": opts.CudaVersion})
	}

	if len(opts.CodeLanguage) > 0 {
		cond = cond.And(builder.Eq{"image.code_language": opts.CodeLanguage})
	}
	if len(opts.Framework) > 0 {
		cond = cond.And(builder.Eq{"image.framework": opts.Framework})
	}

	if opts.IncludePublic {
		cond = cond.And(builder.Eq{"image.is_private": false})
	}
	/**
	if opts.IncludeStarByMe {
		cond = cond.And(builder.Eq{"dataset.status": DatasetStatusPublic})
		cond = cond.And(builder.Eq{"attachment.is_private": false})
		if opts.OwnerID > 0 {
			if len(opts.Keyword) == 0 {
				cond = cond.Or(builder.Eq{"repository.owner_id": opts.OwnerID})
			} else {
				subCon := builder.NewCond()
				subCon = subCon.And(builder.Eq{"repository.owner_id": opts.OwnerID}, builder.Or(builder.Like{"dataset.title", opts.Keyword}, builder.Like{"dataset.description", opts.Keyword}))
				cond = cond.Or(subCon)

			}
		}
	} else if opts.OwnerID > 0 {
		cond = cond.And(builder.Eq{"repository.owner_id": opts.OwnerID})
		if !opts.IsOwner {
			cond = cond.And(builder.Eq{"dataset.status": DatasetStatusPublic})
			cond = cond.And(builder.Eq{"attachment.is_private": false})
		}
	}*/

	return cond
}

/*func SearchImageByCondition(opts *SearchImageOptions, cond builder.Cond) ([]*Image, int64, error) {
	if opts.Page <= 0 {
		opts.Page = 1
	}

	var err error
	sess := x.NewSession()
	defer sess.Close()

	datasets := make(DatasetList, 0, opts.PageSize)
	selectColumnsSql := "distinct dataset.id,dataset.title, dataset.status, dataset.category, dataset.description, dataset.download_times, dataset.license, dataset.task, dataset.release_id, dataset.user_id, dataset.repo_id, dataset.created_unix,dataset.updated_unix,dataset.num_stars"

	count, err := sess.Distinct("dataset.id").Join("INNER", "repository", "repository.id = dataset.repo_id").
		Join("INNER", "attachment", "attachment.dataset_id=dataset.id").
		Where(cond).Count(new(Dataset))

	if err != nil {
		return nil, 0, fmt.Errorf("Count: %v", err)
	}

	sess.Select(selectColumnsSql).Join("INNER", "repository", "repository.id = dataset.repo_id").
		Join("INNER", "attachment", "attachment.dataset_id=dataset.id").
		Where(cond).OrderBy(opts.SearchOrderBy.String())

	if opts.PageSize > 0 {
		sess.Limit(opts.PageSize, (opts.Page-1)*opts.PageSize)
	}
	if err = sess.Find(&datasets); err != nil {
		return nil, 0, fmt.Errorf("Dataset: %v", err)
	}

	if err = datasets.loadAttributes(sess); err != nil {
		return nil, 0, fmt.Errorf("LoadAttributes: %v", err)
	}

	return datasets, count, nil
}*/

func CreateLocalImage(image *Image) error {

	_, err := x.Insert(image)
	return err
}

//star or unstar Image
func StarImage(userID, imageID int64, star bool) error {
	sess := x.NewSession()
	defer sess.Close()

	if err := sess.Begin(); err != nil {
		return err
	}

	if star {
		if isImageStaring(sess, userID, imageID) {
			return nil
		}

		if _, err := sess.Insert(&ImageStar{UID: userID, ImageID: imageID}); err != nil {
			return err
		}
		if _, err := sess.Exec("UPDATE `image` SET num_stars = num_stars + 1 WHERE id = ?", imageID); err != nil {
			return err
		}
		if _, err := sess.Exec("UPDATE `user` SET num_image_stars = num_image_stars + 1 WHERE id = ?", userID); err != nil {
			return err
		}
	} else {
		if !isImageStaring(sess, userID, imageID) {
			return nil
		}

		if _, err := sess.Delete(&ImageStar{0, userID, imageID, 0}); err != nil {
			return err
		}
		if _, err := sess.Exec("UPDATE `image` SET num_stars = num_stars - 1 WHERE id = ?", imageID); err != nil {
			return err
		}
		if _, err := sess.Exec("UPDATE `user` SET num_image_stars = num_image_stars - 1 WHERE id = ?", userID); err != nil {
			return err
		}
	}

	return sess.Commit()
}

func IsImageStaring(userID, datasetID int64) bool {
	return isImageStaring(x, userID, datasetID)

}

func isImageStaring(e Engine, userID, imageID int64) bool {
	has, _ := e.Get(&ImageStar{0, userID, imageID, 0})
	return has
}
func RecommendImage(imageId int64, recommond bool) error {

	image := Image{Type: getRecommondType(recommond)}
	_, err := x.ID(imageId).Cols("type").Update(image)
	return err
}

func getRecommondType(recommond bool) int {
	if recommond {

		return RECOMMOND_TYPE
	} else {
		return NORMAL_TYPE
	}

}