poimen-workflows/internal/history/pruner.go

package history

import (
	"encoding/json"
	"fmt"
	"os"
	"sort"
	"sync"
	"time"
)

// TaskHistory represents a single task execution in history
type TaskHistory struct {
	TaskID      string                 `json:"task_id"`
	Status      string                 `json:"status"` // "pending", "completed", "failed"
	StartTime   time.Time              `json:"start_time"`
	EndTime     time.Time              `json:"end_time"`
	Duration    time.Duration          `json:"duration"`
	Output      map[string]interface{} `json:"output,omitempty"`
	Error       string                 `json:"error,omitempty"`
	Metrics     map[string]interface{} `json:"metrics,omitempty"`
	Size        int64                  `json:"size"` // Estimated size in bytes
}

// PrunePolicy defines how to prune history
type PrunePolicy struct {
	MaxHistorySize int64         // Max total history size in bytes (e.g., 100MB)
	MaxHistoryAge  time.Duration // Max age of history entries (e.g., 24 hours)
	MaxEntries     int           // Max number of entries to keep (e.g., 1000)
	ArchiveDir     string        // Directory to archive pruned items
}

// HistoryPruner manages workflow history with automatic pruning
type HistoryPruner struct {
	mu              sync.RWMutex
	entries         []*TaskHistory
	policy          PrunePolicy
	totalSize       int64
	pruneCount      int
	archiveCount    int
	lastPruneTime   time.Time
	pruneThreshold  int64 // Size threshold that triggers pruning
}

// NewHistoryPruner creates a new history pruner
func NewHistoryPruner(policy PrunePolicy) *HistoryPruner {
	if policy.MaxHistorySize == 0 {
		policy.MaxHistorySize = 100 * 1024 * 1024 // 100MB default
	}
	if policy.MaxHistoryAge == 0 {
		policy.MaxHistoryAge = 24 * time.Hour // 24 hours default
	}
	if policy.MaxEntries == 0 {
		policy.MaxEntries = 1000 // 1000 entries default
	}

	// Set prune threshold at 90% of max size
	pruneThreshold := (policy.MaxHistorySize * 9) / 10

	return &HistoryPruner{
		entries:        make([]*TaskHistory, 0),
		policy:         policy,
		pruneThreshold: pruneThreshold,
	}
}

// AddEntry adds a task history entry
func (hp *HistoryPruner) AddEntry(entry *TaskHistory) error {
	if entry == nil {
		return fmt.Errorf("entry cannot be nil")
	}

	hp.mu.Lock()
	defer hp.mu.Unlock()

	// Estimate size
	data, _ := json.Marshal(entry)
	entry.Size = int64(len(data))

	hp.entries = append(hp.entries, entry)
	hp.totalSize += entry.Size

	// Check if pruning is needed
	if hp.totalSize > hp.pruneThreshold || len(hp.entries) > hp.policy.MaxEntries {
		hp.pruneLocked()
	}

	return nil
}

// pruneLocked prunes old entries based on policy (must be called with lock held)
func (hp *HistoryPruner) pruneLocked() {
	if len(hp.entries) == 0 {
		return
	}

	// Sort by end time (oldest first)
	sort.Slice(hp.entries, func(i, j int) bool {
		return hp.entries[i].EndTime.Before(hp.entries[j].EndTime)
	})

	// Archive old entries
	var toKeep []*TaskHistory
	newTotalSize := int64(0)
	now := time.Now()

	for _, entry := range hp.entries {
		age := now.Sub(entry.EndTime)

		// Keep if:
		// 1. Newer than max age, AND
		// 2. Total size not exceeded, AND
		// 3. Not too many entries
		if age < hp.policy.MaxHistoryAge &&
			newTotalSize+entry.Size <= hp.policy.MaxHistorySize &&
			len(toKeep) < hp.policy.MaxEntries {
			toKeep = append(toKeep, entry)
			newTotalSize += entry.Size
		} else {
			// Archive this entry
			hp.archiveEntry(entry)
			hp.archiveCount++
		}
	}

	hp.entries = toKeep
	hp.totalSize = newTotalSize
	hp.pruneCount++
	hp.lastPruneTime = time.Now()
}

// archiveEntry archives an entry to disk (must be called with lock held)
func (hp *HistoryPruner) archiveEntry(entry *TaskHistory) {
	if hp.policy.ArchiveDir == "" {
		return // No archive directory configured
	}

	// Create archive directory if it doesn't exist
	_ = os.MkdirAll(hp.policy.ArchiveDir, 0755)

	// Save entry to archive file
	timestamp := time.Now().Unix()
	archivePath := fmt.Sprintf("%s/history-%s-%d.json", hp.policy.ArchiveDir, entry.TaskID, timestamp)

	data, _ := json.MarshalIndent(entry, "", "  ")
	_ = os.WriteFile(archivePath, data, 0644)
}

// Prune manually triggers pruning
func (hp *HistoryPruner) Prune() {
	hp.mu.Lock()
	defer hp.mu.Unlock()

	hp.pruneLocked()
}

// GetSize returns total history size
func (hp *HistoryPruner) GetSize() int64 {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	return hp.totalSize
}

// GetEntryCount returns number of entries in history
func (hp *HistoryPruner) GetEntryCount() int {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	return len(hp.entries)
}

// GetStats returns pruning statistics
func (hp *HistoryPruner) GetStats() map[string]interface{} {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	avgEntrySize := int64(0)
	if len(hp.entries) > 0 {
		avgEntrySize = hp.totalSize / int64(len(hp.entries))
	}

	return map[string]interface{}{
		"total_size":          hp.totalSize,
		"entry_count":         len(hp.entries),
		"avg_entry_size":      avgEntrySize,
		"max_allowed_size":    hp.policy.MaxHistorySize,
		"max_allowed_age":     hp.policy.MaxHistoryAge,
		"max_allowed_entries": hp.policy.MaxEntries,
		"prune_count":         hp.pruneCount,
		"archive_count":       hp.archiveCount,
		"last_prune_time":     hp.lastPruneTime,
		"usage_ratio":         float64(hp.totalSize) / float64(hp.policy.MaxHistorySize),
	}
}

// GetEntries returns a copy of all entries
func (hp *HistoryPruner) GetEntries() []*TaskHistory {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	result := make([]*TaskHistory, len(hp.entries))
	copy(result, hp.entries)

	return result
}

// GetEntriesByStatus returns entries filtered by status
func (hp *HistoryPruner) GetEntriesByStatus(status string) []*TaskHistory {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	var result []*TaskHistory
	for _, entry := range hp.entries {
		if entry.Status == status {
			result = append(result, entry)
		}
	}

	return result
}

// GetRecentEntries returns the most recent N entries
func (hp *HistoryPruner) GetRecentEntries(count int) []*TaskHistory {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	if count > len(hp.entries) {
		count = len(hp.entries)
	}

	// Sort by end time descending (newest first)
	sorted := make([]*TaskHistory, len(hp.entries))
	copy(sorted, hp.entries)
	sort.Slice(sorted, func(i, j int) bool {
		return sorted[i].EndTime.After(sorted[j].EndTime)
	})

	return sorted[:count]
}

// Clear clears all history
func (hp *HistoryPruner) Clear() {
	hp.mu.Lock()
	defer hp.mu.Unlock()

	hp.entries = make([]*TaskHistory, 0)
	hp.totalSize = 0
}

// GetEntry returns a specific entry by task ID
func (hp *HistoryPruner) GetEntry(taskID string) (*TaskHistory, bool) {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	for _, entry := range hp.entries {
		if entry.TaskID == taskID {
			return entry, true
		}
	}

	return nil, false
}

// CalculateMemorySavings calculates estimated memory saved by pruning
func (hp *HistoryPruner) CalculateMemorySavings() int64 {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	// Estimated savings: total pruned size minus current size
	// This is an approximation based on how much was archived
	savedSize := int64(hp.archiveCount) * (hp.totalSize / int64(len(hp.entries) + 1))
	return savedSize
}

// ShouldPrune checks if pruning is needed
func (hp *HistoryPruner) ShouldPrune() bool {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	return hp.totalSize > hp.pruneThreshold || len(hp.entries) > hp.policy.MaxEntries
}

// GetMemoryInfo returns memory usage information
func (hp *HistoryPruner) GetMemoryInfo() map[string]interface{} {
	hp.mu.RLock()
	defer hp.mu.RUnlock()

	return map[string]interface{}{
		"current_size":      hp.totalSize,
		"max_size":          hp.policy.MaxHistorySize,
		"current_entries":   len(hp.entries),
		"max_entries":       hp.policy.MaxEntries,
		"usage_percentage":  float64(hp.totalSize*100) / float64(hp.policy.MaxHistorySize),
		"entries_percentage": float64(len(hp.entries)*100) / float64(hp.policy.MaxEntries),
	}
}

// UpdateEntry updates an existing entry
func (hp *HistoryPruner) UpdateEntry(taskID string, updates map[string]interface{}) error {
	hp.mu.Lock()
	defer hp.mu.Unlock()

	for _, entry := range hp.entries {
		if entry.TaskID == taskID {
			// Apply updates
			for key, value := range updates {
				switch key {
				case "status":
					entry.Status = value.(string)
				case "output":
					entry.Output = value.(map[string]interface{})
				case "error":
					entry.Error = value.(string)
				case "end_time":
					entry.EndTime = value.(time.Time)
					entry.Duration = entry.EndTime.Sub(entry.StartTime)
				}
			}

			// Recalculate size
			data, _ := json.Marshal(entry)
			newSize := int64(len(data))
			hp.totalSize = hp.totalSize - entry.Size + newSize
			entry.Size = newSize

			return nil
		}
	}

	return fmt.Errorf("entry not found: %s", taskID)
}