core/scheduler/scheduler.go

// Copyright (c) 2026 Jose Cely <me@jacs.guru>
// Use of this source code is governed by MIT-style
// license that can be found in the LICENSE file.

package scheduler

import (
	"context"
	"fmt"
	"sync"
	"time"

	"git.smarteching.com/goffee/core/logger"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// ─── Status types ────────────────────────────────────────────────────────────

// QueueStatus represents the lifecycle state of a queued item.
type QueueStatus string

const (
	StatusPending    QueueStatus = "pending"
	StatusProcessing QueueStatus = "processing" // claimed by a scheduler tick
)

// ProcessedStatus is the final outcome recorded in processed_items.
type ProcessedStatus string

const (
	ProcessedOK    ProcessedStatus = "ok"
	ProcessedError ProcessedStatus = "error"
)

// ─── Models ──────────────────────────────────────────────────────────────────

// QueueItem is a task waiting to be executed.
// Any external program can insert a row into this table to enqueue work.
//
//	INSERT INTO queue_items (task_type, payload, priority, max_runs, thread) VALUES (?, ?, ?, ?, ?);
type QueueItem struct {
	ID        uint        `gorm:"primaryKey;autoIncrement"`
	TaskType  string      `gorm:"not null;index"` // e.g. "send_email", "sync_records"
	Payload   string      `gorm:"type:text"`      // JSON or plain string, interpreted by the handler
	Status    QueueStatus `gorm:"not null;default:'pending';index"`
	Priority  int         `gorm:"not null;default:0;index"` // higher = runs first within a tick
	MaxRuns   int         `gorm:"not null;default:1"`       // 1 = run once, -1 = repeat indefinitely, N = run N times
	Thread    int         `gorm:"not null;default:0"`       // 0 = sequential, 1 = concurrent per type
	CreatedAt time.Time
	UpdatedAt time.Time
}

// TableName overrides the GORM default.
func (QueueItem) TableName() string { return "queue_items" }

// ProcessedItem is an immutable record of a completed execution.
// Rows are never updated — only inserted.
type ProcessedItem struct {
	ID          uint            `gorm:"primaryKey;autoIncrement"`
	QueueItemID uint            `gorm:"not null;index"` // FK → queue_items.id (soft ref)
	TaskType    string          `gorm:"not null;index"`
	Payload     string          `gorm:"type:text"`
	Status      ProcessedStatus `gorm:"not null;index"`
	ErrorMsg    string          `gorm:"type:text"` // empty when Status == "ok"
	StartedAt   time.Time       `gorm:"not null"`
	FinishedAt  time.Time       `gorm:"not null"`
	DurationMs  int64           `gorm:"not null"` // FinishedAt - StartedAt in ms
	CreatedAt   time.Time
}

// TableName overrides the GORM default.
func (ProcessedItem) TableName() string { return "processed_items" }

// ─── Handler ─────────────────────────────────────────────────────────────────

// HandlerFunc processes a single queue item.
// Return a non-nil error to mark the item as processed with error status.
type HandlerFunc func(ctx context.Context, item *QueueItem) error

// Registry maps task types to their handlers.
// Register all handlers before calling Scheduler.Start.
type Registry struct {
	handlers map[string]HandlerFunc
}

// NewRegistry creates an empty handler registry.
func NewRegistry() *Registry {
	return &Registry{handlers: make(map[string]HandlerFunc)}
}

// Register associates a task type with a handler.
// Panics on duplicate registration to catch misconfiguration at startup.
func (r *Registry) Register(taskType string, h HandlerFunc) {
	if _, exists := r.handlers[taskType]; exists {
		panic(fmt.Sprintf("scheduler: handler already registered for task type %q", taskType))
	}
	r.handlers[taskType] = h
}

// get returns the handler for a task type, or a fallback error handler.
func (r *Registry) get(taskType string) HandlerFunc {
	if h, ok := r.handlers[taskType]; ok {
		return h
	}
	return func(_ context.Context, item *QueueItem) error {
		return fmt.Errorf("no handler registered for task type %q", item.TaskType)
	}
}

// ─── Store ───────────────────────────────────────────────────────────────────

// Store handles all database operations for the scheduler.
type Store struct {
	db *gorm.DB
}

// NewStore creates a Store.
// Tables (queue_items, processed_items)
func NewStore(db *gorm.DB) (*Store, error) {
	if db == nil {
		return nil, fmt.Errorf("scheduler: db is nil")
	}
	return &Store{db: db}, nil
}

// Enqueue inserts a new pending task.
// maxRuns controls how many times the task executes: 1 = once (default), -1 = repeat indefinitely.
// thread controls concurrency: 0 = sequential per type, 1 = concurrent per type.
// Any value of maxRuns < -1 is clamped to 1.
func (s *Store) Enqueue(ctx context.Context, taskType, payload string, priority, maxRuns, thread int) (*QueueItem, error) {
	if maxRuns < -1 {
		maxRuns = 1
	}
	item := &QueueItem{
		TaskType: taskType,
		Payload:  payload,
		Status:   StatusPending,
		Priority: priority,
		MaxRuns:  maxRuns,
		Thread:   thread,
	}
	if err := s.db.WithContext(ctx).Create(item).Error; err != nil {
		return nil, fmt.Errorf("enqueue %s: %w", taskType, err)
	}
	return item, nil
}

// ClaimBatch atomically claims up to `limit` pending tasks and marks them
// as "processing" so that concurrent scheduler instances never double-execute.
func (s *Store) ClaimBatch(ctx context.Context, limit int) ([]*QueueItem, error) {
	var items []*QueueItem

	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		query := tx.
			Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}).
			Where("status = ?", StatusPending).
			Order("priority DESC, id ASC").
			Limit(limit).
			Find(&items)

		if query.Error != nil {
			return query.Error
		}
		if len(items) == 0 {
			return nil
		}

		ids := make([]uint, len(items))
		for i, it := range items {
			ids[i] = it.ID
		}

		return tx.Model(&QueueItem{}).
			Where("id IN ?", ids).
			Update("status", StatusProcessing).
			Error
	})

	if err != nil {
		return nil, fmt.Errorf("claim batch: %w", err)
	}
	return items, nil
}

// RecordResult writes a ProcessedItem, then either deletes the QueueItem (if
// MaxRuns is exhausted) or decrements MaxRuns and resets the status to pending
// for re-execution on the next tick. A maxRuns of -1 means repeat indefinitely.
func (s *Store) RecordResult(ctx context.Context, item *QueueItem, startedAt time.Time, execErr error) error {
	finishedAt := time.Now()

	status := ProcessedOK
	errMsg := ""
	if execErr != nil {
		status = ProcessedError
		errMsg = execErr.Error()
	}

	processed := &ProcessedItem{
		QueueItemID: item.ID,
		TaskType:    item.TaskType,
		Payload:     item.Payload,
		Status:      status,
		ErrorMsg:    errMsg,
		StartedAt:   startedAt,
		FinishedAt:  finishedAt,
		DurationMs:  finishedAt.Sub(startedAt).Milliseconds(),
	}

	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// 1. Insert the immutable result record.
		if err := tx.Create(processed).Error; err != nil {
			return fmt.Errorf("insert processed_item: %w", err)
		}

		// 2. Decide: delete the queue item or re-queue it.
		// MaxRuns == 1 means "run once" — delete.
		// MaxRuns == -1 means "repeat indefinitely" — always re-queue.
		// MaxRuns > 1 means "run N times" — decrement and re-queue.
		if item.MaxRuns == 1 {
			// Exhausted — delete.
			if err := tx.Delete(&QueueItem{}, item.ID).Error; err != nil {
				return fmt.Errorf("delete queue_item %d: %w", item.ID, err)
			}
		} else {
			// Decrement MaxRuns (if not -1/infinite) and set back to pending.
			newMaxRuns := item.MaxRuns
			if newMaxRuns > 1 {
				newMaxRuns--
			}
			if err := tx.Model(&QueueItem{}).
				Where("id = ?", item.ID).
				Updates(map[string]interface{}{
					"status":   StatusPending,
					"max_runs": newMaxRuns,
				}).Error; err != nil {
				return fmt.Errorf("re-queue queue_item %d: %w", item.ID, err)
			}
		}

		return nil
	})
}

// PendingCount returns the number of tasks currently in pending or processing state.
func (s *Store) PendingCount(ctx context.Context) (int64, error) {
	var count int64
	err := s.db.WithContext(ctx).
		Model(&QueueItem{}).
		Where("status IN ?", []QueueStatus{StatusPending, StatusProcessing}).
		Count(&count).Error
	return count, err
}

// RecentProcessed returns the last `n` processed items ordered by newest first.
func (s *Store) RecentProcessed(ctx context.Context, n int) ([]*ProcessedItem, error) {
	var items []*ProcessedItem
	err := s.db.WithContext(ctx).
		Order("id DESC").
		Limit(n).
		Find(&items).Error
	return items, err
}

// RecoverStuck resets any items stuck in "processing" state (e.g. after a
// crash) back to "pending" so they can be re-claimed on the next tick.
// Call this once at startup.
func (s *Store) RecoverStuck(ctx context.Context) (int64, error) {
	result := s.db.WithContext(ctx).
		Model(&QueueItem{}).
		Where("status = ?", StatusProcessing).
		Update("status", StatusPending)
	return result.RowsAffected, result.Error
}

// ─── Semaphore ───────────────────────────────────────────────────────────────

// Semaphore controls whether the scheduler is allowed to execute tasks.
// Green (true) means tasks can run; Red (false) means tasks are blocked.
type Semaphore struct {
	green chan bool
}

// NewSemaphore creates a Semaphore. Pass true to start green, false to start red.
func NewSemaphore(startGreen bool) *Semaphore {
	s := &Semaphore{green: make(chan bool, 1)}
	s.green <- startGreen
	return s
}

// SetGreen allows task execution to proceed.
func (s *Semaphore) SetGreen() {
	<-s.green
	s.green <- true
}

// SetRed blocks task execution.
func (s *Semaphore) SetRed() {
	<-s.green
	s.green <- false
}

// IsGreen reports whether execution is currently allowed.
func (s *Semaphore) IsGreen() bool {
	val := <-s.green
	s.green <- val
	return val
}

// ─── Scheduler ───────────────────────────────────────────────────────────────

// Scheduler polls the database every interval and executes pending tasks.
type Scheduler struct {
	interval  time.Duration
	rateLimit int // max tasks per tick; 0 = unlimited
	semaphore *Semaphore
	store     *Store
	registry  *Registry

	stop chan struct{}
	done chan struct{}
}

// SchedulerConfig holds constructor options.
type SchedulerConfig struct {
	Interval  time.Duration
	RateLimit int // 0 = unlimited
}

// NewScheduler creates a Scheduler.
func NewScheduler(cfg SchedulerConfig, sem *Semaphore, st *Store, reg *Registry) *Scheduler {
	return &Scheduler{
		interval:  cfg.Interval,
		rateLimit: cfg.RateLimit,
		semaphore: sem,
		store:     st,
		registry:  reg,
		stop:      make(chan struct{}),
		done:      make(chan struct{}),
	}
}

// Start launches the scheduler goroutine.
func (s *Scheduler) Start(ctx context.Context) {
	go s.run(ctx)
}

// Stop signals the scheduler and waits for it to exit cleanly.
func (s *Scheduler) Stop() {
	close(s.stop)
	<-s.done
}

// run is the main event loop.
func (s *Scheduler) run(ctx context.Context) {
	defer close(s.done)

	ticker := time.NewTicker(s.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			logger.ResolveLogger().Info("[scheduler] context cancelled, shutting down")
			return
		case <-s.stop:
			logger.ResolveLogger().Info("[scheduler] stop requested, shutting down")
			return
		case t := <-ticker.C:
			s.tick(ctx, t)
		}
	}
}

// tick runs on every ticker event.
func (s *Scheduler) tick(ctx context.Context, t time.Time) {
	ts := t.Format(time.TimeOnly)

	if !s.semaphore.IsGreen() {
		logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] tick %s — semaphore RED, skipping", ts))
		return
	}

	limit := s.rateLimit
	if limit == 0 {
		limit = 1000
	}

	items, err := s.store.ClaimBatch(ctx, limit)
	if err != nil {
		logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] tick %s — claim error: %v", ts, err))
		return
	}
	if len(items) == 0 {
		return
	}

	logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] tick %s — executing %d task(s)", ts, len(items)))

	// Split items into sequential vs concurrent groups.
	var sequential []*QueueItem
	concurrentGroups := make(map[string][]*QueueItem) // taskType -> items

	for _, item := range items {
		if item.Thread != 0 {
			concurrentGroups[item.TaskType] = append(concurrentGroups[item.TaskType], item)
		} else {
			sequential = append(sequential, item)
		}
	}

	// Execute sequential items one at a time.
	for i, item := range sequential {
		s.execute(ctx, item, i+1, len(items))
	}

	// Execute concurrent groups in parallel per type.
	for _, group := range concurrentGroups {
		s.executeGroup(ctx, group, len(items))
	}

	if remaining, err := s.store.PendingCount(ctx); err == nil {
		logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] tick %s — done (%d task(s) still queued)", ts, remaining))
	}
}

// execute runs one task and records the result regardless of outcome.
func (s *Scheduler) execute(ctx context.Context, item *QueueItem, pos, total int) {
	logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] starting task id=%d type=%s", pos, total, item.ID, item.TaskType))

	startedAt := time.Now()
	handler := s.registry.get(item.TaskType)
	execErr := handler(ctx, item)

	if execErr != nil {
		logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] task id=%d FAILED: %v", pos, total, item.ID, execErr))
	} else {
		logger.ResolveLogger().Info(fmt.Sprintf("[scheduler] [%d/%d] task id=%d OK (%.0fms)",
			pos, total, item.ID, float64(time.Since(startedAt).Milliseconds())))
	}

	if err := s.store.RecordResult(ctx, item, startedAt, execErr); err != nil {
		logger.ResolveLogger().Error(fmt.Sprintf("[scheduler] [%d/%d] failed to record result for id=%d: %v", pos, total, item.ID, err))
	}
}

// executeGroup runs all items in a group concurrently using goroutines.
// Each item's result is recorded independently.
func (s *Scheduler) executeGroup(ctx context.Context, items []*QueueItem, total int) {
	var wg sync.WaitGroup
	wg.Add(len(items))

	for _, item := range items {
		item := item // capture range variable
		go func() {
			defer wg.Done()
			// Find the item's position among all items for logging context.
			// Since we don't have global position here, use a local index.
			s.execute(ctx, item, 0, total)
		}()
	}

	wg.Wait()
}