feat: add yt-dlp support for downloading video/audio and enhance related commands

2026-06-05 07:30:32 +08:00 · 2026-01-17 17:42:11 +08:00
parent cd7cf4964d
commit 3ce00884a0
16 changed files with 602 additions and 4 deletions
--- a/core/tasks/ytdlp/progress.go
+++ b/core/tasks/ytdlp/progress.go
@@ -0,0 +1,183 @@
+package ytdlp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync/atomic"
+	"time"
+
+	"github.com/charmbracelet/log"
+	"github.com/gotd/td/telegram/message/entity"
+	"github.com/gotd/td/telegram/message/styling"
+	"github.com/gotd/td/tg"
+
+	"github.com/krau/SaveAny-Bot/common/i18n"
+	"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
+	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
+)
+
+// ProgressTracker defines the interface for tracking ytdlp task progress
+type ProgressTracker interface {
+	OnStart(ctx context.Context, task *Task)
+	OnProgress(ctx context.Context, task *Task, status string)
+	OnDone(ctx context.Context, task *Task, err error)
+}
+
+type Progress struct {
+	msgID             int
+	chatID            int64
+	start             time.Time
+	lastUpdate        atomic.Value // stores time.Time
+	minUpdateInterval time.Duration
+}
+
+// OnStart implements ProgressTracker.
+func (p *Progress) OnStart(ctx context.Context, task *Task) {
+	logger := log.FromContext(ctx)
+	p.start = time.Now()
+	p.lastUpdate.Store(time.Now())
+	p.minUpdateInterval = 2 * time.Second // Avoid too frequent updates
+	logger.Infof("yt-dlp task started: message_id=%d, chat_id=%d, urls=%d", p.msgID, p.chatID, len(task.URLs))
+	ext := tgutil.ExtFromContext(ctx)
+	if ext == nil {
+		return
+	}
+	entityBuilder := entity.Builder{}
+	if err := styling.Perform(&entityBuilder,
+		styling.Plain(i18n.T(i18nk.BotMsgProgressYtdlpStart, map[string]any{
+			"Count": len(task.URLs),
+		})),
+		styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)),
+		styling.Code(fmt.Sprintf("[%s]:%s", task.Storage.Name(), task.StorPath)),
+	); err != nil {
+		log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
+		return
+	}
+	text, entities := entityBuilder.Complete()
+	req := &tg.MessagesEditMessageRequest{
+		ID: p.msgID,
+	}
+	req.SetMessage(text)
+	req.SetEntities(entities)
+	req.SetReplyMarkup(&tg.ReplyInlineMarkup{
+		Rows: []tg.KeyboardButtonRow{
+			{
+				Buttons: []tg.KeyboardButtonClass{
+					tgutil.BuildCancelButton(task.TaskID()),
+				},
+			},
+		}},
+	)
+	ext.EditMessage(p.chatID, req)
+}
+
+// OnProgress implements ProgressTracker.
+func (p *Progress) OnProgress(ctx context.Context, task *Task, status string) {
+	// Throttle updates to avoid flooding Telegram API
+	lastUpdateTime := p.lastUpdate.Load().(time.Time)
+	if time.Since(lastUpdateTime) < p.minUpdateInterval {
+		return
+	}
+	p.lastUpdate.Store(time.Now())
+
+	log.FromContext(ctx).Debugf("yt-dlp progress update: %s", status)
+
+	entityBuilder := entity.Builder{}
+	if err := styling.Perform(&entityBuilder,
+		styling.Plain(i18n.T(i18nk.BotMsgProgressYtdlpDownloading, map[string]any{
+			"Count": len(task.URLs),
+		})),
+		styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)),
+		styling.Code(fmt.Sprintf("[%s]:%s", task.Storage.Name(), task.StorPath)),
+		styling.Plain("\n\n"),
+		styling.Plain(status),
+	); err != nil {
+		log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
+		return
+	}
+	text, entities := entityBuilder.Complete()
+	req := &tg.MessagesEditMessageRequest{
+		ID: p.msgID,
+	}
+	req.SetMessage(text)
+	req.SetEntities(entities)
+	req.SetReplyMarkup(&tg.ReplyInlineMarkup{
+		Rows: []tg.KeyboardButtonRow{
+			{
+				Buttons: []tg.KeyboardButtonClass{
+					tgutil.BuildCancelButton(task.TaskID()),
+				},
+			},
+		}},
+	)
+	ext := tgutil.ExtFromContext(ctx)
+	if ext != nil {
+		ext.EditMessage(p.chatID, req)
+	}
+}
+
+// OnDone implements ProgressTracker.
+func (p *Progress) OnDone(ctx context.Context, task *Task, err error) {
+	logger := log.FromContext(ctx)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			logger.Infof("yt-dlp task %s was canceled", task.TaskID())
+			ext := tgutil.ExtFromContext(ctx)
+			if ext != nil {
+				ext.EditMessage(p.chatID, &tg.MessagesEditMessageRequest{
+					ID: p.msgID,
+					Message: i18n.T(i18nk.BotMsgProgressTaskCanceledWithId, map[string]any{
+						"TaskID": task.TaskID(),
+					}),
+				})
+			}
+		} else {
+			logger.Errorf("yt-dlp task %s failed: %s", task.TaskID(), err)
+			ext := tgutil.ExtFromContext(ctx)
+			if ext != nil {
+				ext.EditMessage(p.chatID, &tg.MessagesEditMessageRequest{
+					ID: p.msgID,
+					Message: i18n.T(i18nk.BotMsgProgressTaskFailedWithError, map[string]any{
+						"Error": err.Error(),
+					}),
+				})
+			}
+		}
+		return
+	}
+	logger.Infof("yt-dlp task %s completed successfully", task.TaskID())
+
+	entityBuilder := entity.Builder{}
+	if err := styling.Perform(&entityBuilder,
+		styling.Plain(i18n.T(i18nk.BotMsgProgressYtdlpDone, map[string]any{
+			"Count": len(task.URLs),
+		})),
+		styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)),
+		styling.Code(fmt.Sprintf("[%s]:%s", task.Storage.Name(), task.StorPath)),
+	); err != nil {
+		logger.Errorf("Failed to build entities: %s", err)
+		return
+	}
+	text, entities := entityBuilder.Complete()
+	req := &tg.MessagesEditMessageRequest{
+		ID: p.msgID,
+	}
+	req.SetMessage(text)
+	req.SetEntities(entities)
+
+	ext := tgutil.ExtFromContext(ctx)
+	if ext != nil {
+		ext.EditMessage(p.chatID, req)
+	}
+}
+
+var _ ProgressTracker = (*Progress)(nil)
+
+func NewProgress(msgID int, userID int64) ProgressTracker {
+	return &Progress{
+		msgID:             msgID,
+		chatID:            userID,
+		minUpdateInterval: 2 * time.Second,
+	}
+}