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" json:"id"`
	Type           int                `xorm:"INDEX NOT NULL" json:"type"`           //0 normal  5官方推荐,中间值保留为后续扩展
	CloudbrainType int                `xorm:"INDEX NOT NULL" json:"cloudbrainType"` //0 云脑一   1云脑二
	UID            int64              `xorm:"INDEX NOT NULL" json:"uid"`
	IsPrivate      bool               `xorm:"INDEX NOT NULL" json:"isPrivate"`
	Tag            string             `xorm:"varchar(100) UNIQUE" json:"tag"`
	Description    string             `xorm:"varchar(765)" json:"description"`
	Topics         []string           `xorm:"TEXT JSON" json:"topics"`
	Place          string             `xorm:"varchar(300)" json:"place"`
	NumStars       int                `xorm:"NOT NULL DEFAULT 0" json:"numStars"`
	Creator        *User              `xorm:"-" json:"creator"`
	CreatedUnix    timeutil.TimeStamp `xorm:"INDEX created" json:"createdUnix"`
	UpdatedUnix    timeutil.TimeStamp `xorm:"INDEX updated" json:"updatedUnix"`
}

type ImageList []*Image

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
	IncludePublicOnly   bool
	IncludeOfficialOnly bool
	IncludePrivateOnly  bool
	IncludeStarByMe     bool
	IncludeCustom       bool
	IncludeOwnerOnly    bool
	Topics              string
	ListOptions
	SearchOrderBy
}
type ErrorImageTagExist struct {
	Tag string
}

type ImagesPageResult struct {
	Count  int64    `json:"count"`
	Images []*Image `json:"images"`
}

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) (ImageList, int64, error) {
	cond := SearchImageCondition(opts)
	return SearchImageByCondition(opts, cond)
}

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

	if len(opts.Keyword) > 0 {

		var subQueryCond = builder.NewCond()
		for _, v := range strings.Split(opts.Keyword, ",") {

			subQueryCond = subQueryCond.Or(builder.Like{"LOWER(image_topic.name)", strings.ToLower(v)})

		}
		subQuery := builder.Select("image_topic_relation.image_id").From("image_topic_relation").
			Join("INNER", "image_topic", "image_topic.id = image_topic_relation.image_id").
			Where(subQueryCond).
			GroupBy("image_topic_relation.image_id")
		var keywordCond = builder.In("id", subQuery)

		var likes = builder.NewCond()
		for _, v := range strings.Split(opts.Keyword, ",") {
			likes = likes.Or(builder.Like{"LOWER(tag)", strings.ToLower(v)})

			likes = likes.Or(builder.Like{"LOWER(description)", strings.ToLower(v)})

		}
		keywordCond = keywordCond.Or(likes)

		cond = cond.And(keywordCond)

	}
	if len(opts.Topics) > 0 { //标签精确匹配
		var subQueryCond = builder.NewCond()
		for _, v := range strings.Split(opts.Keyword, ",") {

			subQueryCond = subQueryCond.Or(builder.Eq{"LOWER(image_topic.name)": strings.ToLower(v)})
			subQuery := builder.Select("image_topic_relation.image_id").From("image_topic_relation").
				Join("INNER", "image_topic", "image_topic.id = image_topic_relation.image_id").
				Where(subQueryCond).
				GroupBy("image_topic_relation.image_id")
			var topicCond = builder.In("id", subQuery)
			cond = cond.And(topicCond)
		}
	}

	if opts.IncludePublicOnly {
		cond = cond.And(builder.Eq{"is_private": false})
	}

	if opts.IncludePrivateOnly {
		cond = cond.And(builder.Eq{"is_private": true})
	}

	if opts.IncludeOwnerOnly {

		cond = cond.And(builder.Eq{"uid": opts.UID})
	}
	if opts.IncludeOfficialOnly {
		cond = cond.And(builder.Eq{"type": RECOMMOND_TYPE})
	}

	if opts.IncludeStarByMe {

		subQuery := builder.Select("image_id").From("image_star").
			Where(builder.Eq{"uid": opts.UID})
		var starCond = builder.In("id", subQuery)
		cond = cond.And(starCond)

	}

	return cond
}

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

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

	images := make(ImageList, 0, opts.PageSize)
	count, err := sess.Where(cond).Count(new(Image))

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

	sess.Where(cond).OrderBy(opts.SearchOrderBy.String())

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

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

	return images, count, nil
}

func (images ImageList) loadAttributes(e Engine) error {
	if len(images) == 0 {
		return nil
	}

	set := make(map[int64]struct{})

	for i := range images {
		set[images[i].UID] = struct{}{}
	}

	// Load creators.
	users := make(map[int64]*User, len(set))
	if err := e.Table("\"user\"").
		Cols("name", "lower_name", "avatar", "email").
		Where("id > 0").
		In("id", keysInt64(set)).
		Find(&users); err != nil {
		return fmt.Errorf("find users: %v", err)
	}

	for i := range images {
		images[i].Creator = users[images[i].UID]

	}

	return nil
}

func CreateLocalImage(image *Image) error {

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

func UpdateLocalImage(image *Image) error {

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

func DeleteLocalImage(id int64) error {
	image := new(Image)
	_, err := x.ID(id).Delete(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
	}

}