SaveAny-Bot/core/tasks/ytdlp/progress.go

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,
	}
}