diff --git a/common/utils/tgutil/message.go b/common/utils/tgutil/message.go index 157c365..0e30845 100644 --- a/common/utils/tgutil/message.go +++ b/common/utils/tgutil/message.go @@ -2,6 +2,7 @@ package tgutil import ( "fmt" + "sort" "strconv" "strings" "unicode" @@ -359,9 +360,16 @@ func GetGroupedMessages(ctx *ext.Context, chatID int64, msg *tg.Message) ([]*tg. groupedMessages = append(groupedMessages, m) } } + sortMessagesByID(groupedMessages) return groupedMessages, nil } +func sortMessagesByID(messages []*tg.Message) { + sort.Slice(messages, func(i, j int) bool { + return messages[i].GetID() < messages[j].GetID() + }) +} + func ExtractMessageEntityUrls(msg *tg.Message) []string { if len(msg.Entities) == 0 { return nil diff --git a/common/utils/tgutil/message_test.go b/common/utils/tgutil/message_test.go new file mode 100644 index 0000000..3823101 --- /dev/null +++ b/common/utils/tgutil/message_test.go @@ -0,0 +1,18 @@ +package tgutil + +import ( + "testing" + + "github.com/gotd/td/tg" +) + +func TestSortMessagesByID(t *testing.T) { + messages := []*tg.Message{{ID: 9}, {ID: 3}, {ID: 7}} + sortMessagesByID(messages) + want := []int{3, 7, 9} + for i := range messages { + if messages[i].GetID() != want[i] { + t.Fatalf("message %d has ID %d, want %d", i, messages[i].GetID(), want[i]) + } + } +} diff --git a/core/tasks/batchtfile/execute.go b/core/tasks/batchtfile/execute.go index c686919..57c4a7f 100644 --- a/core/tasks/batchtfile/execute.go +++ b/core/tasks/batchtfile/execute.go @@ -14,36 +14,47 @@ import ( "github.com/krau/SaveAny-Bot/common/utils/ioutil" "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/pkg/enums/ctxkey" + "github.com/krau/SaveAny-Bot/pkg/storagetypes" "github.com/krau/SaveAny-Bot/pkg/taskevent" + "github.com/krau/SaveAny-Bot/storage" "golang.org/x/sync/errgroup" ) +type executionGroup struct { + elems []*TaskElement + batchSaver storage.StorageBatchSaver +} + +func (g executionGroup) usesBatchSaver() bool { + return g.batchSaver != nil +} + func (t *Task) Execute(ctx context.Context) error { logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("batch_file[%s]", t.ID)) logger.Info("Starting batch file task") t.Progress.OnStart(ctx, t) - workers := config.C().Workers - eg, gctx := errgroup.WithContext(ctx) - eg.SetLimit(workers) - for _, elem := range t.elems { - eg.Go(func() error { - t.processingMu.RLock() - if t.processing[elem.ID] != nil { - return fmt.Errorf("element with ID %s is already being processed", elem.ID) + groups := t.executionGroups() + var err error + for i := 0; i < len(groups); { + if groups[i].usesBatchSaver() { + err = t.processBatch(ctx, groups[i]) + i++ + } else { + end := i + 1 + for end < len(groups) && !groups[end].usesBatchSaver() { + end++ } - t.processingMu.RUnlock() - t.processingMu.Lock() - t.processing[elem.ID] = &elem - t.processingMu.Unlock() - defer func() { - t.processingMu.Lock() - delete(t.processing, elem.ID) - t.processingMu.Unlock() - }() - return t.processElement(gctx, elem) - }) + elems := make([]*TaskElement, 0, end-i) + for _, group := range groups[i:end] { + elems = append(elems, group.elems...) + } + err = t.processElements(ctx, elems) + i = end + } + if err != nil { + break + } } - err := eg.Wait() if err != nil { logger.Errorf("Error during batch file processing: %v", err) } else { @@ -53,6 +64,159 @@ func (t *Task) Execute(ctx context.Context) error { return err } +func (t *Task) executionGroups() []executionGroup { + groups := make([]executionGroup, 0, len(t.elems)) + for i := 0; i < len(t.elems); { + elem := &t.elems[i] + batchSaver, batchCapable := elem.Storage.(storage.StorageBatchSaver) + if !batchCapable || elem.sourceGroupKey == "" { + groups = append(groups, executionGroup{elems: []*TaskElement{elem}}) + i++ + continue + } + + end := i + 1 + for end < len(t.elems) { + next := &t.elems[end] + if next.Storage != elem.Storage || next.sourceGroupKey != elem.sourceGroupKey { + break + } + end++ + } + elems := make([]*TaskElement, 0, end-i) + for j := i; j < end; j++ { + elems = append(elems, &t.elems[j]) + } + groups = append(groups, executionGroup{elems: elems, batchSaver: batchSaver}) + i = end + } + return groups +} + +func (t *Task) processElements(ctx context.Context, elems []*TaskElement) error { + eg, gctx := errgroup.WithContext(ctx) + eg.SetLimit(config.C().Workers) + for _, elem := range elems { + eg.Go(func() error { + if err := t.markProcessing(elem); err != nil { + return err + } + defer t.unmarkProcessing(elem.ID) + return t.processElement(gctx, *elem) + }) + } + return eg.Wait() +} + +func (t *Task) processBatch(ctx context.Context, group executionGroup) error { + defer func() { + for _, elem := range group.elems { + if err := os.Remove(elem.localPath); err != nil && !os.IsNotExist(err) { + log.FromContext(ctx).Warnf("Failed to cleanup batch cache file %s: %v", elem.localPath, err) + } + } + }() + + eg, gctx := errgroup.WithContext(ctx) + eg.SetLimit(config.C().Workers) + for _, elem := range group.elems { + eg.Go(func() error { + if err := t.markProcessing(elem); err != nil { + return err + } + defer t.unmarkProcessing(elem.ID) + return t.downloadElement(gctx, elem) + }) + } + if err := eg.Wait(); err != nil { + return err + } + + items := make([]storagetypes.BatchItem, 0, len(group.elems)) + openFiles := make([]*os.File, 0, len(group.elems)) + defer func() { + for _, file := range openFiles { + if err := file.Close(); err != nil { + log.FromContext(ctx).Warnf("Failed to close batch cache file %s: %v", file.Name(), err) + } + } + }() + for _, elem := range group.elems { + file, err := os.Open(elem.localPath) + if err != nil { + return fmt.Errorf("failed to open cache file: %w", err) + } + stat, err := file.Stat() + if err != nil { + file.Close() + return fmt.Errorf("failed to get cache file stat: %w", err) + } + openFiles = append(openFiles, file) + items = append(items, storagetypes.BatchItem{ + Reader: file, + StoragePath: elem.Path, + Size: stat.Size(), + SourceGroupKey: elem.sourceGroupKey, + Caption: elem.sourceCaption, + PreserveCaption: elem.preserveCaption, + }) + } + if err := group.batchSaver.SaveBatch(ctx, items); err != nil { + return fmt.Errorf("failed to save batch: %w", err) + } + return nil +} + +func (t *Task) markProcessing(elem *TaskElement) error { + t.processingMu.Lock() + defer t.processingMu.Unlock() + if t.processing[elem.ID] != nil { + return fmt.Errorf("element with ID %s is already being processed", elem.ID) + } + t.processing[elem.ID] = elem + return nil +} + +func (t *Task) unmarkProcessing(id string) { + t.processingMu.Lock() + delete(t.processing, id) + t.processingMu.Unlock() +} + +func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error { + logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name())) + logger.Info("Starting file download") + localFile, err := fsutil.CreateFile(elem.localPath) + if err != nil { + return fmt.Errorf("failed to create local file: %w", err) + } + wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) { + downloaded := t.downloaded.Add(int64(n)) + t.Progress.OnProgress(ctx, t) + taskevent.Emit(ctx, taskevent.Event{ + TaskID: t.ID, + Phase: taskevent.PhaseProgress, + TotalBytes: t.totalSize, + DownloadedBytes: downloaded, + }) + }) + _, downloadErr := tdler.NewDownloader(elem.File).Parallel(ctx, wrAt) + closeErr := localFile.Close() + if downloadErr != nil { + return fmt.Errorf("failed to download file: %w", downloadErr) + } + if closeErr != nil { + return fmt.Errorf("failed to close cache file: %w", closeErr) + } + logger.Info("File downloaded successfully") + if path.Ext(elem.FileName()) == "" { + if ext := fsutil.DetectFileExt(elem.localPath); ext != "" { + elem.Path += ext + } + } + return nil +} + func (t *Task) processElement(ctx context.Context, elem TaskElement) error { logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("file[%s]", elem.File.Name())) if elem.stream { diff --git a/core/tasks/batchtfile/execute_group_test.go b/core/tasks/batchtfile/execute_group_test.go new file mode 100644 index 0000000..eaa0f16 --- /dev/null +++ b/core/tasks/batchtfile/execute_group_test.go @@ -0,0 +1,57 @@ +package batchtfile + +import ( + "testing" + + "github.com/gotd/td/tg" + "github.com/krau/SaveAny-Bot/pkg/tfile" + tgstorage "github.com/krau/SaveAny-Bot/storage/telegram" +) + +func TestExecutionGroupsPreserveSourceAlbums(t *testing.T) { + stor := new(tgstorage.Telegram) + otherStor := new(tgstorage.Telegram) + task := Task{elems: []TaskElement{ + {Storage: stor, sourceGroupKey: "album-1"}, + {Storage: stor, sourceGroupKey: "album-1"}, + {Storage: stor}, + {Storage: stor, sourceGroupKey: "album-2"}, + {Storage: stor, sourceGroupKey: "album-2"}, + {Storage: otherStor, sourceGroupKey: "album-2"}, + }} + + groups := task.executionGroups() + wantSizes := []int{2, 1, 2, 1} + wantBatch := []bool{true, false, true, true} + if len(groups) != len(wantSizes) { + t.Fatalf("got %d groups, want %d", len(groups), len(wantSizes)) + } + for i := range groups { + if got := len(groups[i].elems); got != wantSizes[i] { + t.Errorf("group %d has %d elements, want %d", i, got, wantSizes[i]) + } + if got := groups[i].usesBatchSaver(); got != wantBatch[i] { + t.Errorf("group %d batch=%v, want %v", i, got, wantBatch[i]) + } + } +} + +func TestSourceMetadataPreservesAlbumIdentityAndCaption(t *testing.T) { + msg := &tg.Message{ + PeerID: &tg.PeerChannel{ChannelID: 77}, + Message: "original caption", + } + msg.SetGroupedID(42) + file := tfile.NewTGFile(nil, nil, 0, "photo.jpg", tfile.WithMessage(msg)) + + groupKey, caption, preserveCaption := sourceMetadata(file) + if groupKey != "*tg.PeerChannel:77:42" { + t.Fatalf("group key = %q, want %q", groupKey, "*tg.PeerChannel:77:42") + } + if caption != "original caption" { + t.Fatalf("caption = %q, want original caption", caption) + } + if !preserveCaption { + t.Fatal("preserveCaption = false, want true") + } +} diff --git a/core/tasks/batchtfile/task.go b/core/tasks/batchtfile/task.go index c00d510..3737983 100644 --- a/core/tasks/batchtfile/task.go +++ b/core/tasks/batchtfile/task.go @@ -7,6 +7,7 @@ import ( "sync" "sync/atomic" + "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/pkg/enums/tasktype" @@ -18,12 +19,15 @@ import ( var _ core.Executable = (*Task)(nil) type TaskElement struct { - ID string - Storage storage.Storage - Path string - File tfile.TGFile - localPath string - stream bool + ID string + Storage storage.Storage + Path string + File tfile.TGFile + localPath string + stream bool + sourceGroupKey string + sourceCaption string + preserveCaption bool } type Task struct { @@ -54,6 +58,7 @@ func NewTaskElement( file tfile.TGFile, ) (*TaskElement, error) { id := xid.New().String() + groupKey, caption, preserveCaption := sourceMetadata(file) _, ok := stor.(storage.StorageCannotStream) if !config.C().Stream || ok { cachePath, err := filepath.Abs(filepath.Join(config.C().Temp.BasePath, fmt.Sprintf("%s_%s", id, file.Name()))) @@ -61,22 +66,42 @@ func NewTaskElement( return nil, fmt.Errorf("failed to get absolute path for cache: %w", err) } return &TaskElement{ - ID: id, - Storage: stor, - Path: path, - File: file, - localPath: cachePath, + ID: id, + Storage: stor, + Path: path, + File: file, + localPath: cachePath, + sourceGroupKey: groupKey, + sourceCaption: caption, + preserveCaption: preserveCaption, }, nil } return &TaskElement{ - ID: id, - Storage: stor, - Path: path, - File: file, - stream: true, + ID: id, + Storage: stor, + Path: path, + File: file, + stream: true, + sourceGroupKey: groupKey, + sourceCaption: caption, + preserveCaption: preserveCaption, }, nil } +func sourceMetadata(file tfile.TGFile) (groupKey, caption string, preserveCaption bool) { + messageFile, ok := file.(tfile.TGFileMessage) + if !ok || messageFile.Message() == nil { + return "", "", false + } + msg := messageFile.Message() + groupID, grouped := msg.GetGroupedID() + if !grouped || groupID == 0 { + return "", "", false + } + chatID := tgutil.ChatIdFromPeer(msg.GetPeerID()) + return fmt.Sprintf("%T:%d:%d", msg.GetPeerID(), chatID, groupID), msg.GetMessage(), true +} + func NewBatchTGFileTask( id string, ctx context.Context, diff --git a/pkg/storagetypes/batch.go b/pkg/storagetypes/batch.go new file mode 100644 index 0000000..795bd0c --- /dev/null +++ b/pkg/storagetypes/batch.go @@ -0,0 +1,17 @@ +package storagetypes + +import "io" + +// BatchItem describes one seekable file in a logical batch storage operation. +type BatchItem struct { + Reader io.ReadSeeker + StoragePath string + Size int64 + + // SourceGroupKey is empty for standalone source messages. + SourceGroupKey string + Caption string + // PreserveCaption distinguishes an intentionally empty source caption from + // the storage backend's default caption. + PreserveCaption bool +} diff --git a/storage/storage.go b/storage/storage.go index 1d14cc1..285b1ce 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -31,6 +31,13 @@ type StorageCannotStream interface { CannotStream() string } +// StorageBatchSaver can preserve relationships between files when saving a +// logical batch, such as a Telegram media album. +type StorageBatchSaver interface { + Storage + SaveBatch(ctx context.Context, items []storagetypes.BatchItem) error +} + // StorageListable 表示支持列举目录内容的存储 type StorageListable interface { Storage diff --git a/storage/telegram/media_group_test.go b/storage/telegram/media_group_test.go new file mode 100644 index 0000000..4b2d394 --- /dev/null +++ b/storage/telegram/media_group_test.go @@ -0,0 +1,140 @@ +package telegram + +import ( + "bytes" + "io" + "testing" + + "github.com/krau/SaveAny-Bot/pkg/storagetypes" +) + +func TestPlanMediaGroups(t *testing.T) { + tests := []struct { + name string + items []batchMediaItem + wantSizes []int + }{ + { + name: "same source album", + items: []batchMediaItem{ + albumItem("a", 1, true), + albumItem("a", 1, true), + }, + wantSizes: []int{2}, + }, + { + name: "different source albums", + items: []batchMediaItem{ + albumItem("a", 1, true), + albumItem("b", 1, true), + }, + wantSizes: []int{1, 1}, + }, + { + name: "ungrouped messages", + items: []batchMediaItem{ + albumItem("", 1, true), + albumItem("", 1, true), + }, + wantSizes: []int{1, 1}, + }, + { + name: "different target chats", + items: []batchMediaItem{ + albumItem("a", 1, true), + albumItem("a", 2, true), + }, + wantSizes: []int{1, 1}, + }, + { + name: "ineligible media does not bridge albums", + items: []batchMediaItem{ + albumItem("a", 1, true), + albumItem("a", 1, false), + albumItem("a", 1, true), + }, + wantSizes: []int{1, 1, 1}, + }, + { + name: "maximum album size", + items: repeatedAlbumItems(11), + wantSizes: []int{10, 1}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + groups := planMediaGroups(tt.items) + if len(groups) != len(tt.wantSizes) { + t.Fatalf("got %d groups, want %d", len(groups), len(tt.wantSizes)) + } + for i, want := range tt.wantSizes { + if got := len(groups[i]); got != want { + t.Errorf("group %d has %d items, want %d", i, got, want) + } + } + }) + } +} + +func TestMediaCaption(t *testing.T) { + empty := "" + original := "original caption" + tests := []struct { + name string + override *string + wantLen int + }{ + {name: "filename fallback", wantLen: 1}, + {name: "preserve empty source caption", override: &empty, wantLen: 0}, + {name: "preserve source caption", override: &original, wantLen: 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := len(mediaCaption("file.jpg", tt.override)); got != tt.wantLen { + t.Fatalf("got %d caption options, want %d", got, tt.wantLen) + } + }) + } +} + +func TestInspectBatchItemRewindsBeforeMimetypeDetection(t *testing.T) { + data := []byte("\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR") + reader := bytes.NewReader(data) + if _, err := reader.Seek(4, io.SeekStart); err != nil { + t.Fatalf("failed to set initial reader offset: %v", err) + } + + mediaItem, err := new(Telegram).inspectBatchItem(nil, storagetypes.BatchItem{ + Reader: reader, + StoragePath: "photo.png", + Size: int64(len(data)), + }) + if err != nil { + t.Fatalf("inspectBatchItem returned an error: %v", err) + } + if !mediaItem.albumEligible { + t.Fatal("albumEligible = false, want true for PNG input") + } + if offset, err := reader.Seek(0, io.SeekCurrent); err != nil { + t.Fatalf("failed to get final reader offset: %v", err) + } else if offset != 0 { + t.Fatalf("reader offset = %d, want 0", offset) + } +} + +func albumItem(group string, chatID int64, eligible bool) batchMediaItem { + return batchMediaItem{ + item: storagetypes.BatchItem{SourceGroupKey: group}, + chatID: chatID, + albumEligible: eligible, + } +} + +func repeatedAlbumItems(count int) []batchMediaItem { + items := make([]batchMediaItem, count) + for i := range items { + items[i] = albumItem("a", 1, true) + } + return items +} diff --git a/storage/telegram/telegram.go b/storage/telegram/telegram.go index 6e94c86..6938e87 100644 --- a/storage/telegram/telegram.go +++ b/storage/telegram/telegram.go @@ -12,6 +12,7 @@ import ( "github.com/celestix/gotgproto/ext" "github.com/charmbracelet/log" + "github.com/duke-git/lancet/v2/retry" "github.com/duke-git/lancet/v2/slice" "github.com/duke-git/lancet/v2/validator" "github.com/gabriel-vasile/mimetype" @@ -26,6 +27,7 @@ import ( "github.com/krau/SaveAny-Bot/pkg/consts/tglimit" "github.com/krau/SaveAny-Bot/pkg/enums/ctxkey" storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage" + "github.com/krau/SaveAny-Bot/pkg/storagetypes" "github.com/rs/xid" "golang.org/x/time/rate" ) @@ -41,6 +43,19 @@ type Telegram struct { limiter *rate.Limiter } +type preparedMedia struct { + peer tg.InputPeerClass + uploader *uploader.Uploader + media message.MultiMediaOption +} + +type batchMediaItem struct { + item storagetypes.BatchItem + chatID int64 + albumEligible bool + useSingleSave bool +} + func (t *Telegram) Init(ctx context.Context, cfg storconfig.StorageConfig) error { telegramConfig, ok := cfg.(*storconfig.TelegramStorageConfig) if !ok { @@ -71,37 +86,76 @@ func (t *Telegram) Exists(ctx context.Context, storagePath string) bool { } func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) error { - storagePath = path.Clean(storagePath) tctx := tgutil.ExtFromContext(ctx) if tctx == nil { return fmt.Errorf("failed to get telegram context") } - size := func() int64 { - if length := ctx.Value(ctxkey.ContentLength); length != nil { - if l, ok := length.(int64); ok { - return l - } - } - return -1 // unknown size - }() + size := contentLength(ctx) if t.config.SkipLarge && size > MaxUploadFileSize { log.FromContext(ctx).Warnf("Skipping file larger than Telegram limit (%d bytes): %d bytes", MaxUploadFileSize, size) return nil } - rs, seekable := r.(io.ReadSeeker) - splitSize := t.config.SplitSizeMB * 1024 * 1024 - if splitSize <= 0 { - splitSize = DefaultSplitSize + if size > t.splitSize() { + filename, chatID := t.target(tctx, path.Clean(storagePath)) + if filename == "" { + if rs, ok := r.(io.ReadSeeker); ok { + mtype, err := mimetype.DetectReader(rs) + if err != nil { + return fmt.Errorf("failed to detect mimetype: %w", err) + } + filename = xid.New().String() + mtype.Extension() + if _, err := rs.Seek(0, io.SeekStart); err != nil { + return fmt.Errorf("failed to seek reader: %w", err) + } + } + } + upler := t.newUploader(tctx, size) + peer := tryGetInputPeer(tctx, chatID) + if peer == nil || peer.Zero() { + return fmt.Errorf("failed to get input peer for chat ID %d", chatID) + } + if err := t.limiter.Wait(ctx); err != nil { + return fmt.Errorf("rate limit failed: %w", err) + } + return t.splitUpload(tctx, r, filename, upler, peer, size, t.splitSize()) } if err := t.limiter.Wait(ctx); err != nil { return fmt.Errorf("rate limit failed: %w", err) } + prepared, err := t.prepareMedia(ctx, tctx, r, storagePath, size, nil) + if err != nil { + return err + } + _, err = tctx.Sender. + WithUploader(prepared.uploader). + To(prepared.peer). + Media(ctx, prepared.media) + return err +} +func contentLength(ctx context.Context) int64 { + if length := ctx.Value(ctxkey.ContentLength); length != nil { + if size, ok := length.(int64); ok { + return size + } + } + return -1 +} + +func (t *Telegram) splitSize() int64 { + splitSize := t.config.SplitSizeMB * 1024 * 1024 + if splitSize <= 0 { + return DefaultSplitSize + } + return splitSize +} + +func (t *Telegram) target(tctx *ext.Context, storagePath string) (string, int64) { // 去除前导斜杠并分隔路径, 当 len(parts): // ==0, 存储到配置文件中的 chat_id, 随机文件名 // ==1, 视作只有文件名, 存储到配置文件中的 chat_id - // ==2, parts[0]: 视作要存储到的 chat_id, parts[1]: filename + // >=2, parts[0]: 视作要存储到的 chat_id, 最后一项为 filename parts := slice.Compact(strings.Split(strings.TrimPrefix(storagePath, "/"), "/")) filename := "" chatID := t.config.ChatID @@ -111,38 +165,54 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er if len(parts) >= 2 && validator.IsAlphaNumeric(parts[0]) { cid, err := tgutil.ParseChatID(tctx, parts[0]) if err != nil { - // id不合法时使用配置文件中的 chat_id - log.FromContext(ctx).Warnf("Failed to parse chat ID from path, using configured chat_id: %s", err) + log.FromContext(tctx).Warnf("Failed to parse chat ID from path, using configured chat_id: %s", err) cid = chatID } chatID = cid } - upler := uploader.NewUploader(tctx.Raw). + return filename, chatID +} + +func (t *Telegram) newUploader(tctx *ext.Context, size int64) *uploader.Uploader { + return uploader.NewUploader(tctx.Raw). WithPartSize(tglimit.MaxUploadPartSize). WithThreads(dlutil.BestThreads(size, config.C().Threads)) +} + +func mediaCaption(filename string, override *string) []message.StyledTextOption { + if override == nil { + return []message.StyledTextOption{styling.Plain(filename)} + } + if *override == "" { + return nil + } + return []message.StyledTextOption{styling.Plain(*override)} +} + +func (t *Telegram) prepareMedia(ctx context.Context, tctx *ext.Context, r io.Reader, storagePath string, size int64, captionOverride *string) (*preparedMedia, error) { + storagePath = path.Clean(storagePath) + filename, chatID := t.target(tctx, storagePath) + upler := t.newUploader(tctx, size) peer := tryGetInputPeer(tctx, chatID) if peer == nil || peer.Zero() { - return fmt.Errorf("failed to get input peer for chat ID %d", chatID) + return nil, fmt.Errorf("failed to get input peer for chat ID %d", chatID) } + + rs, seekable := r.(io.ReadSeeker) var mtype *mimetype.MIME if seekable { var err error mtype, err = mimetype.DetectReader(rs) if err != nil { - return fmt.Errorf("failed to detect mimetype: %w", err) + return nil, fmt.Errorf("failed to detect mimetype: %w", err) } if filename == "" { filename = xid.New().String() + mtype.Extension() } - if _, err := rs.Seek(0, io.SeekStart); err != nil { - return fmt.Errorf("failed to seek reader: %w", err) + return nil, fmt.Errorf("failed to seek reader: %w", err) } } - if size > splitSize { - // large file, use split uploader - return t.splitUpload(tctx, r, filename, upler, peer, size, splitSize) - } var file tg.InputFileClass var err error @@ -152,21 +222,20 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er file, err = upler.Upload(ctx, uploader.NewUpload(filename, r, size)) } if err != nil { - return fmt.Errorf("failed to upload file to telegram: %w", err) + return nil, fmt.Errorf("failed to upload file to telegram: %w", err) } - caption := styling.Plain(filename) + caption := mediaCaption(filename, captionOverride) forceFile := t.config.ForceFile - if mtype != nil && strings.HasPrefix(mtype.String(), "image/") && size >= tglimit.MaxPhotoSize { forceFile = true } - doc := message.UploadedDocument(file, caption). + doc := message.UploadedDocument(file, caption...). Filename(filename). ForceFile(forceFile) if mtype != nil { doc = doc.MIME(mtype.String()) } - var media message.MediaOption = doc + var media message.MultiMediaOption = doc if mtype != nil && rs != nil { switch mtypeStr := mtype.String(); { case strings.HasPrefix(mtypeStr, "video/"): @@ -205,12 +274,131 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er case strings.HasPrefix(mtypeStr, "audio/"): media = doc.Audio().Title(filename) case strings.HasPrefix(mtypeStr, "image/") && !strings.HasSuffix(mtypeStr, "webp"): - media = message.UploadedPhoto(file, caption) + media = message.UploadedPhoto(file, caption...) } } - sender := tctx.Sender - _, err = sender.WithUploader(upler).To(peer).Media(ctx, media) - return err + return &preparedMedia{ + peer: peer, + uploader: upler, + media: media, + }, nil +} + +// SaveBatch preserves each source photo/video group as a Telegram album. +func (t *Telegram) SaveBatch(ctx context.Context, items []storagetypes.BatchItem) error { + tctx := tgutil.ExtFromContext(ctx) + if tctx == nil { + return fmt.Errorf("failed to get telegram context") + } + + inspected := make([]batchMediaItem, 0, len(items)) + for _, item := range items { + mediaItem, err := t.inspectBatchItem(tctx, item) + if err != nil { + return err + } + inspected = append(inspected, mediaItem) + } + for _, group := range planMediaGroups(inspected) { + if err := t.saveMediaGroup(ctx, tctx, group); err != nil { + return err + } + } + return nil +} + +func (t *Telegram) inspectBatchItem(tctx *ext.Context, item storagetypes.BatchItem) (batchMediaItem, error) { + _, chatID := t.target(tctx, path.Clean(item.StoragePath)) + result := batchMediaItem{item: item, chatID: chatID} + if (t.config.SkipLarge && item.Size > MaxUploadFileSize) || item.Size > t.splitSize() { + result.useSingleSave = true + return result, nil + } + if _, err := item.Reader.Seek(0, io.SeekStart); err != nil { + return result, fmt.Errorf("failed to seek batch item before mimetype detection: %w", err) + } + mtype, err := mimetype.DetectReader(item.Reader) + if err != nil { + return result, fmt.Errorf("failed to detect batch item mimetype: %w", err) + } + if _, err := item.Reader.Seek(0, io.SeekStart); err != nil { + return result, fmt.Errorf("failed to seek batch item: %w", err) + } + mtypeStr := mtype.String() + forceFile := t.config.ForceFile || strings.HasPrefix(mtypeStr, "image/") && item.Size >= tglimit.MaxPhotoSize + result.albumEligible = !forceFile && (strings.HasPrefix(mtypeStr, "video/") || + strings.HasPrefix(mtypeStr, "image/") && mtypeStr != "image/webp" && mtypeStr != "image/gif") + return result, nil +} + +func planMediaGroups(items []batchMediaItem) [][]batchMediaItem { + groups := make([][]batchMediaItem, 0, len(items)) + for i := 0; i < len(items); { + item := items[i] + if item.useSingleSave || !item.albumEligible || item.item.SourceGroupKey == "" { + groups = append(groups, items[i:i+1]) + i++ + continue + } + end := i + 1 + for end < len(items) && end-i < 10 { + next := items[end] + if next.useSingleSave || !next.albumEligible || next.chatID != item.chatID || next.item.SourceGroupKey != item.item.SourceGroupKey { + break + } + end++ + } + groups = append(groups, items[i:end]) + i = end + } + return groups +} + +func (t *Telegram) saveMediaGroup(ctx context.Context, tctx *ext.Context, group []batchMediaItem) error { + return retry.Retry(func() error { + if len(group) == 1 && group[0].useSingleSave { + item := group[0].item + if _, err := item.Reader.Seek(0, io.SeekStart); err != nil { + return fmt.Errorf("failed to seek batch item: %w", err) + } + itemCtx := context.WithValue(ctx, ctxkey.ContentLength, item.Size) + return t.Save(itemCtx, item.Reader, item.StoragePath) + } + if err := t.limiter.Wait(ctx); err != nil { + return fmt.Errorf("rate limit failed: %w", err) + } + + prepared := make([]preparedMedia, 0, len(group)) + for _, mediaItem := range group { + item := mediaItem.item + if _, err := item.Reader.Seek(0, io.SeekStart); err != nil { + return fmt.Errorf("failed to seek batch item: %w", err) + } + var captionOverride *string + if item.PreserveCaption { + captionOverride = &item.Caption + } + media, err := t.prepareMedia(ctx, tctx, item.Reader, item.StoragePath, item.Size, captionOverride) + if err != nil { + return err + } + prepared = append(prepared, *media) + } + + builder := tctx.Sender.WithUploader(prepared[0].uploader).To(prepared[0].peer) + if len(prepared) == 1 { + _, err := builder.Media(ctx, prepared[0].media) + return err + } + media := make([]message.MultiMediaOption, len(prepared)) + for i := range prepared { + media[i] = prepared[i].media + } + if _, err := builder.Album(ctx, media[0], media[1:]...); err != nil { + return fmt.Errorf("failed to send media album: %w", err) + } + return nil + }, retry.Context(ctx), retry.RetryTimes(uint(config.C().Retry))) } func (t *Telegram) CannotStream() string {