gems/opal-task/internal/engine/ws.go

package engine

import (
	"fmt"

	"github.com/google/uuid"
)

// WorkingSet is a mapping from small integers to task UUIDs for filtered tasks.
// The small integers are meant to be stable, easily-typed identifiers for users to interact with
// important tasks.
type WorkingSet struct {
	byUUID map[uuid.UUID]*Task
	byID   map[int]uuid.UUID // display_id -> UUID
}

// BuildWorkingSet creates a working set from a list of tasks and persists to DB
// Tasks should be pre-sorted in the desired display order
// IDs are assigned sequentially (1, 2, 3, ...) based on task order in the slice
func BuildWorkingSet(tasks []*Task) (*WorkingSet, error) {
	ws := &WorkingSet{
		byUUID: make(map[uuid.UUID]*Task),
		byID:   make(map[int]uuid.UUID),
	}

	db := GetDB()
	if db == nil {
		return nil, fmt.Errorf("database not initialized")
	}

	tx, err := db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	// Clear existing working set
	if _, err := tx.Exec("DELETE FROM working_set"); err != nil {
		return nil, err
	}

	// Insert new working set
	stmt, err := tx.Prepare("INSERT INTO working_set (display_id, task_uuid) VALUES (?, ?)")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	for i, task := range tasks {
		displayID := i + 1 // 1-based
		ws.byUUID[task.UUID] = task
		ws.byID[displayID] = task.UUID

		if _, err := stmt.Exec(displayID, task.UUID.String()); err != nil {
			return nil, err
		}
	}

	if err := tx.Commit(); err != nil {
		return nil, err
	}

	return ws, nil
}

// LoadWorkingSet loads the current working set from DB
func LoadWorkingSet() (*WorkingSet, error) {
	db := GetDB()
	if db == nil {
		return nil, fmt.Errorf("database not initialized")
	}

	rows, err := db.Query(`
		SELECT ws.display_id, ws.task_uuid
		FROM working_set ws
		ORDER BY ws.display_id
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ws := &WorkingSet{
		byUUID: make(map[uuid.UUID]*Task),
		byID:   make(map[int]uuid.UUID),
	}

	for rows.Next() {
		var displayID int
		var taskUUIDStr string

		if err := rows.Scan(&displayID, &taskUUIDStr); err != nil {
			return nil, err
		}

		taskUUID, err := uuid.Parse(taskUUIDStr)
		if err != nil {
			return nil, fmt.Errorf("failed to parse UUID: %w", err)
		}

		// Load the actual task
		task, err := GetTask(taskUUID)
		if err != nil {
			// Task might have been deleted, skip it
			continue
		}

		ws.byUUID[task.UUID] = task
		ws.byID[displayID] = task.UUID
	}

	return ws, nil
}

// GetTaskByDisplayID resolves display ID to task
func (ws *WorkingSet) GetTaskByDisplayID(id int) (*Task, error) {
	taskUUID, exists := ws.byID[id]
	if !exists {
		return nil, fmt.Errorf("invalid task ID: %d (not in current working set)", id)
	}
	task, exists := ws.byUUID[taskUUID]
	if !exists {
		return nil, fmt.Errorf("task UUID not found in working set")
	}
	return task, nil
}

// GetTasks returns all tasks in the working set
func (ws *WorkingSet) GetTasks() []*Task {
	tasks := make([]*Task, 0, len(ws.byID))
	for i := 1; i <= len(ws.byID); i++ {
		if taskUUID, ok := ws.byID[i]; ok {
			if task, ok := ws.byUUID[taskUUID]; ok {
				tasks = append(tasks, task)
			}
		}
	}
	return tasks
}

// Size returns number of tasks in working set
func (ws *WorkingSet) Size() int {
	return len(ws.byID)
}

// ByID returns the display_id -> UUID mapping.
func (ws *WorkingSet) ByID() map[int]uuid.UUID {
	return ws.byID
}

// AppendTask inserts a task into the working set at MAX(display_id) + 1.
// Returns the assigned display ID. The ID is valid until the next report
// render, at which point the entire working set is rebuilt.
func AppendTask(task *Task) (int, error) {
	db := GetDB()
	if db == nil {
		return 0, fmt.Errorf("database not initialized")
	}

	var maxID int
	err := db.QueryRow("SELECT COALESCE(MAX(display_id), 0) FROM working_set").Scan(&maxID)
	if err != nil {
		return 0, fmt.Errorf("failed to query max display_id: %w", err)
	}

	newID := maxID + 1
	_, err = db.Exec("INSERT INTO working_set (display_id, task_uuid) VALUES (?, ?)", newID, task.UUID.String())
	if err != nil {
		return 0, fmt.Errorf("failed to append to working set: %w", err)
	}

	return newID, nil
}