gooseek/backend/internal/computer/memory.go

package computer

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
)

type MemoryStore struct {
	repo  MemoryRepository
	cache map[string][]MemoryEntry
	mu    sync.RWMutex
}

func NewMemoryStore(repo MemoryRepository) *MemoryStore {
	return &MemoryStore{
		repo:  repo,
		cache: make(map[string][]MemoryEntry),
	}
}

func (m *MemoryStore) Store(ctx context.Context, userID string, entry *MemoryEntry) error {
	if entry.ID == "" {
		entry.ID = uuid.New().String()
	}
	entry.UserID = userID
	if entry.CreatedAt.IsZero() {
		entry.CreatedAt = time.Now()
	}

	if m.repo != nil {
		if err := m.repo.Store(ctx, entry); err != nil {
			return err
		}
	}

	m.mu.Lock()
	m.cache[userID] = append(m.cache[userID], *entry)
	if len(m.cache[userID]) > 1000 {
		m.cache[userID] = m.cache[userID][len(m.cache[userID])-500:]
	}
	m.mu.Unlock()

	return nil
}

func (m *MemoryStore) StoreResult(ctx context.Context, userID, taskID, key string, value interface{}) error {
	valueJSON, _ := json.Marshal(value)

	entry := &MemoryEntry{
		UserID:    userID,
		TaskID:    taskID,
		Key:       key,
		Value:     string(valueJSON),
		Type:      MemoryTypeResult,
		CreatedAt: time.Now(),
	}

	return m.Store(ctx, userID, entry)
}

func (m *MemoryStore) StoreFact(ctx context.Context, userID, key string, value interface{}, tags []string) error {
	entry := &MemoryEntry{
		UserID:    userID,
		Key:       key,
		Value:     value,
		Type:      MemoryTypeFact,
		Tags:      tags,
		CreatedAt: time.Now(),
	}

	return m.Store(ctx, userID, entry)
}

func (m *MemoryStore) StorePreference(ctx context.Context, userID, key string, value interface{}) error {
	entry := &MemoryEntry{
		UserID:    userID,
		Key:       key,
		Value:     value,
		Type:      MemoryTypePreference,
		CreatedAt: time.Now(),
	}

	return m.Store(ctx, userID, entry)
}

func (m *MemoryStore) StoreContext(ctx context.Context, userID, taskID, key string, value interface{}, ttl time.Duration) error {
	expiresAt := time.Now().Add(ttl)

	entry := &MemoryEntry{
		UserID:    userID,
		TaskID:    taskID,
		Key:       key,
		Value:     value,
		Type:      MemoryTypeContext,
		CreatedAt: time.Now(),
		ExpiresAt: &expiresAt,
	}

	return m.Store(ctx, userID, entry)
}

func (m *MemoryStore) Recall(ctx context.Context, userID string, query string, limit int) ([]MemoryEntry, error) {
	if m.repo != nil {
		entries, err := m.repo.Search(ctx, userID, query, limit)
		if err == nil && len(entries) > 0 {
			return entries, nil
		}
	}

	m.mu.RLock()
	cached := m.cache[userID]
	m.mu.RUnlock()

	if len(cached) == 0 {
		return nil, nil
	}

	queryLower := strings.ToLower(query)
	queryTerms := strings.Fields(queryLower)

	type scored struct {
		entry MemoryEntry
		score int
	}

	var results []scored
	now := time.Now()

	for _, entry := range cached {
		if entry.ExpiresAt != nil && entry.ExpiresAt.Before(now) {
			continue
		}

		score := 0

		keyLower := strings.ToLower(entry.Key)
		for _, term := range queryTerms {
			if strings.Contains(keyLower, term) {
				score += 3
			}
		}

		if valueStr, ok := entry.Value.(string); ok {
			valueLower := strings.ToLower(valueStr)
			for _, term := range queryTerms {
				if strings.Contains(valueLower, term) {
					score += 1
				}
			}
		}

		for _, tag := range entry.Tags {
			tagLower := strings.ToLower(tag)
			for _, term := range queryTerms {
				if strings.Contains(tagLower, term) {
					score += 2
				}
			}
		}

		if score > 0 {
			results = append(results, scored{entry: entry, score: score})
		}
	}

	for i := 0; i < len(results)-1; i++ {
		for j := i + 1; j < len(results); j++ {
			if results[j].score > results[i].score {
				results[i], results[j] = results[j], results[i]
			}
		}
	}

	if len(results) > limit {
		results = results[:limit]
	}

	entries := make([]MemoryEntry, len(results))
	for i, r := range results {
		entries[i] = r.entry
	}

	return entries, nil
}

func (m *MemoryStore) GetByUser(ctx context.Context, userID string, limit int) ([]MemoryEntry, error) {
	if m.repo != nil {
		return m.repo.GetByUser(ctx, userID, limit)
	}

	m.mu.RLock()
	cached := m.cache[userID]
	m.mu.RUnlock()

	if len(cached) > limit {
		return cached[len(cached)-limit:], nil
	}

	return cached, nil
}

func (m *MemoryStore) GetByTask(ctx context.Context, taskID string) ([]MemoryEntry, error) {
	if m.repo != nil {
		return m.repo.GetByTask(ctx, taskID)
	}

	var result []MemoryEntry

	m.mu.RLock()
	for _, entries := range m.cache {
		for _, e := range entries {
			if e.TaskID == taskID {
				result = append(result, e)
			}
		}
	}
	m.mu.RUnlock()

	return result, nil
}

func (m *MemoryStore) GetTaskContext(ctx context.Context, taskID string) (map[string]interface{}, error) {
	entries, err := m.GetByTask(ctx, taskID)
	if err != nil {
		return nil, err
	}

	context := make(map[string]interface{})
	for _, e := range entries {
		context[e.Key] = e.Value
	}

	return context, nil
}

func (m *MemoryStore) GetUserContext(ctx context.Context, userID string) (map[string]interface{}, error) {
	entries, err := m.GetByUser(ctx, userID, 100)
	if err != nil {
		return nil, err
	}

	context := make(map[string]interface{})
	for _, e := range entries {
		if e.Type == MemoryTypePreference || e.Type == MemoryTypeFact {
			context[e.Key] = e.Value
		}
	}

	return context, nil
}

func (m *MemoryStore) GetPreferences(ctx context.Context, userID string) (map[string]interface{}, error) {
	entries, err := m.GetByUser(ctx, userID, 100)
	if err != nil {
		return nil, err
	}

	prefs := make(map[string]interface{})
	for _, e := range entries {
		if e.Type == MemoryTypePreference {
			prefs[e.Key] = e.Value
		}
	}

	return prefs, nil
}

func (m *MemoryStore) GetFacts(ctx context.Context, userID string) ([]MemoryEntry, error) {
	entries, err := m.GetByUser(ctx, userID, 100)
	if err != nil {
		return nil, err
	}

	var facts []MemoryEntry
	for _, e := range entries {
		if e.Type == MemoryTypeFact {
			facts = append(facts, e)
		}
	}

	return facts, nil
}

func (m *MemoryStore) Delete(ctx context.Context, id string) error {
	if m.repo != nil {
		return m.repo.Delete(ctx, id)
	}

	m.mu.Lock()
	for userID, entries := range m.cache {
		for i, e := range entries {
			if e.ID == id {
				m.cache[userID] = append(entries[:i], entries[i+1:]...)
				break
			}
		}
	}
	m.mu.Unlock()

	return nil
}

func (m *MemoryStore) Clear(ctx context.Context, userID string) error {
	m.mu.Lock()
	delete(m.cache, userID)
	m.mu.Unlock()

	return nil
}

func (m *MemoryStore) ClearTask(ctx context.Context, taskID string) error {
	m.mu.Lock()
	for userID, entries := range m.cache {
		var filtered []MemoryEntry
		for _, e := range entries {
			if e.TaskID != taskID {
				filtered = append(filtered, e)
			}
		}
		m.cache[userID] = filtered
	}
	m.mu.Unlock()

	return nil
}

func (m *MemoryStore) Cleanup(ctx context.Context) error {
	now := time.Now()

	m.mu.Lock()
	for userID, entries := range m.cache {
		var valid []MemoryEntry
		for _, e := range entries {
			if e.ExpiresAt == nil || e.ExpiresAt.After(now) {
				valid = append(valid, e)
			}
		}
		m.cache[userID] = valid
	}
	m.mu.Unlock()

	return nil
}

func (m *MemoryStore) Stats(userID string) map[string]int {
	m.mu.RLock()
	entries := m.cache[userID]
	m.mu.RUnlock()

	stats := map[string]int{
		"total":       len(entries),
		"facts":       0,
		"preferences": 0,
		"context":     0,
		"results":     0,
	}

	for _, e := range entries {
		switch e.Type {
		case MemoryTypeFact:
			stats["facts"]++
		case MemoryTypePreference:
			stats["preferences"]++
		case MemoryTypeContext:
			stats["context"]++
		case MemoryTypeResult:
			stats["results"]++
		}
	}

	return stats
}