SaveAny-Bot/core/tasks/aria2dl/execute.go

package aria2dl

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/charmbracelet/log"
	"github.com/krau/SaveAny-Bot/config"
	"github.com/krau/SaveAny-Bot/pkg/aria2"
	"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
)

// Execute implements core.Executable.
func (t *Task) Execute(ctx context.Context) error {
	logger := log.FromContext(ctx)
	logger.Infof("Starting aria2 download task %s (GID: %s)", t.ID, t.GID)

	if t.Progress != nil {
		t.Progress.OnStart(ctx, t)
	}

	// Wait for aria2 download to complete
	if err := t.waitForDownload(ctx); err != nil {
		// If context was canceled, also cancel the aria2 download
		if errors.Is(err, context.Canceled) {
			t.cancelAria2Download()
		}
		logger.Errorf("Aria2 download failed: %v", err)
		if t.Progress != nil {
			t.Progress.OnDone(ctx, t, err)
		}
		return err
	}

	// Transfer downloaded files to storage
	if err := t.transferFiles(ctx); err != nil {
		logger.Errorf("File transfer failed: %v", err)
		if t.Progress != nil {
			t.Progress.OnDone(ctx, t, err)
		}
		return err
	}

	logger.Infof("Aria2 task %s completed successfully", t.ID)
	if t.Progress != nil {
		t.Progress.OnDone(ctx, t, nil)
	}

	// Clean up aria2 download result
	if _, err := t.Aria2Client.RemoveDownloadResult(context.Background(), t.GID); err != nil {
		logger.Warnf("Failed to remove aria2 download result: %v", err)
	}

	return nil
}

// waitForDownload waits for aria2 to complete the download
func (t *Task) waitForDownload(ctx context.Context) error {
	logger := log.FromContext(ctx)
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			status, err := t.getStatus(ctx)
			if err != nil {
				return err
			}

			if t.Progress != nil {
				t.Progress.OnProgress(ctx, t, status)
			}

			// Check if download is complete
			if status.IsDownloadComplete() {
				// Handle metadata downloads (torrent/magnet) that spawn follow-up downloads
				if len(status.FollowedBy) > 0 {
					logger.Infof("Switching from metadata GID %s to actual download GID: %s", t.GID, status.FollowedBy[0])
					t.GID = status.FollowedBy[0]
					continue
				}
				logger.Infof("Download completed for GID %s", t.GID)
				return nil
			}

			// Check for errors
			if status.IsDownloadError() {
				return fmt.Errorf("aria2 download error: %s (code: %s)", status.ErrorMessage, status.ErrorCode)
			}

			if status.IsDownloadRemoved() {
				return errors.New("aria2 download was removed")
			}
		}
	}
}

// getStatus retrieves the current status of the download
func (t *Task) getStatus(ctx context.Context) (*aria2.Status, error) {
	logger := log.FromContext(ctx)

	// Try active/waiting queue first
	status, err := t.Aria2Client.TellStatus(ctx, t.GID)
	if err == nil {
		return status, nil
	}

	// Check stopped queue
	logger.Debugf("Task not in active queue, checking stopped queue")
	stoppedTasks, stopErr := t.Aria2Client.TellStopped(ctx, -1, 100)
	if stopErr != nil {
		return nil, fmt.Errorf("failed to get aria2 status: %w", err)
	}

	for _, task := range stoppedTasks {
		if task.GID == t.GID {
			logger.Debugf("Found task in stopped queue with status: %s", task.Status)
			return &task, nil
		}
	}

	return nil, fmt.Errorf("task GID %s not found: %w", t.GID, err)
}

// transferFiles transfers downloaded files from aria2 to storage
func (t *Task) transferFiles(ctx context.Context) error {
	logger := log.FromContext(ctx)

	status, err := t.Aria2Client.TellStatus(ctx, t.GID)
	if err != nil {
		return fmt.Errorf("failed to get final status: %w", err)
	}

	if len(status.Files) == 0 {
		return errors.New("no files in aria2 download")
	}

	logger.Infof("Transferring %d file(s) to storage %s", len(status.Files), t.Storage.Name())
	transferredCount := 0

	for _, file := range status.Files {
		if file.Selected != "true" {
			logger.Debugf("Skipping unselected file: %s", file.Path)
			continue
		}

		fileName := filepath.Base(file.Path)

		// Skip torrent metadata files
		if filepath.Ext(fileName) == ".torrent" {
			logger.Debugf("Skipping torrent metadata file: %s", fileName)
			t.removeFileIfNeeded(file.Path)
			continue
		}

		if err := t.transferFile(ctx, file.Path); err != nil {
			return err
		}

		transferredCount++
		t.removeFileIfNeeded(file.Path)
	}

	if transferredCount == 0 {
		return errors.New("no files were transferred")
	}

	return nil
}

// transferFile transfers a single file to storage
func (t *Task) transferFile(ctx context.Context, filePath string) error {
	logger := log.FromContext(ctx)

	// Check if file exists
	fileInfo, err := os.Stat(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			logger.Warnf("Downloaded file not found: %s", filePath)
			return nil // Not a fatal error, continue with other files
		}
		return fmt.Errorf("failed to stat file %s: %w", filePath, err)
	}

	// Open file
	f, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("failed to open file %s: %w", filePath, err)
	}
	defer f.Close()

	// Set content length in context for storage
	ctx = context.WithValue(ctx, ctxkey.ContentLength, fileInfo.Size())

	// Save to storage
	fileName := filepath.Base(filePath)
	destPath := filepath.Join(t.StorPath, fileName)

	logger.Infof("Transferring file %s to %s:%s", fileName, t.Storage.Name(), destPath)

	if err := t.Storage.Save(ctx, f, destPath); err != nil {
		return fmt.Errorf("failed to save file %s to storage: %w", fileName, err)
	}

	logger.Infof("Successfully transferred file %s", fileName)
	return nil
}

// removeFileIfNeeded removes a file if RemoveAfterTransfer is enabled
func (t *Task) removeFileIfNeeded(filePath string) {
	if config.C().Aria2.KeepFile {
		return
	}

	logger := log.FromContext(t.ctx)
	if err := os.Remove(filePath); err != nil {
		logger.Warnf("Failed to remove local file %s: %v", filePath, err)
	} else {
		logger.Debugf("Removed local file %s", filePath)
	}
}

// cancelAria2Download cancels the aria2 download task
func (t *Task) cancelAria2Download() {
	logger := log.FromContext(t.ctx)
	logger.Infof("Canceling aria2 download GID: %s", t.GID)

	// Use a background context with timeout for cleanup
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// Try to force remove the download
	if _, err := t.Aria2Client.ForceRemove(ctx, t.GID); err != nil {
		logger.Warnf("Failed to cancel aria2 download %s: %v", t.GID, err)
	} else {
		logger.Infof("Successfully canceled aria2 download %s", t.GID)
	}

	// Also remove the download result to clean up
	if _, err := t.Aria2Client.RemoveDownloadResult(ctx, t.GID); err != nil {
		logger.Debugf("Failed to remove download result for %s: %v", t.GID, err)
	}
}