poimen-workflows/internal/pause/snapshot.go

package pause

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// WorkflowSnapshot represents a complete snapshot of workflow state
type WorkflowSnapshot struct {
	WorkflowID        string                 `json:"workflow_id"`
	Timestamp         time.Time              `json:"timestamp"`
	Stage             string                 `json:"stage"`
	CompletedTasks    []string               `json:"completed_tasks"`
	PendingTasks      []string               `json:"pending_tasks"`
	FailedTasks       []string               `json:"failed_tasks"`
	CurrentTaskID     string                 `json:"current_task_id"`
	CurrentActivityID string                 `json:"current_activity_id"`
	TaskMetrics       map[string]interface{} `json:"task_metrics"`
	WorkflowMetrics   map[string]interface{} `json:"workflow_metrics"`
	Configuration    map[string]interface{} `json:"configuration"`
	Error             string                 `json:"error,omitempty"`
	PausedAt          time.Time              `json:"paused_at"`
	ResumedAt         *time.Time             `json:"resumed_at,omitempty"`
}

// SnapshotManager manages workflow state snapshots for pause/resume
type SnapshotManager struct {
	mu         sync.RWMutex
	basePath   string
	snapshots  map[string]*WorkflowSnapshot
	lastSnapshot *WorkflowSnapshot
}

// NewSnapshotManager creates a new snapshot manager
func NewSnapshotManager(basePath string) *SnapshotManager {
	return &SnapshotManager{
		basePath:  basePath,
		snapshots: make(map[string]*WorkflowSnapshot),
	}
}

// CreateSnapshot creates and persists a workflow snapshot
func (sm *SnapshotManager) CreateSnapshot(
	workflowID string,
	stage string,
	completedTasks, pendingTasks, failedTasks []string,
	currentTaskID, currentActivityID string,
	taskMetrics, workflowMetrics, configuration map[string]interface{},
) (*WorkflowSnapshot, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	snapshot := &WorkflowSnapshot{
		WorkflowID:       workflowID,
		Timestamp:        time.Now(),
		Stage:            stage,
		CompletedTasks:   completedTasks,
		PendingTasks:     pendingTasks,
		FailedTasks:      failedTasks,
		CurrentTaskID:    currentTaskID,
		CurrentActivityID: currentActivityID,
		TaskMetrics:      taskMetrics,
		WorkflowMetrics:  workflowMetrics,
		Configuration:    configuration,
		PausedAt:         time.Now(),
	}

	sm.snapshots[workflowID] = snapshot
	sm.lastSnapshot = snapshot

	return snapshot, sm.persistLocked(workflowID, snapshot)
}

// GetLatestSnapshot retrieves the latest snapshot for a workflow
func (sm *SnapshotManager) GetLatestSnapshot(workflowID string) *WorkflowSnapshot {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	return sm.snapshots[workflowID]
}

// HasSnapshot checks if a snapshot exists for a workflow
func (sm *SnapshotManager) HasSnapshot(workflowID string) bool {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	_, exists := sm.snapshots[workflowID]
	return exists
}

// RestoreFromSnapshot restores workflow state from a snapshot
func (sm *SnapshotManager) RestoreFromSnapshot(workflowID string) (*WorkflowSnapshot, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	snapshot, exists := sm.snapshots[workflowID]
	if !exists {
		// Try to load from disk
		return nil, fmt.Errorf("no snapshot found for workflow: %s", workflowID)
	}

	// Mark as resumed
	now := time.Now()
	snapshot.ResumedAt = &now

	return snapshot, nil
}

// MarkResumed updates a snapshot as resumed
func (sm *SnapshotManager) MarkResumed(workflowID string) error {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	snapshot, exists := sm.snapshots[workflowID]
	if !exists {
		return fmt.Errorf("no snapshot found for workflow: %s", workflowID)
	}

	now := time.Now()
	snapshot.ResumedAt = &now

	return sm.persistLocked(workflowID, snapshot)
}

// Load loads snapshots from disk
func (sm *SnapshotManager) Load() error {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	snapshotDir := filepath.Join(sm.basePath, "snapshots")
	entries, err := os.ReadDir(snapshotDir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil // Directory doesn't exist yet
		}
		return err
	}

	for _, entry := range entries {
		if !entry.IsDir() && filepath.Ext(entry.Name()) == ".json" {
			data, err := os.ReadFile(filepath.Join(snapshotDir, entry.Name()))
			if err != nil {
				continue
			}

			var snapshot WorkflowSnapshot
			if err := json.Unmarshal(data, &snapshot); err != nil {
				continue
			}

			sm.snapshots[snapshot.WorkflowID] = &snapshot
		}
	}

	return nil
}

// persistLocked saves a snapshot to disk (must be called with lock held)
func (sm *SnapshotManager) persistLocked(workflowID string, snapshot *WorkflowSnapshot) error {
	snapshotDir := filepath.Join(sm.basePath, "snapshots")

	// Create directory if it doesn't exist
	if err := os.MkdirAll(snapshotDir, 0755); err != nil {
		return err
	}

	snapshotPath := filepath.Join(snapshotDir, fmt.Sprintf("%s.snapshot.json", workflowID))

	data, err := json.MarshalIndent(snapshot, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(snapshotPath, data, 0644)
}

// DeleteSnapshot deletes a snapshot (after successful completion)
func (sm *SnapshotManager) DeleteSnapshot(workflowID string) error {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	delete(sm.snapshots, workflowID)

	snapshotPath := filepath.Join(sm.basePath, "snapshots", fmt.Sprintf("%s.snapshot.json", workflowID))
	if _, err := os.Stat(snapshotPath); err == nil {
		return os.Remove(snapshotPath)
	}

	return nil
}

// GetAllSnapshots returns all snapshots
func (sm *SnapshotManager) GetAllSnapshots() []*WorkflowSnapshot {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	snapshots := make([]*WorkflowSnapshot, 0, len(sm.snapshots))
	for _, snapshot := range sm.snapshots {
		snapshots = append(snapshots, snapshot)
	}

	return snapshots
}

// GetLastSnapshot returns the last snapshot created
func (sm *SnapshotManager) GetLastSnapshot() *WorkflowSnapshot {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	return sm.lastSnapshot
}

// GetSnapshotStats returns statistics about snapshots
func (sm *SnapshotManager) GetSnapshotStats() map[string]interface{} {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	paused := 0
	resumed := 0

	for _, snapshot := range sm.snapshots {
		if snapshot.ResumedAt != nil {
			resumed++
		} else {
			paused++
		}
	}

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

// ClearOldSnapshots removes snapshots older than the specified duration
func (sm *SnapshotManager) ClearOldSnapshots(maxAge time.Duration) (int, error) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	now := time.Now()
	toDelete := make([]string, 0)

	for wfID, snapshot := range sm.snapshots {
		if now.Sub(snapshot.PausedAt) > maxAge {
			toDelete = append(toDelete, wfID)
		}
	}

	for _, wfID := range toDelete {
		delete(sm.snapshots, wfID)
		snapshotPath := filepath.Join(sm.basePath, "snapshots", fmt.Sprintf("%s.snapshot.json", wfID))
		_ = os.Remove(snapshotPath)
	}

	return len(toDelete), nil
}