poimen-workflows/internal/pause/handler.go

package pause

import (
	"fmt"
	"sync"
	"time"
)

// PauseSignal represents a pause request
type PauseSignal struct {
	WorkflowID   string    `json:"workflow_id"`
	Reason       string    `json:"reason"`
	RequestedAt  time.Time `json:"requested_at"`
	GracePeriod  time.Duration `json:"grace_period"`
}

// ResumeSignal represents a resume request
type ResumeSignal struct {
	WorkflowID  string    `json:"workflow_id"`
	Reason      string    `json:"reason"`
	RequestedAt time.Time `json:"requested_at"`
}

// PauseState represents the current pause/resume state
type PauseState struct {
	WorkflowID      string
	IsPaused        bool
	PausedAt        time.Time
	ResumedAt       *time.Time
	PauseReason     string
	ResumeReason    string
	CurrentSnapshot *WorkflowSnapshot
}

// PauseHandler manages workflow pause/resume operations
type PauseHandler struct {
	mu               sync.RWMutex
	snapshotManager  *SnapshotManager
	pauseStates      map[string]*PauseState
	pauseChannels    map[string]chan bool
}

// NewPauseHandler creates a new pause handler
func NewPauseHandler(snapshotManager *SnapshotManager) *PauseHandler {
	return &PauseHandler{
		snapshotManager: snapshotManager,
		pauseStates:     make(map[string]*PauseState),
		pauseChannels:   make(map[string]chan bool),
	}
}

// RequestPause requests that a workflow pause
func (ph *PauseHandler) RequestPause(signal *PauseSignal) error {
	if signal == nil {
		return fmt.Errorf("pause signal cannot be nil")
	}

	ph.mu.Lock()
	defer ph.mu.Unlock()

	state, exists := ph.pauseStates[signal.WorkflowID]
	if !exists {
		state = &PauseState{
			WorkflowID: signal.WorkflowID,
		}
		ph.pauseStates[signal.WorkflowID] = state
	}

	state.IsPaused = true
	state.PausedAt = signal.RequestedAt
	state.PauseReason = signal.Reason

	// Notify the workflow if it's listening
	if ch, exists := ph.pauseChannels[signal.WorkflowID]; exists {
		select {
		case ch <- true:
		default:
			// Channel not ready, that's OK
		}
	}

	return nil
}

// RequestResume requests that a workflow resume
func (ph *PauseHandler) RequestResume(signal *ResumeSignal) error {
	if signal == nil {
		return fmt.Errorf("resume signal cannot be nil")
	}

	ph.mu.Lock()
	defer ph.mu.Unlock()

	state, exists := ph.pauseStates[signal.WorkflowID]
	if !exists {
		return fmt.Errorf("no pause state found for workflow: %s", signal.WorkflowID)
	}

	if !state.IsPaused {
		return fmt.Errorf("workflow is not paused: %s", signal.WorkflowID)
	}

	state.IsPaused = false
	now := time.Now()
	state.ResumedAt = &now
	state.ResumeReason = signal.Reason

	// Notify the workflow if it's listening
	if ch, exists := ph.pauseChannels[signal.WorkflowID]; exists {
		select {
		case ch <- false:
		default:
			// Channel not ready, that's OK
		}
	}

	return nil
}

// IsPaused checks if a workflow is paused
func (ph *PauseHandler) IsPaused(workflowID string) bool {
	ph.mu.RLock()
	defer ph.mu.RUnlock()

	state, exists := ph.pauseStates[workflowID]
	if !exists {
		return false
	}

	return state.IsPaused
}

// GetPauseState retrieves the pause state of a workflow
func (ph *PauseHandler) GetPauseState(workflowID string) *PauseState {
	ph.mu.RLock()
	defer ph.mu.RUnlock()

	if state, exists := ph.pauseStates[workflowID]; exists {
		// Return a copy to avoid external mutations
		stateCopy := *state
		return &stateCopy
	}

	return nil
}

// WaitForPauseOrResume blocks until a pause or resume signal is received
// Returns true if paused, false if resumed
func (ph *PauseHandler) WaitForPauseOrResume(workflowID string, timeout time.Duration) (bool, error) {
	ph.mu.Lock()

	// Create or reuse channel
	var ch chan bool
	if existingCh, exists := ph.pauseChannels[workflowID]; exists {
		ch = existingCh
	} else {
		ch = make(chan bool, 1)
		ph.pauseChannels[workflowID] = ch
	}

	ph.mu.Unlock()

	// Wait for signal with timeout
	if timeout > 0 {
		select {
		case isPaused := <-ch:
			return isPaused, nil
		case <-time.After(timeout):
			return false, fmt.Errorf("pause/resume timeout")
		}
	} else {
		isPaused := <-ch
		return isPaused, nil
	}
}

// SaveSnapshot saves the current workflow state before pausing
func (ph *PauseHandler) SaveSnapshot(
	workflowID string,
	stage string,
	completedTasks, pendingTasks, failedTasks []string,
	currentTaskID, currentActivityID string,
	taskMetrics, workflowMetrics, configuration map[string]interface{},
) (*WorkflowSnapshot, error) {
	ph.mu.Lock()
	defer ph.mu.Unlock()

	snapshot, err := ph.snapshotManager.CreateSnapshot(
		workflowID,
		stage,
		completedTasks, pendingTasks, failedTasks,
		currentTaskID, currentActivityID,
		taskMetrics, workflowMetrics, configuration,
	)

	if err == nil {
		// Create or update pause state with snapshot
		if state, exists := ph.pauseStates[workflowID]; exists {
			state.CurrentSnapshot = snapshot
		} else {
			// Create a new pause state if it doesn't exist
			ph.pauseStates[workflowID] = &PauseState{
				WorkflowID:      workflowID,
				CurrentSnapshot: snapshot,
			}
		}
	}

	return snapshot, err
}

// RestoreSnapshot restores workflow state from a snapshot
func (ph *PauseHandler) RestoreSnapshot(workflowID string) (*WorkflowSnapshot, error) {
	ph.mu.Lock()
	defer ph.mu.Unlock()

	snapshot, err := ph.snapshotManager.RestoreFromSnapshot(workflowID)
	if err != nil {
		return nil, err
	}

	// Update pause state
	if state, exists := ph.pauseStates[workflowID]; exists {
		state.CurrentSnapshot = snapshot
	}

	return snapshot, nil
}

// ResetPauseState clears pause state for a workflow (after successful completion)
func (ph *PauseHandler) ResetPauseState(workflowID string) error {
	ph.mu.Lock()
	defer ph.mu.Unlock()

	delete(ph.pauseStates, workflowID)

	// Close and remove channel if exists
	if ch, exists := ph.pauseChannels[workflowID]; exists {
		close(ch)
		delete(ph.pauseChannels, workflowID)
	}

	// Delete snapshot
	return ph.snapshotManager.DeleteSnapshot(workflowID)
}

// GetAllPauseStates returns all pause states
func (ph *PauseHandler) GetAllPauseStates() []*PauseState {
	ph.mu.RLock()
	defer ph.mu.RUnlock()

	states := make([]*PauseState, 0, len(ph.pauseStates))
	for _, state := range ph.pauseStates {
		stateCopy := *state
		states = append(states, &stateCopy)
	}

	return states
}

// GetPauseStats returns statistics about pause states
func (ph *PauseHandler) GetPauseStats() map[string]interface{} {
	ph.mu.RLock()
	defer ph.mu.RUnlock()

	paused := 0
	resumed := 0

	for _, state := range ph.pauseStates {
		if state.IsPaused {
			paused++
		} else if state.ResumedAt != nil {
			resumed++
		}
	}

	return map[string]interface{}{
		"total":   len(ph.pauseStates),
		"paused":  paused,
		"resumed": resumed,
	}
}