poimen-workflows/internal/cache/result_cache.go

package cache

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

// CacheKey represents a cache key for an activity result
type CacheKey struct {
	ActivityType string // "implementer", "judge", "planner"
	TaskID       string
	InputHash    string // MD5 hash of input
	ModelID      string // LLM model used
}

// String returns a string representation of the cache key
func (ck *CacheKey) String() string {
	return fmt.Sprintf("%s:%s:%s:%s", ck.ActivityType, ck.TaskID, ck.InputHash, ck.ModelID)
}

// CacheEntry represents a cached activity result
type CacheEntry struct {
	Key       CacheKey              `json:"key"`
	Result    map[string]interface{} `json:"result"`
	CreatedAt time.Time             `json:"created_at"`
	HitCount  int                   `json:"hit_count"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`
}

// ResultCache caches activity results to avoid redundant computations
type ResultCache struct {
	mu        sync.RWMutex
	basePath  string
	cache     map[string]*CacheEntry
	maxSize   int
	ttl       time.Duration
}

// NewResultCache creates a new result cache
func NewResultCache(basePath string, maxSize int, ttl time.Duration) *ResultCache {
	return &ResultCache{
		basePath: basePath,
		cache:    make(map[string]*CacheEntry),
		maxSize:  maxSize,
		ttl:      ttl,
	}
}

// ComputeHash computes a hash of the input data
func ComputeHash(data interface{}) (string, error) {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return "", err
	}

	hash := md5.Sum(jsonData)
	return fmt.Sprintf("%x", hash), nil
}

// Set stores a result in the cache
func (rc *ResultCache) Set(key *CacheKey, result map[string]interface{}) error {
	if key == nil {
		return fmt.Errorf("cache key cannot be nil")
	}

	rc.mu.Lock()
	defer rc.mu.Unlock()

	keyStr := key.String()

	entry := &CacheEntry{
		Key:       *key,
		Result:    result,
		CreatedAt: time.Now(),
		Metadata:  make(map[string]interface{}),
	}

	// Check size limit
	if len(rc.cache) >= rc.maxSize && rc.cache[keyStr] == nil {
		// Evict oldest entry (simple FIFO)
		var oldestKey string
		var oldestTime time.Time

		for k, v := range rc.cache {
			if oldestTime.IsZero() || v.CreatedAt.Before(oldestTime) {
				oldestKey = k
				oldestTime = v.CreatedAt
			}
		}

		if oldestKey != "" {
			delete(rc.cache, oldestKey)
		}
	}

	rc.cache[keyStr] = entry
	return rc.persistLocked(keyStr, entry)
}

// Get retrieves a result from the cache
func (rc *ResultCache) Get(key *CacheKey) (map[string]interface{}, bool, error) {
	if key == nil {
		return nil, false, fmt.Errorf("cache key cannot be nil")
	}

	rc.mu.Lock()
	defer rc.mu.Unlock()

	keyStr := key.String()
	entry, exists := rc.cache[keyStr]

	if !exists {
		return nil, false, nil
	}

	// Check TTL
	if rc.ttl > 0 && time.Since(entry.CreatedAt) > rc.ttl {
		delete(rc.cache, keyStr)
		return nil, false, nil
	}

	// Increment hit count
	entry.HitCount++
	_ = rc.persistLocked(keyStr, entry)

	return entry.Result, true, nil
}

// Invalidate removes a cache entry
func (rc *ResultCache) Invalidate(key *CacheKey) error {
	if key == nil {
		return fmt.Errorf("cache key cannot be nil")
	}

	rc.mu.Lock()
	defer rc.mu.Unlock()

	keyStr := key.String()
	delete(rc.cache, keyStr)

	// Delete from disk
	cacheFile := filepath.Join(rc.basePath, "cache", fmt.Sprintf("%s.json", keyStr))
	_ = os.Remove(cacheFile)

	return nil
}

// Clear clears all cache entries
func (rc *ResultCache) Clear() error {
	rc.mu.Lock()
	defer rc.mu.Unlock()

	rc.cache = make(map[string]*CacheEntry)

	// Clear disk cache
	cacheDir := filepath.Join(rc.basePath, "cache")
	_ = os.RemoveAll(cacheDir)

	return nil
}

// GetStats returns cache statistics
func (rc *ResultCache) GetStats() map[string]interface{} {
	rc.mu.RLock()
	defer rc.mu.RUnlock()

	totalHits := 0
	for _, entry := range rc.cache {
		totalHits += entry.HitCount
	}

	return map[string]interface{}{
		"size":        len(rc.cache),
		"max_size":    rc.maxSize,
		"total_hits":  totalHits,
		"usage_ratio": float64(len(rc.cache)) / float64(rc.maxSize),
	}
}

// GetSize returns the current cache size
func (rc *ResultCache) GetSize() int {
	rc.mu.RLock()
	defer rc.mu.RUnlock()

	return len(rc.cache)
}

// persistLocked saves a cache entry to disk (must be called with lock held)
func (rc *ResultCache) persistLocked(keyStr string, entry *CacheEntry) error {
	cacheDir := filepath.Join(rc.basePath, "cache")

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

	cacheFile := filepath.Join(cacheDir, fmt.Sprintf("%s.json", keyStr))

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

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

// Load loads cache from disk
func (rc *ResultCache) Load() error {
	rc.mu.Lock()
	defer rc.mu.Unlock()

	cacheDir := filepath.Join(rc.basePath, "cache")
	entries, err := os.ReadDir(cacheDir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil // Cache doesn't exist yet
		}
		return err
	}

	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		filePath := filepath.Join(cacheDir, entry.Name())
		data, err := os.ReadFile(filePath)
		if err != nil {
			continue
		}

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

		// Skip expired entries
		if rc.ttl > 0 && time.Since(cacheEntry.CreatedAt) > rc.ttl {
			continue
		}

		keyStr := cacheEntry.Key.String()
		rc.cache[keyStr] = &cacheEntry
	}

	return nil
}

// InvalidateByActivity invalidates all cache entries for an activity type
func (rc *ResultCache) InvalidateByActivity(activityType string) error {
	rc.mu.Lock()
	defer rc.mu.Unlock()

	keysToDelete := make([]string, 0)
	for keyStr, entry := range rc.cache {
		if entry.Key.ActivityType == activityType {
			keysToDelete = append(keysToDelete, keyStr)
		}
	}

	for _, keyStr := range keysToDelete {
		delete(rc.cache, keyStr)

		// Delete from disk
		cacheFile := filepath.Join(rc.basePath, "cache", fmt.Sprintf("%s.json", keyStr))
		_ = os.Remove(cacheFile)
	}

	return nil
}

// InvalidateByTask invalidates all cache entries for a task
func (rc *ResultCache) InvalidateByTask(taskID string) error {
	rc.mu.Lock()
	defer rc.mu.Unlock()

	keysToDelete := make([]string, 0)
	for keyStr, entry := range rc.cache {
		if entry.Key.TaskID == taskID {
			keysToDelete = append(keysToDelete, keyStr)
		}
	}

	for _, keyStr := range keysToDelete {
		delete(rc.cache, keyStr)

		// Delete from disk
		cacheFile := filepath.Join(rc.basePath, "cache", fmt.Sprintf("%s.json", keyStr))
		_ = os.Remove(cacheFile)
	}

	return nil
}

// GetHitRate returns the cache hit rate
func (rc *ResultCache) GetHitRate() (float64, int) {
	rc.mu.RLock()
	defer rc.mu.RUnlock()

	if len(rc.cache) == 0 {
		return 0, 0
	}

	totalHits := 0
	for _, entry := range rc.cache {
		totalHits += entry.HitCount
	}

	if totalHits == 0 {
		return 0, len(rc.cache)
	}

	return float64(totalHits) / float64(len(rc.cache)), len(rc.cache)
}