Fix performance issues, add media group support, and improve filename handling

Co-authored-by: krau <71133316+krau@users.noreply.github.com>

api/handlers.go

@@ -4,18 +4,25 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"io"
+	"math"
 	"net/http"
+	"net/url"
 	"path"
 	"sync"
+	"sync/atomic"
 	"time"
 
+	"github.com/celestix/gotgproto/ext"
 	"github.com/charmbracelet/log"
+	"github.com/gotd/td/tg"
 	"github.com/krau/SaveAny-Bot/client/bot"
 	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/core"
 	tftask "github.com/krau/SaveAny-Bot/core/tasks/tfile"
+	"github.com/krau/SaveAny-Bot/database"
 	"github.com/krau/SaveAny-Bot/pkg/queue"
 	"github.com/krau/SaveAny-Bot/pkg/tfile"
 	"github.com/krau/SaveAny-Bot/storage"
@@ -72,9 +79,9 @@ type taskStatus struct {
 	Title       string
 	CreatedAt   time.Time
 	Error       string
-	Downloaded  int64
+	Downloaded  atomic.Int64 // Use atomic for lock-free updates
-	Total       int64
+	Total       atomic.Int64 // Use atomic for lock-free updates
-	ProgressPct float64
+	ProgressPct uint64       // Store as uint64 bits of float64 for atomic access
 }
 
 func handleHealth(w http.ResponseWriter, r *http.Request) {
@@ -102,9 +109,16 @@ func handleCreateTask(w http.ResponseWriter, r *http.Request) {
 
 	logger := log.FromContext(r.Context()).WithPrefix("api")
 
+	// Get user from database
+	userDB, err := database.GetUserByChatID(r.Context(), req.UserID)
+	if err != nil {
+		logger.Errorf("Failed to get user: %v", err)
+		respondError(w, "user not found", http.StatusBadRequest)
+		return
+	}
+
 	// Get storage
 	var stor storage.Storage
-	var err error
 	if req.StorageName != "" {
 		stor, err = storage.GetStorageByUserIDAndName(r.Context(), req.UserID, req.StorageName)
 		if err != nil {
@@ -129,6 +143,13 @@ func handleCreateTask(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	linkUrl, err := url.Parse(req.TelegramURL)
+	if err != nil {
+		logger.Errorf("Failed to parse URL: %v", err)
+		respondError(w, "invalid telegram URL format", http.StatusBadRequest)
+		return
+	}
+
 	chatID, msgID, err := tgutil.ParseMessageLink(botCtx, req.TelegramURL)
 	if err != nil {
 		logger.Errorf("Failed to parse Telegram URL: %v", err)
@@ -151,26 +172,72 @@ func handleCreateTask(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Create TGFile from message media
+	// Collect files - handle both single and grouped messages
-	tgFile, err := tfile.FromMediaMessage(media, botCtx.Raw, msg)
+	files := make([]tfile.TGFileMessage, 0)
+
+	// Check for grouped messages (media group)
+	groupID, isGroup := msg.GetGroupedID()
+	if isGroup && groupID != 0 && !linkUrl.Query().Has("single") {
+		// Handle media group
+		gmsgs, err := tgutil.GetGroupedMessages(botCtx, chatID, msg)
+		if err != nil {
+			logger.Errorf("Failed to get grouped messages: %v", err)
+			// Fall back to single message
+			file, err := createTGFileWithMedia(botCtx, msg, media, userDB)
 			if err != nil {
 				logger.Errorf("Failed to create TGFile: %v", err)
 				respondError(w, "invalid message format", http.StatusBadRequest)
 				return
 			}
+			files = append(files, file)
+		} else {
+			// Process all messages in the group
+			for _, gmsg := range gmsgs {
+				if gmsg.Media == nil {
+					continue
+				}
+				gMedia, ok := gmsg.GetMedia()
+				if !ok {
+					continue
+				}
+				file, err := createTGFileWithMedia(botCtx, gmsg, gMedia, userDB)
+				if err != nil {
+					logger.Warnf("Failed to create TGFile for grouped message: %v", err)
+					continue
+				}
+				files = append(files, file)
+			}
+		}
+	} else {
+		// Single message
+		file, err := createTGFileWithMedia(botCtx, msg, media, userDB)
+		if err != nil {
+			logger.Errorf("Failed to create TGFile: %v", err)
+			respondError(w, "invalid message format", http.StatusBadRequest)
+			return
+		}
+		files = append(files, file)
+	}
 
-	// Create task
+	if len(files) == 0 {
+		respondError(w, "no savable files found", http.StatusBadRequest)
+		return
+	}
+
+	// Create tasks for all files
+	taskIDs := make([]string, 0, len(files))
 	dirPath := req.DirPath
 	if dirPath == "" {
 		dirPath = "/"
 	}
 
-	storagePath := stor.JoinStoragePath(path.Join(dirPath, tgFile.Name()))
-	taskID := xid.New().String()
-
 	// Create context with bot extension
 	injectCtx := tgutil.ExtWithContext(r.Context(), botCtx)
 
+	for _, tgFile := range files {
+		storagePath := stor.JoinStoragePath(path.Join(dirPath, tgFile.Name()))
+		taskID := xid.New().String()
+
 		task, err := tftask.NewTGFileTask(taskID, injectCtx, tgFile, stor, storagePath, &apiProgressTracker{
 			taskID: taskID,
 		})
@@ -191,13 +258,32 @@ func handleCreateTask(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
+		taskIDs = append(taskIDs, taskID)
+	}
+
 	// Send success response
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(http.StatusCreated)
+
+	// Return first task ID for single file, or all task IDs for media group
+	if len(taskIDs) == 1 {
 		json.NewEncoder(w).Encode(CreateTaskResponse{
-		TaskID:  taskID,
+			TaskID:  taskIDs[0],
 			Message: "task created successfully",
 		})
+	} else {
+		json.NewEncoder(w).Encode(map[string]interface{}{
+			"task_ids": taskIDs,
+			"message":  fmt.Sprintf("%d tasks created successfully", len(taskIDs)),
+		})
+	}
+}
+
+// createTGFileWithMedia creates a TGFile with proper filename handling using user's strategy
+func createTGFileWithMedia(botCtx *ext.Context, msg *tg.Message, media tg.MessageMediaClass, userDB *database.User) (tfile.TGFileMessage, error) {
+	// Use the same filename generation logic as bot handlers
+	opts := []tfile.TGFileOption{tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*msg))}
+	return tfile.FromMediaMessage(media, botCtx.Raw, msg, opts...)
 }
 
 func handleGetTask(w http.ResponseWriter, r *http.Request) {
@@ -223,9 +309,9 @@ func handleGetTask(w http.ResponseWriter, r *http.Request) {
 		Title:       status.Title,
 		CreatedAt:   status.CreatedAt,
 		Error:       status.Error,
-		Downloaded:  status.Downloaded,
+		Downloaded:  status.Downloaded.Load(),
-		Total:       status.Total,
+		Total:       status.Total.Load(),
-		ProgressPct: status.ProgressPct,
+		ProgressPct: math.Float64frombits(atomic.LoadUint64((*uint64)(&status.ProgressPct))),
 	})
 }
 
@@ -310,13 +396,18 @@ func (a *apiProgressTracker) OnStart(ctx context.Context, info tftask.TaskInfo)
 }
 
 func (a *apiProgressTracker) OnProgress(ctx context.Context, info tftask.TaskInfo, downloaded int64, total int64) {
-	taskStatusesMu.Lock()
+	// Use atomic operations to avoid mutex locks for better performance
-	defer taskStatusesMu.Unlock()
+	// OnProgress is called very frequently during downloads
-	if ts, exists := taskStatuses[a.taskID]; exists {
+	taskStatusesMu.RLock()
-		ts.Downloaded = downloaded
+	ts, exists := taskStatuses[a.taskID]
-		ts.Total = total
+	taskStatusesMu.RUnlock()
+	
+	if exists {
+		ts.Downloaded.Store(downloaded)
+		ts.Total.Store(total)
 		if total > 0 {
-			ts.ProgressPct = float64(downloaded) / float64(total) * 100.0
+			progressPct := float64(downloaded) / float64(total) * 100.0
+			atomic.StoreUint64((*uint64)(&ts.ProgressPct), math.Float64bits(progressPct))
 		}
 	}
 }
@@ -354,9 +445,9 @@ func sendWebhook(taskID, status, errorMsg string) {
 		Title:       ts.Title,
 		CreatedAt:   ts.CreatedAt,
 		Error:       errorMsg,
-		Downloaded:  ts.Downloaded,
+		Downloaded:  ts.Downloaded.Load(),
-		Total:       ts.Total,
+		Total:       ts.Total.Load(),
-		ProgressPct: ts.ProgressPct,
+		ProgressPct: math.Float64frombits(atomic.LoadUint64((*uint64)(&ts.ProgressPct))),
 	}
 
 	body, err := json.Marshal(payload)