From d703f11ea075396c6c8eb65f659433e4822be51e Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Fri, 4 Jul 2025 16:01:29 +0800 Subject: [PATCH] feat: implement album handling rules and refactor related logic --- client/bot/handlers/utils/ruleutil/rule.go | 40 +++++++++- client/bot/handlers/utils/shortcut/tftask.go | 80 ++++++++++++++++---- pkg/consts/specific.go | 3 +- pkg/enums/rule/ruletype.go | 3 +- pkg/rule/is_album.go | 38 ++++++++++ 5 files changed, 142 insertions(+), 22 deletions(-) create mode 100644 pkg/rule/is_album.go diff --git a/client/bot/handlers/utils/ruleutil/rule.go b/client/bot/handlers/utils/ruleutil/rule.go index 3c5cf34..28a321c 100644 --- a/client/bot/handlers/utils/ruleutil/rule.go +++ b/client/bot/handlers/utils/ruleutil/rule.go @@ -3,6 +3,8 @@ package ruleutil import ( "context" + "github.com/duke-git/lancet/v2/convertor" + "github.com/charmbracelet/log" "github.com/krau/SaveAny-Bot/database" "github.com/krau/SaveAny-Bot/pkg/consts" @@ -33,11 +35,22 @@ func (m matchedStorName) String() string { return string(m) } -func (m matchedStorName) IsValid() bool { +// can we use this storage name directly? +func (m matchedStorName) IsUsable() bool { return m != "" && m != consts.RuleStorNameChosen } -func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (matchedStorageName matchedStorName, dirPath string) { +type MatchedDirPath string + +func (m MatchedDirPath) String() string { + return string(m) +} + +func (m MatchedDirPath) NeedNewForAlbum() bool { + return m != "" && m == consts.RuleDirPathNewForAlbum +} + +func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (matchedStorageName matchedStorName, dirPath MatchedDirPath) { if inputs == nil || len(rules) == 0 { return "", "" } @@ -56,7 +69,7 @@ func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (m continue } if ok { - dirPath = ru.StoragePath() + dirPath = MatchedDirPath(ru.StoragePath()) matchedStorageName = matchedStorName(ru.StorageName()) } case ruleenum.MessageRegex.String(): @@ -71,7 +84,26 @@ func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (m continue } if ok { - dirPath = ru.StoragePath() + dirPath = MatchedDirPath(ru.StoragePath()) + matchedStorageName = matchedStorName(ru.StorageName()) + } + case ruleenum.IsAlbum.String(): + matchAlbum, err := convertor.ToBool(ur.Data) + if err != nil { + matchAlbum = false + } + ru, err := rule.NewRuleMediaType(ur.StorageName, ur.DirPath, matchAlbum) + if err != nil { + logger.Errorf("Failed to create rule: %s", err) + continue + } + ok, err := ru.Match(inputs.File.Message().GroupedID != 0) + if err != nil { + logger.Errorf("Failed to match rule: %s", err) + continue + } + if ok { + dirPath = MatchedDirPath(ru.StoragePath()) matchedStorageName = matchedStorName(ru.StorageName()) } } diff --git a/client/bot/handlers/utils/shortcut/tftask.go b/client/bot/handlers/utils/shortcut/tftask.go index abe8a27..1208c20 100644 --- a/client/bot/handlers/utils/shortcut/tftask.go +++ b/client/bot/handlers/utils/shortcut/tftask.go @@ -3,6 +3,7 @@ package shortcut import ( "fmt" "path" + "strings" "github.com/celestix/gotgproto/dispatcher" "github.com/celestix/gotgproto/ext" @@ -34,8 +35,8 @@ func CreateAndAddTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage } if user.ApplyRule && user.Rules != nil { matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file)) - dirPath = matchedDirPath - if matchedStorageName.IsValid() { + dirPath = matchedDirPath.String() + if matchedStorageName.IsUsable() { stor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, matchedStorageName.String()) if err != nil { logger.Errorf("Failed to get storage by user ID and name: %s", err) @@ -93,19 +94,28 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st }) return dispatcher.EndGroups } + useRule := user.ApplyRule && user.Rules != nil - applyRule := func(file tfile.TGFileMessage) (string, string) { + + applyRule := func(file tfile.TGFileMessage) (string, ruleutil.MatchedDirPath) { if !useRule { - return stor.Name(), dirPath + return stor.Name(), ruleutil.MatchedDirPath(dirPath) } storName, dirP := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file)) - if !storName.IsValid() { - return stor.Name(), dirP + + storname := storName.String() + if !storName.IsUsable() { + storname = stor.Name() } - return storName.String(), dirP + return storname, dirP } elems := make([]batchtftask.TaskElement, 0, len(files)) + type albumFile struct { + file tfile.TGFileMessage + storage storage.Storage + } + albumFiles := make(map[int64][]albumFile, 0) for _, file := range files { storName, dirPath := applyRule(file) fileStor := stor @@ -120,18 +130,56 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st return dispatcher.EndGroups } } - storPath := fileStor.JoinStoragePath(path.Join(dirPath, file.Name())) - elem, err := batchtftask.NewTaskElement(fileStor, storPath, file) - if err != nil { - logger.Errorf("Failed to create task element: %s", err) - ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{ - ID: trackMsgID, - Message: "任务创建失败: " + err.Error(), + if !dirPath.NeedNewForAlbum() { + storPath := fileStor.JoinStoragePath(path.Join(dirPath.String(), file.Name())) + elem, err := batchtftask.NewTaskElement(fileStor, storPath, file) + if err != nil { + logger.Errorf("Failed to create task element: %s", err) + ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{ + ID: trackMsgID, + Message: "任务创建失败: " + err.Error(), + }) + return dispatcher.EndGroups + } + elems = append(elems, *elem) + } else { + groupId, isGroup := file.Message().GetGroupedID() + if !isGroup || groupId == 0 { + logger.Warnf("File %s is not in a group, skipping album handling", file.Name()) + continue + } + if _, ok := albumFiles[groupId]; !ok { + albumFiles[groupId] = make([]albumFile, 0) + } + albumFiles[groupId] = append(albumFiles[groupId], albumFile{ + file: file, + storage: fileStor, }) - return dispatcher.EndGroups } - elems = append(elems, *elem) } + for _, afiles := range albumFiles { + if len(afiles) <= 1 { + continue + } + // 对于需要新建目录的文件, 将第一个文件的文件名(去除扩展名)作为目录名 + // 存储以第一个文件的存储为准 + albumDir := strings.TrimSuffix(path.Base(afiles[0].file.Name()), path.Ext(afiles[0].file.Name())) + albumStor := afiles[0].storage + for _, af := range afiles { + afstorPath := af.storage.JoinStoragePath(path.Join(dirPath, albumDir, af.file.Name())) + elem, err := batchtftask.NewTaskElement(albumStor, afstorPath, af.file) + if err != nil { + logger.Errorf("Failed to create task element for album file: %s", err) + ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{ + ID: trackMsgID, + Message: "任务创建失败: " + err.Error(), + }) + return dispatcher.EndGroups + } + elems = append(elems, *elem) + } + } + injectCtx := tgutil.ExtWithContext(ctx.Context, ctx) taskid := xid.New().String() task := batchtftask.NewBatchTGFileTask(taskid, injectCtx, elems, batchtftask.NewProgressTracker(trackMsgID, userID), true) diff --git a/pkg/consts/specific.go b/pkg/consts/specific.go index 3db7125..87d7140 100644 --- a/pkg/consts/specific.go +++ b/pkg/consts/specific.go @@ -1,5 +1,6 @@ package consts const ( - RuleStorNameChosen = "CHOSEN" + RuleStorNameChosen = "CHOSEN" + RuleDirPathNewForAlbum = "NEW-FOR-ALBUM" // create a new directory for album files ) diff --git a/pkg/enums/rule/ruletype.go b/pkg/enums/rule/ruletype.go index 62b88ee..201ca5a 100644 --- a/pkg/enums/rule/ruletype.go +++ b/pkg/enums/rule/ruletype.go @@ -5,6 +5,7 @@ type RuleType string const ( FileNameRegex RuleType = "FILENAME-REGEX" MessageRegex RuleType = "MESSAGE-REGEX" + IsAlbum RuleType = "IS-ALBUM" ) func (r RuleType) String() string { @@ -12,5 +13,5 @@ func (r RuleType) String() string { } func Values() []RuleType { - return []RuleType{FileNameRegex, MessageRegex} + return []RuleType{FileNameRegex, MessageRegex, IsAlbum} } diff --git a/pkg/rule/is_album.go b/pkg/rule/is_album.go new file mode 100644 index 0000000..f86bc96 --- /dev/null +++ b/pkg/rule/is_album.go @@ -0,0 +1,38 @@ +package rule + +import ( + ruleenum "github.com/krau/SaveAny-Bot/pkg/enums/rule" +) + +var _ RuleClass[bool] = (*RuleMediaType)(nil) + +type RuleMediaType struct { + storInfo + matchAlbum bool +} + +func (r RuleMediaType) Type() ruleenum.RuleType { + return ruleenum.IsAlbum +} + +func (r RuleMediaType) Match(input bool) (bool, error) { + return r.matchAlbum == input, nil +} + +func (r RuleMediaType) StorageName() string { + return r.storName +} + +func (r RuleMediaType) StoragePath() string { + return r.storPath +} + +func NewRuleMediaType(storName, storPath string, matchAlbum bool) (*RuleMediaType, error) { + return &RuleMediaType{ + storInfo: storInfo{ + storName: storName, + storPath: storPath, + }, + matchAlbum: matchAlbum, + }, nil +}