osmedeus/pkg/cli/worker_queue.go

package cli

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/j3ssie/osmedeus/v5/internal/config"
	"github.com/j3ssie/osmedeus/v5/internal/core"
	"github.com/j3ssie/osmedeus/v5/internal/database"
	"github.com/j3ssie/osmedeus/v5/internal/database/repository"
	"github.com/j3ssie/osmedeus/v5/internal/distributed"
	"github.com/j3ssie/osmedeus/v5/internal/executor"
	"github.com/j3ssie/osmedeus/v5/internal/json"
	"github.com/j3ssie/osmedeus/v5/internal/logger"
	"github.com/j3ssie/osmedeus/v5/internal/parser"
	"github.com/j3ssie/osmedeus/v5/internal/terminal"
	"github.com/spf13/cobra"
	"go.uber.org/zap"
)

// QueuedTask represents a task picked from either Redis or DB
type QueuedTask struct {
	Source        string // "redis" or "db"
	RunUUID       string
	RunID         int64
	WorkflowName  string
	WorkflowKind  string
	Target        string
	Params        map[string]interface{}
	InputIsFile   bool
	InputFilePath string
}

// QueuePollerConfig holds configuration for a QueuePoller instance.
type QueuePollerConfig struct {
	DBPollInterval    time.Duration // e.g. 5s for worker queue run, 30s for server
	RedisBRPOPTimeout time.Duration // default 5s
	Concurrency       int           // number of parallel executors
	RedisURL          string        // optional override (applied before Start)
}

// QueuePoller polls DB (and optionally Redis) for queued runs and executes them.
type QueuePoller struct {
	cfg       *config.Config
	pollerCfg QueuePollerConfig
	cancel    context.CancelFunc
	wg        sync.WaitGroup
	mu        sync.Mutex
	running   bool
}

// NewQueuePoller creates a new QueuePoller. Call Start to begin processing.
func NewQueuePoller(cfg *config.Config, pollerCfg QueuePollerConfig) *QueuePoller {
	if pollerCfg.DBPollInterval == 0 {
		pollerCfg.DBPollInterval = 5 * time.Second
	}
	if pollerCfg.RedisBRPOPTimeout == 0 {
		pollerCfg.RedisBRPOPTimeout = 5 * time.Second
	}
	if pollerCfg.Concurrency <= 0 {
		pollerCfg.Concurrency = 1
	}
	return &QueuePoller{cfg: cfg, pollerCfg: pollerCfg}
}

// Start connects to the database, starts pollers and executor goroutines.
// It returns immediately after launching goroutines. Use Stop to shut down.
func (qp *QueuePoller) Start(ctx context.Context) error {
	qp.mu.Lock()
	defer qp.mu.Unlock()
	if qp.running {
		return fmt.Errorf("queue poller already running")
	}

	log := logger.Get()

	// Apply Redis URL override
	if qp.pollerCfg.RedisURL != "" {
		redisCfg, err := distributed.ParseRedisURL(qp.pollerCfg.RedisURL)
		if err != nil {
			return fmt.Errorf("invalid redis URL: %w", err)
		}
		qp.cfg.Redis = *redisCfg
	}

	// Connect to database and migrate
	dbConn, err := database.Connect(qp.cfg)
	if err != nil {
		return fmt.Errorf("failed to connect to database: %w", err)
	}
	if err := database.Migrate(ctx); err != nil {
		return fmt.Errorf("failed to migrate database: %w", err)
	}

	innerCtx, cancel := context.WithCancel(ctx)
	qp.cancel = cancel
	qp.running = true

	taskCh := make(chan QueuedTask, qp.pollerCfg.Concurrency*2)

	// Deduplication
	var seenMu sync.Mutex
	seen := make(map[string]bool)
	markSeen := func(runUUID string) bool {
		seenMu.Lock()
		defer seenMu.Unlock()
		if seen[runUUID] {
			return false
		}
		seen[runUUID] = true
		return true
	}

	runRepo := repository.NewRunRepository(dbConn)
	p := terminal.NewPrinter()

	// DB poller goroutine
	qp.wg.Add(1)
	go func() {
		defer qp.wg.Done()
		ticker := time.NewTicker(qp.pollerCfg.DBPollInterval)
		defer ticker.Stop()

		for {
			select {
			case <-innerCtx.Done():
				return
			case <-ticker.C:
				runs, err := runRepo.ListQueued(innerCtx)
				if err != nil {
					if innerCtx.Err() != nil {
						return
					}
					log.Debug("DB poller error", zap.Error(err))
					continue
				}

				for _, run := range runs {
					if !markSeen(run.RunUUID) {
						continue
					}

					claimed, err := runRepo.ClaimQueuedRun(innerCtx, run.ID)
					if err != nil || !claimed {
						continue
					}

					taskCh <- QueuedTask{
						Source:        "db",
						RunUUID:       run.RunUUID,
						RunID:         run.ID,
						WorkflowName:  run.WorkflowName,
						WorkflowKind:  run.WorkflowKind,
						Target:        run.Target,
						Params:        run.Params,
						InputIsFile:   run.InputIsFile,
						InputFilePath: run.InputFilePath,
					}
				}
			}
		}
	}()

	// Redis poller goroutine (if configured)
	if qp.cfg.IsRedisConfigured() {
		qp.wg.Add(1)
		go func() {
			defer qp.wg.Done()
			client, err := distributed.NewClientFromConfig(qp.cfg)
			if err != nil {
				log.Debug("Redis poller: failed to connect", zap.Error(err))
				return
			}
			defer client.Close()

			if err := client.Ping(innerCtx); err != nil {
				log.Debug("Redis poller: ping failed", zap.Error(err))
				return
			}

			p.Info("Redis poller %s", terminal.Green("connected"))

			for {
				select {
				case <-innerCtx.Done():
					return
				default:
					task, err := client.PopTask(innerCtx, qp.pollerCfg.RedisBRPOPTimeout)
					if err != nil {
						if innerCtx.Err() != nil {
							return
						}
						log.Debug("Redis poller error", zap.Error(err))
						time.Sleep(time.Second)
						continue
					}
					if task == nil {
						continue
					}

					runUUID := task.ScanID
					if runUUID == "" {
						runUUID = task.ID
					}

					if !markSeen(runUUID) {
						continue
					}

					// Dedup with DB
					if task.ScanID != "" {
						var dbRun database.Run
						err := dbConn.NewSelect().
							Model(&dbRun).
							Where("run_uuid = ?", task.ScanID).
							Scan(innerCtx)
						if err == nil && (dbRun.Status == "running" || dbRun.Status == "completed") {
							log.Debug("Skipping Redis task - already processed in DB",
								zap.String("run_uuid", task.ScanID))
							continue
						}
					}

					taskCh <- QueuedTask{
						Source:        "redis",
						RunUUID:       runUUID,
						RunID:         0,
						WorkflowName:  task.WorkflowName,
						WorkflowKind:  task.WorkflowKind,
						Target:        task.Target,
						Params:        task.Params,
						InputIsFile:   task.InputIsFile,
						InputFilePath: task.InputFilePath,
					}
				}
			}
		}()
	}

	// Executor goroutines
	for i := 0; i < qp.pollerCfg.Concurrency; i++ {
		qp.wg.Add(1)
		go func(workerNum int) {
			defer qp.wg.Done()
			executeQueuedTasks(innerCtx, qp.cfg, taskCh, workerNum, p)
		}(i)
	}

	return nil
}

// Stop cancels the poller context and waits for all goroutines to finish.
func (qp *QueuePoller) Stop() {
	qp.mu.Lock()
	if !qp.running {
		qp.mu.Unlock()
		return
	}
	qp.running = false
	qp.mu.Unlock()

	if qp.cancel != nil {
		qp.cancel()
	}
	qp.wg.Wait()
}

var (
	queueConcurrency int
	queueRedisURL    string

	// Flags for worker queue new
	queueFlow       string
	queueModule     string
	queueTargets    []string
	queueTargetFile string
	queueParams     []string
)

// workerQueueCmd is the parent command for queue subcommands
var workerQueueCmd = &cobra.Command{
	Use:   "queue",
	Short: "Manage and process queued tasks",
	Long:  UsageWorkerQueue(),
}

// workerQueueListCmd lists queued tasks
var workerQueueListCmd = &cobra.Command{
	Use:     "list",
	Aliases: []string{"ls"},
	Short:   "List queued tasks",
	Long:    UsageWorkerQueueList(),
	RunE:    runWorkerQueueList,
}

// workerQueueNewCmd creates new queued tasks
var workerQueueNewCmd = &cobra.Command{
	Use:   "new",
	Short: "Queue a new task for later processing",
	Long:  UsageWorkerQueueNew(),
	RunE:  runWorkerQueueNew,
}

// workerQueueRunCmd processes queued tasks
var workerQueueRunCmd = &cobra.Command{
	Use:   "run",
	Short: "Process queued tasks from DB and Redis",
	Long:  UsageWorkerQueueRun(),
	RunE:  runWorkerQueueRun,
}

func init() {
	// Queue list flags
	workerQueueListCmd.Flags().StringVar(&queueRedisURL, "redis-url", "", "Redis connection URL (overrides settings)")

	// Queue new flags
	workerQueueNewCmd.Flags().StringVarP(&queueFlow, "flow", "f", "", "flow workflow name")
	workerQueueNewCmd.Flags().StringVarP(&queueModule, "module", "m", "", "module workflow name")
	workerQueueNewCmd.Flags().StringArrayVarP(&queueTargets, "target", "t", nil, "target(s) to queue")
	workerQueueNewCmd.Flags().StringVarP(&queueTargetFile, "target-file", "T", "", "file containing targets")
	workerQueueNewCmd.Flags().StringArrayVarP(&queueParams, "params", "p", nil, "additional parameters (key=value)")
	workerQueueNewCmd.Flags().StringVar(&queueRedisURL, "redis-url", "", "Redis connection URL (overrides settings)")

	// Queue run flags
	workerQueueRunCmd.Flags().IntVar(&queueConcurrency, "concurrency", 1, "number of concurrent task executors")
	workerQueueRunCmd.Flags().StringVar(&queueRedisURL, "redis-url", "", "Redis connection URL (overrides settings)")

	// Register subcommands
	workerQueueCmd.AddCommand(workerQueueListCmd)
	workerQueueCmd.AddCommand(workerQueueNewCmd)
	workerQueueCmd.AddCommand(workerQueueRunCmd)

	// Register queue under worker
	workerCmd.AddCommand(workerQueueCmd)
}

// runWorkerQueueList lists all queued runs from the database
func runWorkerQueueList(cmd *cobra.Command, args []string) error {
	cfg := config.Get()
	if cfg == nil {
		return errConfigNotLoaded
	}

	ctx := context.Background()

	// Connect to database
	dbConn, err := database.Connect(cfg)
	if err != nil {
		return fmt.Errorf("failed to connect to database: %w", err)
	}
	if err := database.Migrate(ctx); err != nil {
		return fmt.Errorf("failed to migrate database: %w", err)
	}

	// Query all queued runs (all statuses, not just "queued")
	var runs []*database.Run
	err = dbConn.NewSelect().
		Model(&runs).
		Where("is_queued = ?", true).
		Order("created_at ASC").
		Scan(ctx)
	if err != nil {
		return fmt.Errorf("failed to query queued runs: %w", err)
	}

	// JSON output
	if globalJSON {
		if len(runs) == 0 {
			fmt.Println("[]")
			return nil
		}
		jsonBytes, err := json.MarshalIndent(runs, "", "  ")
		if err != nil {
			return fmt.Errorf("failed to marshal runs: %w", err)
		}
		fmt.Println(string(jsonBytes))
		return nil
	}

	if len(runs) == 0 {
		printer.Info("No queued tasks found")
		return nil
	}

	// Display table
	printer.Section("Queued Tasks")

	columns := []string{"id", "uuid", "workflow", "target", "status", "priority", "is_file", "created"}
	var records []map[string]interface{}
	for _, r := range runs {
		isFileStr := ""
		if r.InputIsFile {
			isFileStr = "yes"
		}
		records = append(records, map[string]interface{}{
			"id":       fmt.Sprintf("%d", r.ID),
			"uuid":     truncateUUID(r.RunUUID),
			"workflow": r.WorkflowName,
			"target":   truncateQueueStr(r.Target, 40),
			"status":   r.Status,
			"priority": r.RunPriority,
			"is_file":  isFileStr,
			"created":  r.CreatedAt.Format("2006-01-02 15:04"),
		})
	}

	renderTableWithTablewriter("", records, columns, globalWidth, false, nil)
	fmt.Println()
	printer.Info("Total: %d queued task(s)", len(runs))
	return nil
}

// runWorkerQueueNew creates new queued tasks
func runWorkerQueueNew(cmd *cobra.Command, args []string) error {
	cfg := config.Get()
	if cfg == nil {
		return errConfigNotLoaded
	}

	// Override Redis config from URL if provided
	if queueRedisURL != "" {
		redisCfg, err := distributed.ParseRedisURL(queueRedisURL)
		if err != nil {
			return err
		}
		cfg.Redis = *redisCfg
	}

	// Determine workflow name and kind
	workflowName := queueFlow
	workflowKind := "flow"
	if workflowName == "" && queueModule != "" {
		workflowName = queueModule
		workflowKind = "module"
	}
	if workflowName == "" {
		return fmt.Errorf("workflow name required (use -f or -m)")
	}

	// Collect targets
	var allTargets []string
	allTargets = append(allTargets, queueTargets...)
	if queueTargetFile != "" {
		fileTargets, err := readTargetsFromFile(queueTargetFile)
		if err != nil {
			return fmt.Errorf("failed to read target file: %w", err)
		}
		allTargets = append(allTargets, fileTargets...)
	}
	allTargets = deduplicateTargets(allTargets)

	if len(allTargets) == 0 {
		return fmt.Errorf("no targets specified (use -t or -T)")
	}

	// Parse params
	params := make(map[string]interface{})
	for _, flag := range queueParams {
		parts := strings.SplitN(flag, "=", 2)
		if len(parts) == 2 {
			params[parts[0]] = parts[1]
		}
	}

	p := terminal.NewPrinter()
	return queueRuns(context.Background(), cfg, workflowName, workflowKind, allTargets, queueTargetFile, params, p)
}

// runWorkerQueueRun processes queued tasks with dual-source polling (DB + Redis)
func runWorkerQueueRun(cmd *cobra.Command, args []string) error {
	cfg := config.Get()
	if cfg == nil {
		return errConfigNotLoaded
	}

	// Ensure external binaries are in PATH
	ensureExternalBinariesInPath(cfg)

	if queueConcurrency <= 0 {
		queueConcurrency = 1
	}

	p := terminal.NewPrinter()

	// Setup graceful shutdown
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		p.Info("Shutting down queue runner...")
		cancel()
	}()

	qp := NewQueuePoller(cfg, QueuePollerConfig{
		DBPollInterval: 5 * time.Second,
		Concurrency:    queueConcurrency,
		RedisURL:       queueRedisURL,
	})

	if err := qp.Start(ctx); err != nil {
		return err
	}

	p.Success("Queue runner started (concurrency: %s)", terminal.Cyan(fmt.Sprintf("%d", queueConcurrency)))
	p.Info("Polling DB every %s%s...", terminal.Cyan("5s"), redisStatusSuffix(cfg))

	// Block until context is cancelled
	<-ctx.Done()
	qp.Stop()

	p.Info("Queue runner stopped")
	return nil
}

// executeQueuedTasks processes tasks from the channel
func executeQueuedTasks(ctx context.Context, cfg *config.Config, taskCh <-chan QueuedTask, workerNum int, p *terminal.Printer) {
	loader := parser.NewLoader(cfg.WorkflowsPath)

	for {
		select {
		case <-ctx.Done():
			return
		case task, ok := <-taskCh:
			if !ok {
				return
			}

			p.Info("[worker-%d] Processing: %s -> %s (source: %s)",
				workerNum, terminal.Yellow(task.WorkflowName), terminal.Green(task.Target), task.Source)

			// If InputIsFile, check if file exists locally
			if task.InputIsFile && task.InputFilePath != "" {
				if _, err := os.Stat(task.InputFilePath); os.IsNotExist(err) {
					p.Warning("[worker-%d] Input file %s not found, skipping task", workerNum, terminal.Yellow(task.InputFilePath))
					updateQueuedRunStatus(ctx, task.RunUUID, "failed", "input file not found: "+task.InputFilePath)
					continue
				}
			}

			// Load workflow
			workflow, err := loader.LoadWorkflow(task.WorkflowName)
			if err != nil {
				p.Error("[worker-%d] Failed to load workflow %s: %s", workerNum, task.WorkflowName, err)
				updateQueuedRunStatus(ctx, task.RunUUID, "failed", fmt.Sprintf("failed to load workflow: %v", err))
				continue
			}

			// Convert params to string map
			params := make(map[string]string)
			params["target"] = task.Target
			for k, v := range task.Params {
				if s, ok := v.(string); ok {
					params[k] = s
				}
			}

			// Create executor
			exec := executor.NewExecutor()
			exec.SetLoader(loader)
			exec.SetDBRunUUID(task.RunUUID)

			// Execute workflow
			var result *core.WorkflowResult
			if workflow.IsFlow() {
				result, err = exec.ExecuteFlow(ctx, workflow, params, cfg)
			} else {
				result, err = exec.ExecuteModule(ctx, workflow, params, cfg)
			}

			// Determine outcome
			if err != nil {
				p.Error("[worker-%d] Task failed for %s: %s", workerNum, task.Target, err)
				updateQueuedRunStatus(ctx, task.RunUUID, "failed", err.Error())
			} else if result != nil && result.Status == core.RunStatusFailed {
				errMsg := "workflow execution failed"
				if result.Error != nil {
					errMsg = result.Error.Error()
				}
				p.Error("[worker-%d] Task failed for %s: %s", workerNum, task.Target, errMsg)
				updateQueuedRunStatus(ctx, task.RunUUID, "failed", errMsg)
			} else {
				p.Success("[worker-%d] Task completed for %s", workerNum, terminal.Green(task.Target))
				updateQueuedRunStatus(ctx, task.RunUUID, "completed", "")
			}
		}
	}
}

// updateQueuedRunStatus updates the status of a queued run in the database
func updateQueuedRunStatus(ctx context.Context, runUUID, status, errorMsg string) {
	if runUUID == "" {
		return
	}
	_ = database.UpdateRunStatus(ctx, runUUID, status, errorMsg)
}

// truncateUUID returns a truncated UUID for display
func truncateUUID(u string) string {
	if len(u) > 8 {
		return u[:8]
	}
	return u
}

// truncateQueueStr truncates a string to maxLen with ellipsis
func truncateQueueStr(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen-3] + "..."
}

// colorizeQueueStatus applies color to queue status
func colorizeQueueStatus(status string) string {
	switch status {
	case "queued":
		return terminal.Yellow(status)
	case "running":
		return terminal.Cyan(status)
	case "completed":
		return terminal.Green(status)
	case "failed":
		return terminal.Red(status)
	default:
		return terminal.Gray(status)
	}
}

// redisStatusSuffix returns a suffix string describing Redis connectivity
func redisStatusSuffix(cfg *config.Config) string {
	if cfg.IsRedisConfigured() {
		return " + " + terminal.Cyan("Redis")
	}
	return ""
}