SaveAny-Bot/api/progress.go

package api

import (
	"sync"
	"time"

	"github.com/krau/SaveAny-Bot/pkg/taskevent"
)

// TaskProgressInfo stores the progress of an API-submitted task. All fields are
// guarded by mu. It implements taskevent.Sink so the task layer can update it
// without knowing about the API.
type TaskProgressInfo struct {
	mu              sync.Mutex
	TaskID          string
	Type            string
	Status          TaskStatus
	Title           string
	TotalBytes      int64
	DownloadedBytes int64
	TotalFiles      int
	DownloadedFiles int
	Storage         string
	Path            string
	Error           string
	CreatedAt       time.Time
	UpdatedAt       time.Time
	StartedAt       time.Time
	Webhook         string
	webhookNotified bool
}

// progressStore holds all API tasks. Entries are removed a fixed duration after
// they reach a terminal state to bound memory usage.
type progressStore struct {
	mu        sync.RWMutex
	tasks     map[string]*TaskProgressInfo
	retention time.Duration
}

var store = &progressStore{
	tasks:     make(map[string]*TaskProgressInfo),
	retention: 24 * time.Hour,
}

// RegisterTask registers a new API task and returns its progress info.
func RegisterTask(taskID, taskType, storage, path, title, webhook string) *TaskProgressInfo {
	info := &TaskProgressInfo{
		TaskID:    taskID,
		Type:      taskType,
		Status:    TaskStatusQueued,
		Title:     title,
		Storage:   storage,
		Path:      path,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
		Webhook:   webhook,
	}

	store.mu.Lock()
	store.tasks[taskID] = info
	store.mu.Unlock()

	return info
}

// GetTask returns the progress info for a task.
func GetTask(taskID string) (*TaskProgressInfo, bool) {
	store.mu.RLock()
	defer store.mu.RUnlock()
	info, ok := store.tasks[taskID]
	return info, ok
}

// GetAllTasks returns all tracked tasks.
func GetAllTasks() []*TaskProgressInfo {
	store.mu.RLock()
	defer store.mu.RUnlock()

	tasks := make([]*TaskProgressInfo, 0, len(store.tasks))
	for _, info := range store.tasks {
		tasks = append(tasks, info)
	}
	return tasks
}

// DeleteTask removes a task record.
func DeleteTask(taskID string) {
	store.mu.Lock()
	defer store.mu.Unlock()
	delete(store.tasks, taskID)
}

// CleanupExpired removes tasks that reached a terminal state more than the
// store's retention duration ago. It is safe to call periodically.
func CleanupExpired() {
	now := time.Now()
	store.mu.Lock()
	defer store.mu.Unlock()
	for id, info := range store.tasks {
		info.mu.Lock()
		terminal := info.Status == TaskStatusCompleted || info.Status == TaskStatusFailed || info.Status == TaskStatusCancelled
		stale := terminal && now.Sub(info.UpdatedAt) > store.retention
		info.mu.Unlock()
		if stale {
			delete(store.tasks, id)
		}
	}
}

// StartCleanupLoop runs CleanupExpired on a fixed interval until ctx is done.
// It should be started once during API server initialization.
func StartCleanupLoop(ctx interface{ Done() <-chan struct{} }) {
	go func() {
		ticker := time.NewTicker(10 * time.Minute)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				CleanupExpired()
			}
		}
	}()
}

// UpdateStatus sets the task status.
func (t *TaskProgressInfo) UpdateStatus(status TaskStatus) {
	t.mu.Lock()
	t.Status = status
	t.UpdatedAt = time.Now()
	if status == TaskStatusRunning && t.StartedAt.IsZero() {
		t.StartedAt = t.UpdatedAt
	}
	t.mu.Unlock()
}

// SetError marks the task failed with an error message.
func (t *TaskProgressInfo) SetError(err string) {
	t.mu.Lock()
	t.Error = err
	t.Status = TaskStatusFailed
	t.UpdatedAt = time.Now()
	t.mu.Unlock()
}

// snapshot returns a point-in-time copy of the fields needed to render a
// response, so callers never touch the mutex directly.
func (t *TaskProgressInfo) snapshot() (status TaskStatus, total, downloaded int64, totalFiles, downloadedFiles int, startedAt time.Time, err string, updatedAt time.Time) {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.Status, t.TotalBytes, t.DownloadedBytes, t.TotalFiles, t.DownloadedFiles, t.StartedAt, t.Error, t.UpdatedAt
}

// Emit implements taskevent.Sink. It translates task lifecycle events into
// status/progress updates and fires the webhook on terminal transitions.
func (t *TaskProgressInfo) Emit(e taskevent.Event) {
	t.mu.Lock()
	switch e.Phase {
	case taskevent.PhaseStart:
		t.Status = TaskStatusRunning
		if t.StartedAt.IsZero() {
			t.StartedAt = time.Now()
		}
		if e.TotalBytes > 0 {
			t.TotalBytes = e.TotalBytes
		}
	case taskevent.PhaseProgress:
		t.Status = TaskStatusRunning
		if e.TotalBytes > 0 {
			t.TotalBytes = e.TotalBytes
		}
		t.DownloadedBytes = e.DownloadedBytes
		if e.TotalFiles > 0 {
			t.TotalFiles = e.TotalFiles
		}
		if e.DownloadedFiles > 0 {
			t.DownloadedFiles = e.DownloadedFiles
		}
	case taskevent.PhaseDone:
		if e.Err != nil {
			t.Status = TaskStatusFailed
			t.Error = e.Err.Error()
		} else {
			t.Status = TaskStatusCompleted
		}
	}
	t.UpdatedAt = time.Now()
	notify := t.Webhook != "" && !t.webhookNotified && (t.Status == TaskStatusCompleted || t.Status == TaskStatusFailed)
	if notify {
		t.webhookNotified = true
	}
	t.mu.Unlock()

	if notify {
		payload := CreateWebhookPayload(t.TaskID, t.Type, t.Status, t.Storage, t.Path, e.Err)
		SendWebhook(nil, payload)
	}
}