From f947ee6fc781b62dd2a1789979d5c6ecc1af0cd7 Mon Sep 17 00:00:00 2001
From: krau <71133316+krau@users.noreply.github.com>
Date: Sun, 28 Sep 2025 16:49:32 +0800
Subject: [PATCH] fix: implement filename strategy in userbot listen mode

---
 client/bot/handlers/register.go | 58 +++++++++++++++++++++++++++++++++
 pkg/tfile/tgfile.go             |  5 +++
 2 files changed, 63 insertions(+)

diff --git a/client/bot/handlers/register.go b/client/bot/handlers/register.go
index ca4be2c..d7304ff 100644
--- a/client/bot/handlers/register.go
+++ b/client/bot/handlers/register.go
@@ -1,9 +1,12 @@
 package handlers
 
 import (
+	"fmt"
 	"path"
 	"regexp"
 	"strings"
+	"text/template"
+	"time"
 
 	"github.com/celestix/gotgproto/dispatcher"
 	"github.com/celestix/gotgproto/dispatcher/handlers"
@@ -11,14 +14,17 @@ import (
 	"github.com/celestix/gotgproto/ext"
 	"github.com/charmbracelet/log"
 	sabotfilters "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/filters"
+	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/mediautil"
 	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/re"
 	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/ruleutil"
 	userclient "github.com/krau/SaveAny-Bot/client/user"
+	"github.com/krau/SaveAny-Bot/common/utils/strutil"
 	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/core"
 	"github.com/krau/SaveAny-Bot/core/tasks/tfile"
 	"github.com/krau/SaveAny-Bot/database"
+	"github.com/krau/SaveAny-Bot/pkg/enums/fnamest"
 	"github.com/krau/SaveAny-Bot/pkg/tcbdata"
 	"github.com/krau/SaveAny-Bot/storage"
 	"github.com/rs/xid"
@@ -118,6 +124,58 @@ func listenMediaMessageEvent(ch chan userclient.MediaMessageEvent) {
 				logger.Errorf("Failed to get storage by user ID %d and name %s: %v", user.ChatID, user.DefaultStorage, err)
 				continue
 			}
+			switch user.FilenameStrategy {
+			case fnamest.Message.String():
+				file.SetName(tgutil.GenFileNameFromMessage(*file.Message()))
+			case fnamest.Template.String():
+				if user.FilenameTemplate == "" {
+					logger.Warnf("Empty filename template for user %d, using default filename", user.ChatID)
+					break
+				}
+				// [TODO] refactor this
+				message := file.Message()
+				tmpl, err := template.New("filename").Parse(user.FilenameTemplate)
+				if err != nil {
+					logger.Errorf("Failed to parse filename template for user %d: %s", user.ChatID, err)
+					break
+				}
+				data := mediautil.FilenameTemplateData{
+					MsgID: func() string {
+						id := message.GetID()
+						if id == 0 {
+							return ""
+						}
+						return fmt.Sprintf("%d", id)
+					}(),
+					MsgTags: func() string {
+						tags := strutil.ExtractTagsFromText(message.GetMessage())
+						if len(tags) == 0 {
+							return ""
+						}
+						return strings.Join(tags, "_")
+					}(),
+					MsgGen: tgutil.GenFileNameFromMessage(*message),
+					OrigName: func() string {
+						f, _ := tgutil.GetMediaFileName(message.Media)
+						return f
+					}(),
+					MsgDate: func() string {
+						date := message.GetDate()
+						if date == 0 {
+							return ""
+						}
+						t := time.Unix(int64(date), 0)
+						return t.Format("2006-01-02_15-04-05")
+					}(),
+				}.ToMap()
+				var sb strings.Builder
+				err = tmpl.Execute(&sb, data)
+				if err != nil {
+					log.FromContext(ctx).Errorf("failed to execute filename template: %s", err)
+					break
+				}
+				file.SetName(sb.String())
+			}
 			var dirPath string
 			if user.ApplyRule && user.Rules != nil {
 				matched, matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
diff --git a/pkg/tfile/tgfile.go b/pkg/tfile/tgfile.go
index 0465ed7..5c3ed9b 100644
--- a/pkg/tfile/tgfile.go
+++ b/pkg/tfile/tgfile.go
@@ -14,6 +14,7 @@ type TGFile interface {
 	Dler() downloader.Client // witch client to use for downloading
 	Size() int64
 	Name() string
+	SetName(name string)
 }
 
 type TGFileMessage interface {
@@ -29,6 +30,10 @@ type tgFile struct {
 	dler     downloader.Client
 }
 
+func (f *tgFile) SetName(name string) {
+	f.name = name
+}
+
 func (f *tgFile) Location() tg.InputFileLocationClass {
 	return f.location
 }