internal/memory/working.go

// Package memory provides working memory management for current session state.
package memory

import (
	"sync"
	"time"
)

// WorkingMemory manages the current session's active memory state.
// It holds recent entries and provides fast access to current context.
type WorkingMemory struct {
	mu sync.RWMutex

	// SessionID identifies the current session
	SessionID string

	// UserID identifies the user
	UserID string

	// TurnCount is the number of turns in this session
	TurnCount int

	// RecentEntries holds the most recent memory entries
	RecentEntries []*MemoryEntry

	// Summary is a condensed summary of older entries
	Summary string

	// SizeBytes tracks the current memory size
	SizeBytes int64

	// MaxSizeBytes is the maximum allowed size
	MaxSizeBytes int64

	// CreatedAt is when this working memory was created
	CreatedAt time.Time

	// LastAccessedAt is the last access time
	LastAccessedAt time.Time

	// MaxEntries limits the number of entries
	MaxEntries int
}

// NewWorkingMemory creates a new working memory instance
func NewWorkingMemory(sessionID, userID string, maxSizeBytes int64) *WorkingMemory {
	return &WorkingMemory{
		SessionID:      sessionID,
		UserID:         userID,
		TurnCount:      0,
		RecentEntries:  make([]*MemoryEntry, 0, 100),
		MaxSizeBytes:   maxSizeBytes,
		CreatedAt:      time.Now(),
		LastAccessedAt: time.Now(),
		MaxEntries:     100,
	}
}

// Add adds a new entry to working memory, pruning if necessary
func (wm *WorkingMemory) Add(entry *MemoryEntry) {
	wm.mu.Lock()
	defer wm.mu.Unlock()

	wm.RecentEntries = append(wm.RecentEntries, entry)
	wm.TurnCount++
	wm.SizeBytes += int64(len(entry.Content))
	wm.LastAccessedAt = time.Now()

	// Prune if exceeding size limit
	for wm.SizeBytes > wm.MaxSizeBytes && len(wm.RecentEntries) > 1 {
		removed := wm.RecentEntries[0]
		wm.RecentEntries = wm.RecentEntries[1:]
		wm.SizeBytes -= int64(len(removed.Content))
	}

	// Prune if exceeding entry limit
	for len(wm.RecentEntries) > wm.MaxEntries {
		removed := wm.RecentEntries[0]
		wm.RecentEntries = wm.RecentEntries[1:]
		wm.SizeBytes -= int64(len(removed.Content))
	}
}

// AddBatch adds multiple entries efficiently
func (wm *WorkingMemory) AddBatch(entries []*MemoryEntry) {
	wm.mu.Lock()
	defer wm.mu.Unlock()

	for _, entry := range entries {
		wm.RecentEntries = append(wm.RecentEntries, entry)
		wm.TurnCount++
		wm.SizeBytes += int64(len(entry.Content))
	}

	wm.LastAccessedAt = time.Now()

	// Prune after batch add
	wm.pruneUnlocked()
}

// pruneUnlocked handles pruning without acquiring locks (caller must hold lock)
func (wm *WorkingMemory) pruneUnlocked() {
	for wm.SizeBytes > wm.MaxSizeBytes && len(wm.RecentEntries) > 1 {
		removed := wm.RecentEntries[0]
		wm.RecentEntries = wm.RecentEntries[1:]
		wm.SizeBytes -= int64(len(removed.Content))
	}

	for len(wm.RecentEntries) > wm.MaxEntries {
		removed := wm.RecentEntries[0]
		wm.RecentEntries = wm.RecentEntries[1:]
		wm.SizeBytes -= int64(len(removed.Content))
	}
}

// GetRecent returns the N most recent entries
func (wm *WorkingMemory) GetRecent(n int) []*MemoryEntry {
	wm.mu.RLock()
	defer wm.mu.RUnlock()

	wm.LastAccessedAt = time.Now()

	if n >= len(wm.RecentEntries) {
		result := make([]*MemoryEntry, len(wm.RecentEntries))
		copy(result, wm.RecentEntries)
		return result
	}

	start := len(wm.RecentEntries) - n
	result := make([]*MemoryEntry, n)
	copy(result, wm.RecentEntries[start:])
	return result
}

// GetAll returns all entries in working memory
func (wm *WorkingMemory) GetAll() []*MemoryEntry {
	wm.mu.RLock()
	defer wm.mu.RUnlock()

	result := make([]*MemoryEntry, len(wm.RecentEntries))
	copy(result, wm.RecentEntries)
	return result
}

// GetByType returns entries filtered by memory type
func (wm *WorkingMemory) GetByType(memType MemoryType) []*MemoryEntry {
	wm.mu.RLock()
	defer wm.mu.RUnlock()

	var result []*MemoryEntry
	for _, entry := range wm.RecentEntries {
		if entry.Type == memType {
			result = append(result, entry)
		}
	}
	return result
}

// GetBySource returns entries filtered by source
func (wm *WorkingMemory) GetBySource(source string) []*MemoryEntry {
	wm.mu.RLock()
	defer wm.mu.RUnlock()

	var result []*MemoryEntry
	for _, entry := range wm.RecentEntries {
		if entry.Source == source {
			result = append(result, entry)
		}
	}
	return result
}

// Clear removes all entries from working memory
func (wm *WorkingMemory) Clear() {
	wm.mu.Lock()
	defer wm.mu.Unlock()

	wm.RecentEntries = make([]*MemoryEntry, 0, 100)
	wm.SizeBytes = 0
	wm.Summary = ""
}

// SetSummary updates the summary of older entries
func (wm *WorkingMemory) SetSummary(summary string) {
	wm.mu.Lock()
	defer wm.mu.Unlock()
	wm.Summary = summary
}

// GetSummary returns the current summary
func (wm *WorkingMemory) GetSummary() string {
	wm.mu.RLock()
	defer wm.mu.RUnlock()
	return wm.Summary
}

// Size returns the current size in bytes
func (wm *WorkingMemory) Size() int64 {
	wm.mu.RLock()
	defer wm.mu.RUnlock()
	return wm.SizeBytes
}

// Count returns the number of entries
func (wm *WorkingMemory) Count() int {
	wm.mu.RLock()
	defer wm.mu.RUnlock()
	return len(wm.RecentEntries)
}

// IsStale returns true if the working memory hasn't been accessed recently
func (wm *WorkingMemory) IsStale(threshold time.Duration) bool {
	wm.mu.RLock()
	defer wm.mu.RUnlock()
	return time.Since(wm.LastAccessedAt) > threshold
}

// NeedsConsolidation returns true if memory should be consolidated
func (wm *WorkingMemory) NeedsConsolidation(cfg ConsolidationConfig) bool {
	wm.mu.RLock()
	defer wm.mu.RUnlock()

	// Check turn threshold
	if wm.TurnCount >= cfg.TurnThreshold {
		return true
	}

	// Check time threshold
	if time.Since(wm.CreatedAt) >= cfg.TimeThreshold {
		return true
	}

	return false
}

// Stats returns working memory statistics
func (wm *WorkingMemory) Stats() WorkingMemoryStats {
	wm.mu.RLock()
	defer wm.mu.RUnlock()

	return WorkingMemoryStats{
		SessionID:      wm.SessionID,
		UserID:         wm.UserID,
		TurnCount:      wm.TurnCount,
		EntryCount:     len(wm.RecentEntries),
		SizeBytes:      wm.SizeBytes,
		MaxSizeBytes:   wm.MaxSizeBytes,
		HasSummary:     wm.Summary != "",
		CreatedAt:      wm.CreatedAt,
		LastAccessedAt: wm.LastAccessedAt,
		AgeSeconds:     time.Since(wm.CreatedAt).Seconds(),
		IdleSeconds:    time.Since(wm.LastAccessedAt).Seconds(),
		UtilizationPct: float64(wm.SizeBytes) / float64(wm.MaxSizeBytes) * 100,
	}
}

// WorkingMemoryStats holds statistics about working memory
type WorkingMemoryStats struct {
	SessionID      string    `json:"session_id"`
	UserID         string    `json:"user_id"`
	TurnCount      int       `json:"turn_count"`
	EntryCount     int       `json:"entry_count"`
	SizeBytes      int64     `json:"size_bytes"`
	MaxSizeBytes   int64     `json:"max_size_bytes"`
	HasSummary     bool      `json:"has_summary"`
	CreatedAt      time.Time `json:"created_at"`
	LastAccessedAt time.Time `json:"last_accessed_at"`
	AgeSeconds     float64   `json:"age_seconds"`
	IdleSeconds    float64   `json:"idle_seconds"`
	UtilizationPct float64   `json:"utilization_pct"`
}

// WorkingMemoryStore manages multiple working memory instances
type WorkingMemoryStore struct {
	mu       sync.RWMutex
	memories map[string]*WorkingMemory // keyed by sessionID
	config   *MemoryConfig
}

// NewWorkingMemoryStore creates a new store for working memories
func NewWorkingMemoryStore(config *MemoryConfig) *WorkingMemoryStore {
	return &WorkingMemoryStore{
		memories: make(map[string]*WorkingMemory),
		config:   config,
	}
}

// Get retrieves or creates a working memory for a session
func (s *WorkingMemoryStore) Get(sessionID, userID string) *WorkingMemory {
	s.mu.Lock()
	defer s.mu.Unlock()

	if wm, exists := s.memories[sessionID]; exists {
		return wm
	}

	wm := NewWorkingMemory(sessionID, userID, s.config.MaxWorkingMemorySize)
	s.memories[sessionID] = wm
	return wm
}

// Remove removes a working memory instance
func (s *WorkingMemoryStore) Remove(sessionID string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.memories, sessionID)
}

// GetStale returns session IDs of stale working memories
func (s *WorkingMemoryStore) GetStale(threshold time.Duration) []string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var stale []string
	for sessionID, wm := range s.memories {
		if wm.IsStale(threshold) {
			stale = append(stale, sessionID)
		}
	}
	return stale
}

// GetNeedingConsolidation returns session IDs that need consolidation
func (s *WorkingMemoryStore) GetNeedingConsolidation() []string {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var needConsolidation []string
	for sessionID, wm := range s.memories {
		if wm.NeedsConsolidation(s.config.Consolidation) {
			needConsolidation = append(needConsolidation, sessionID)
		}
	}
	return needConsolidation
}

// Count returns the number of active working memories
func (s *WorkingMemoryStore) Count() int {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return len(s.memories)
}

// CleanupStale removes stale working memories
func (s *WorkingMemoryStore) CleanupStale(threshold time.Duration) int {
	stale := s.GetStale(threshold)

	s.mu.Lock()
	defer s.mu.Unlock()

	for _, sessionID := range stale {
		delete(s.memories, sessionID)
	}

	return len(stale)
}

// AllStats returns stats for all working memories
func (s *WorkingMemoryStore) AllStats() []WorkingMemoryStats {
	s.mu.RLock()
	defer s.mu.RUnlock()

	stats := make([]WorkingMemoryStats, 0, len(s.memories))
	for _, wm := range s.memories {
		stats = append(stats, wm.Stats())
	}
	return stats
}