mq/dedup_and_flow.go

package mq

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/oarkflow/mq/logger"
)

// DedupEntry represents a deduplication cache entry
type DedupEntry struct {
	MessageID   string
	ContentHash string
	FirstSeen   time.Time
	LastSeen    time.Time
	Count       int
}

// DeduplicationManager manages message deduplication
type DeduplicationManager struct {
	cache           map[string]*DedupEntry
	mu              sync.RWMutex
	window          time.Duration
	cleanupInterval time.Duration
	shutdown        chan struct{}
	logger          logger.Logger
	persistent      DedupStorage
	onDuplicate     func(*DedupEntry)
}

// DedupStorage interface for persistent deduplication storage
type DedupStorage interface {
	Store(ctx context.Context, entry *DedupEntry) error
	Get(ctx context.Context, key string) (*DedupEntry, error)
	Delete(ctx context.Context, key string) error
	DeleteOlderThan(ctx context.Context, duration time.Duration) (int, error)
	Close() error
}

// DedupConfig holds configuration for deduplication
type DedupConfig struct {
	Window          time.Duration // Time window for deduplication
	CleanupInterval time.Duration
	Persistent      DedupStorage // Optional persistent storage
	Logger          logger.Logger
}

// NewDeduplicationManager creates a new deduplication manager
func NewDeduplicationManager(config DedupConfig) *DeduplicationManager {
	if config.Window == 0 {
		config.Window = 5 * time.Minute
	}
	if config.CleanupInterval == 0 {
		config.CleanupInterval = 1 * time.Minute
	}

	dm := &DeduplicationManager{
		cache:           make(map[string]*DedupEntry),
		window:          config.Window,
		cleanupInterval: config.CleanupInterval,
		shutdown:        make(chan struct{}),
		logger:          config.Logger,
		persistent:      config.Persistent,
	}

	go dm.cleanupLoop()

	return dm
}

// CheckDuplicate checks if a message is a duplicate
func (dm *DeduplicationManager) CheckDuplicate(ctx context.Context, task *Task) (bool, error) {
	// Generate dedup key from task
	dedupKey := dm.generateDedupKey(task)

	dm.mu.Lock()
	defer dm.mu.Unlock()

	// Check in-memory cache
	if entry, exists := dm.cache[dedupKey]; exists {
		// Check if within window
		if time.Since(entry.FirstSeen) < dm.window {
			entry.LastSeen = time.Now()
			entry.Count++

			if dm.onDuplicate != nil {
				go dm.onDuplicate(entry)
			}

			dm.logger.Debug("Duplicate message detected",
				logger.Field{Key: "dedupKey", Value: dedupKey},
				logger.Field{Key: "count", Value: entry.Count},
				logger.Field{Key: "taskID", Value: task.ID})

			return true, nil
		}

		// Entry expired, remove it
		delete(dm.cache, dedupKey)
	}

	// Check persistent storage if available
	if dm.persistent != nil {
		entry, err := dm.persistent.Get(ctx, dedupKey)
		if err == nil && time.Since(entry.FirstSeen) < dm.window {
			entry.LastSeen = time.Now()
			entry.Count++
			dm.cache[dedupKey] = entry

			if dm.onDuplicate != nil {
				go dm.onDuplicate(entry)
			}

			return true, nil
		}
	}

	// Not a duplicate, add to cache
	entry := &DedupEntry{
		MessageID:   task.ID,
		ContentHash: dedupKey,
		FirstSeen:   time.Now(),
		LastSeen:    time.Now(),
		Count:       1,
	}

	dm.cache[dedupKey] = entry

	// Persist if storage available
	if dm.persistent != nil {
		go dm.persistent.Store(ctx, entry)
	}

	return false, nil
}

// generateDedupKey generates a deduplication key from a task
func (dm *DeduplicationManager) generateDedupKey(task *Task) string {
	// If task has explicit dedup key, use it
	if task.DedupKey != "" {
		return task.DedupKey
	}

	// Otherwise, hash the content
	hasher := sha256.New()
	hasher.Write([]byte(task.Topic))
	hasher.Write(task.Payload)

	// Include headers in hash for more precise deduplication
	if task.Headers != nil {
		headerBytes, _ := json.Marshal(task.Headers)
		hasher.Write(headerBytes)
	}

	return hex.EncodeToString(hasher.Sum(nil))
}

// cleanupLoop periodically cleans up expired entries
func (dm *DeduplicationManager) cleanupLoop() {
	ticker := time.NewTicker(dm.cleanupInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			dm.cleanup()
		case <-dm.shutdown:
			return
		}
	}
}

// cleanup removes expired deduplication entries
func (dm *DeduplicationManager) cleanup() {
	dm.mu.Lock()
	defer dm.mu.Unlock()

	cutoff := time.Now().Add(-dm.window)
	removed := 0

	for key, entry := range dm.cache {
		if entry.FirstSeen.Before(cutoff) {
			delete(dm.cache, key)
			removed++
		}
	}

	if removed > 0 {
		dm.logger.Debug("Cleaned up expired dedup entries",
			logger.Field{Key: "removed", Value: removed})
	}

	// Cleanup persistent storage
	if dm.persistent != nil {
		go dm.persistent.DeleteOlderThan(context.Background(), dm.window)
	}
}

// SetOnDuplicate sets callback for duplicate detection
func (dm *DeduplicationManager) SetOnDuplicate(fn func(*DedupEntry)) {
	dm.onDuplicate = fn
}

// GetStats returns deduplication statistics
func (dm *DeduplicationManager) GetStats() map[string]interface{} {
	dm.mu.RLock()
	defer dm.mu.RUnlock()

	totalDuplicates := 0
	for _, entry := range dm.cache {
		totalDuplicates += entry.Count - 1 // Subtract 1 for original message
	}

	return map[string]interface{}{
		"cache_size":       len(dm.cache),
		"total_duplicates": totalDuplicates,
		"window":           dm.window,
	}
}

// Shutdown stops the deduplication manager
func (dm *DeduplicationManager) Shutdown(ctx context.Context) error {
	close(dm.shutdown)

	if dm.persistent != nil {
		return dm.persistent.Close()
	}

	return nil
}

// FlowController manages backpressure and flow control
type FlowController struct {
	credits          int64
	maxCredits       int64
	minCredits       int64
	creditRefillRate int64
	mu               sync.Mutex
	logger           logger.Logger
	shutdown         chan struct{}
	refillInterval   time.Duration
	onCreditLow      func(current, max int64)
	onCreditHigh     func(current, max int64)
}

// FlowControlConfig holds flow control configuration
type FlowControlConfig struct {
	MaxCredits     int64
	MinCredits     int64
	RefillRate     int64 // Credits to add per interval
	RefillInterval time.Duration
	Logger         logger.Logger
}

// NewFlowController creates a new flow controller
func NewFlowController(config FlowControlConfig) *FlowController {
	if config.MaxCredits == 0 {
		config.MaxCredits = 1000
	}
	if config.MinCredits == 0 {
		config.MinCredits = 100
	}
	if config.RefillRate == 0 {
		config.RefillRate = 10
	}
	if config.RefillInterval == 0 {
		config.RefillInterval = 100 * time.Millisecond
	}

	fc := &FlowController{
		credits:          config.MaxCredits,
		maxCredits:       config.MaxCredits,
		minCredits:       config.MinCredits,
		creditRefillRate: config.RefillRate,
		refillInterval:   config.RefillInterval,
		logger:           config.Logger,
		shutdown:         make(chan struct{}),
	}

	go fc.refillLoop()

	return fc
}

// AcquireCredit attempts to acquire credits for processing
func (fc *FlowController) AcquireCredit(ctx context.Context, amount int64) error {
	for {
		fc.mu.Lock()
		if fc.credits >= amount {
			fc.credits -= amount

			// Check if credits are low
			if fc.credits < fc.minCredits && fc.onCreditLow != nil {
				go fc.onCreditLow(fc.credits, fc.maxCredits)
			}

			fc.mu.Unlock()
			return nil
		}
		fc.mu.Unlock()

		// Wait before retrying
		select {
		case <-time.After(10 * time.Millisecond):
			continue
		case <-ctx.Done():
			return ctx.Err()
		case <-fc.shutdown:
			return fmt.Errorf("flow controller shutting down")
		}
	}
}

// ReleaseCredit returns credits after processing
func (fc *FlowController) ReleaseCredit(amount int64) {
	fc.mu.Lock()
	defer fc.mu.Unlock()

	fc.credits += amount
	if fc.credits > fc.maxCredits {
		fc.credits = fc.maxCredits
	}

	// Check if credits recovered
	if fc.credits > fc.maxCredits/2 && fc.onCreditHigh != nil {
		go fc.onCreditHigh(fc.credits, fc.maxCredits)
	}
}

// refillLoop periodically refills credits
func (fc *FlowController) refillLoop() {
	ticker := time.NewTicker(fc.refillInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			fc.mu.Lock()
			fc.credits += fc.creditRefillRate
			if fc.credits > fc.maxCredits {
				fc.credits = fc.maxCredits
			}
			fc.mu.Unlock()
		case <-fc.shutdown:
			return
		}
	}
}

// GetAvailableCredits returns the current available credits
func (fc *FlowController) GetAvailableCredits() int64 {
	fc.mu.Lock()
	defer fc.mu.Unlock()
	return fc.credits
}

// SetOnCreditLow sets callback for low credit warning
func (fc *FlowController) SetOnCreditLow(fn func(current, max int64)) {
	fc.onCreditLow = fn
}

// SetOnCreditHigh sets callback for credit recovery
func (fc *FlowController) SetOnCreditHigh(fn func(current, max int64)) {
	fc.onCreditHigh = fn
}

// AdjustMaxCredits dynamically adjusts maximum credits
func (fc *FlowController) AdjustMaxCredits(newMax int64) {
	fc.mu.Lock()
	defer fc.mu.Unlock()

	fc.maxCredits = newMax
	if fc.credits > newMax {
		fc.credits = newMax
	}

	fc.logger.Info("Adjusted max credits",
		logger.Field{Key: "newMax", Value: newMax})
}

// GetStats returns flow control statistics
func (fc *FlowController) GetStats() map[string]interface{} {
	fc.mu.Lock()
	defer fc.mu.Unlock()

	utilization := float64(fc.maxCredits-fc.credits) / float64(fc.maxCredits) * 100

	return map[string]interface{}{
		"credits":     fc.credits,
		"max_credits": fc.maxCredits,
		"min_credits": fc.minCredits,
		"utilization": utilization,
		"refill_rate": fc.creditRefillRate,
	}
}

// Shutdown stops the flow controller
func (fc *FlowController) Shutdown() {
	close(fc.shutdown)
}

// BackpressureMonitor monitors system backpressure
type BackpressureMonitor struct {
	queueDepthThreshold    int
	memoryThreshold        uint64
	errorRateThreshold     float64
	checkInterval          time.Duration
	logger                 logger.Logger
	shutdown               chan struct{}
	onBackpressureApplied  func(reason string)
	onBackpressureRelieved func()
}

// BackpressureConfig holds backpressure configuration
type BackpressureConfig struct {
	QueueDepthThreshold int
	MemoryThreshold     uint64
	ErrorRateThreshold  float64
	CheckInterval       time.Duration
	Logger              logger.Logger
}

// NewBackpressureMonitor creates a new backpressure monitor
func NewBackpressureMonitor(config BackpressureConfig) *BackpressureMonitor {
	if config.CheckInterval == 0 {
		config.CheckInterval = 5 * time.Second
	}
	if config.ErrorRateThreshold == 0 {
		config.ErrorRateThreshold = 0.5 // 50% error rate
	}

	bm := &BackpressureMonitor{
		queueDepthThreshold: config.QueueDepthThreshold,
		memoryThreshold:     config.MemoryThreshold,
		errorRateThreshold:  config.ErrorRateThreshold,
		checkInterval:       config.CheckInterval,
		logger:              config.Logger,
		shutdown:            make(chan struct{}),
	}

	go bm.monitorLoop()

	return bm
}

// monitorLoop continuously monitors for backpressure conditions
func (bm *BackpressureMonitor) monitorLoop() {
	ticker := time.NewTicker(bm.checkInterval)
	defer ticker.Stop()

	backpressureActive := false

	for {
		select {
		case <-ticker.C:
			shouldApply, reason := bm.shouldApplyBackpressure()

			if shouldApply && !backpressureActive {
				backpressureActive = true
				bm.logger.Warn("Applying backpressure",
					logger.Field{Key: "reason", Value: reason})
				if bm.onBackpressureApplied != nil {
					bm.onBackpressureApplied(reason)
				}
			} else if !shouldApply && backpressureActive {
				backpressureActive = false
				bm.logger.Info("Relieving backpressure")
				if bm.onBackpressureRelieved != nil {
					bm.onBackpressureRelieved()
				}
			}
		case <-bm.shutdown:
			return
		}
	}
}

// shouldApplyBackpressure checks if backpressure should be applied
func (bm *BackpressureMonitor) shouldApplyBackpressure() (bool, string) {
	var memStats runtime.MemStats
	runtime.ReadMemStats(&memStats)

	// Check memory threshold
	if bm.memoryThreshold > 0 && memStats.Alloc > bm.memoryThreshold {
		return true, fmt.Sprintf("memory threshold exceeded: %d > %d",
			memStats.Alloc, bm.memoryThreshold)
	}

	return false, ""
}

// SetOnBackpressureApplied sets callback for backpressure application
func (bm *BackpressureMonitor) SetOnBackpressureApplied(fn func(reason string)) {
	bm.onBackpressureApplied = fn
}

// SetOnBackpressureRelieved sets callback for backpressure relief
func (bm *BackpressureMonitor) SetOnBackpressureRelieved(fn func()) {
	bm.onBackpressureRelieved = fn
}

// Shutdown stops the backpressure monitor
func (bm *BackpressureMonitor) Shutdown() {
	close(bm.shutdown)
}