internal/memory/integration.go

// Package memory provides integration with the agent orchestrator.
package memory

import (
	"context"
	"log/slog"
	"time"
)

// AgentMemoryIntegration integrates the memory system with the agent orchestrator.
// It provides a high-level API for the agent to interact with memory.
type AgentMemoryIntegration struct {
	// orchestrator handles memory operations
	orchestrator *MemoryOrchestrator

	// consolidationWorker handles background consolidation
	consolidationWorker *ConsolidationWorker

	// gdprManager handles GDPR operations
	gdprManager *GDPRManager

	// logger for integration events
	logger *slog.Logger

	// config
	config *MemoryConfig
}

// NewAgentMemoryIntegration creates a new integration instance
func NewAgentMemoryIntegration(
	config *MemoryConfig,
	embedder EmbeddingClient,
	llmClient LLMClient,
) (*AgentMemoryIntegration, error) {
	if config == nil {
		config = DefaultMemoryConfig()
	}

	// Create backend based on configuration
	var backend MemoryManager

	if config.RustServiceEnabled {
		// Use Rust memory service with local fallback
		localBackend := NewLocalMemoryBackend(config)
		rustClient := NewRustMemoryClient(RustClientConfig{
			Host:           config.RustServiceAddr,
			Port:           config.RustServicePort,
			Compression:    true,
			PoolSize:       config.ConnectionPoolSize,
			RequestTimeout: config.RequestTimeout,
			Fallback:       localBackend,
		})
		backend = rustClient
	} else {
		// Use local backend only
		backend = NewLocalMemoryBackend(config)
	}

	// Create orchestrator
	orchestrator := NewMemoryOrchestrator(backend, embedder, llmClient, config)

	// Create consolidation worker
	consolidationWorker := NewConsolidationWorker(orchestrator, config.Consolidation)

	// Create GDPR manager
	gdprManager := NewGDPRManager(backend, nil)

	integration := &AgentMemoryIntegration{
		orchestrator:        orchestrator,
		consolidationWorker: consolidationWorker,
		gdprManager:         gdprManager,
		logger:              slog.Default().With("component", "memory-integration"),
		config:              config,
	}

	return integration, nil
}

// Start starts background workers
func (i *AgentMemoryIntegration) Start() {
	i.logger.Info("starting memory integration")
	i.consolidationWorker.Start()
}

// Stop stops background workers and cleans up
func (i *AgentMemoryIntegration) Stop() error {
	i.logger.Info("stopping memory integration")
	i.consolidationWorker.Stop()
	return i.orchestrator.Close()
}

// GetContextForQuery retrieves relevant memory context for a query
func (i *AgentMemoryIntegration) GetContextForQuery(
	ctx context.Context,
	userID, sessionID string,
	query string,
	maxTokens int,
) (*MemoryContext, error) {
	return i.orchestrator.ProcessQuery(ctx, &QueryRequest{
		Query:                query,
		UserID:               userID,
		SessionID:            sessionID,
		MaxTokens:            maxTokens,
		IncludeWorkingMemory: true,
		StoreToWorkingMemory: false,
	})
}

// StoreConversationTurn stores a conversation turn in memory
func (i *AgentMemoryIntegration) StoreConversationTurn(
	ctx context.Context,
	userID, sessionID string,
	userMessage, assistantResponse string,
) error {
	now := time.Now()

	// Store user message as episodic memory
	userEntry := &MemoryEntry{
		ID:         generateID("episodic", userID, now),
		Type:       EpisodicMemory,
		Content:    "User: " + userMessage,
		Timestamp:  now,
		Confidence: 1.0,
		UserID:     userID,
		SessionID:  sessionID,
		Source:     "conversation",
		Tags:       []string{"user-message"},
	}

	if err := i.orchestrator.Store(ctx, userEntry); err != nil {
		return err
	}

	// Store assistant response
	assistantEntry := &MemoryEntry{
		ID:         generateID("episodic", userID, now.Add(time.Millisecond)),
		Type:       EpisodicMemory,
		Content:    "Assistant: " + assistantResponse,
		Timestamp:  now.Add(time.Millisecond),
		Confidence: 1.0,
		UserID:     userID,
		SessionID:  sessionID,
		Source:     "conversation",
		Tags:       []string{"assistant-response"},
	}

	return i.orchestrator.Store(ctx, assistantEntry)
}

// StoreKnowledge stores semantic knowledge extracted from conversations
func (i *AgentMemoryIntegration) StoreKnowledge(
	ctx context.Context,
	userID, sessionID string,
	knowledge string,
	confidence float64,
	tags []string,
) error {
	entry := &MemoryEntry{
		ID:         generateID("semantic", userID, time.Now()),
		Type:       SemanticMemory,
		Content:    knowledge,
		Timestamp:  time.Now(),
		Confidence: confidence,
		UserID:     userID,
		SessionID:  sessionID,
		Source:     "extraction",
		Tags:       tags,
	}

	return i.orchestrator.Store(ctx, entry)
}

// StoreUserPreference stores a learned user preference
func (i *AgentMemoryIntegration) StoreUserPreference(
	ctx context.Context,
	userID, sessionID string,
	preference string,
	confidence float64,
) error {
	entry := &MemoryEntry{
		ID:         generateID("procedural", userID, time.Now()),
		Type:       ProceduralMemory,
		Content:    preference,
		Timestamp:  time.Now(),
		Confidence: confidence,
		UserID:     userID,
		SessionID:  sessionID,
		Source:     "learning",
		Tags:       []string{"preference"},
	}

	return i.orchestrator.Store(ctx, entry)
}

// GetWorkingMemory returns the current session's working memory
func (i *AgentMemoryIntegration) GetWorkingMemory(sessionID, userID string) *WorkingMemory {
	return i.orchestrator.GetWorkingMemory(sessionID, userID)
}

// ConsolidateSession consolidates a session's memory
func (i *AgentMemoryIntegration) ConsolidateSession(ctx context.Context, sessionID string) error {
	return i.orchestrator.ConsolidateSession(ctx, sessionID)
}

// DeleteUserData handles GDPR deletion requests
func (i *AgentMemoryIntegration) DeleteUserData(ctx context.Context, userID, requestedBy string) error {
	return i.gdprManager.DeleteUserData(ctx, userID, requestedBy)
}

// ExportUserData handles GDPR export requests
func (i *AgentMemoryIntegration) ExportUserData(ctx context.Context, userID, requestedBy string) (*DataExport, error) {
	return i.gdprManager.ExportUserData(ctx, userID, requestedBy)
}

// GetMetrics returns current memory metrics
func (i *AgentMemoryIntegration) GetMetrics() IntegrationMetrics {
	orchestratorMetrics := i.orchestrator.GetMetrics()
	consolidationMetrics := i.consolidationWorker.GetMetrics()

	return IntegrationMetrics{
		TotalQueries:          orchestratorMetrics.TotalQueries,
		TotalStores:           orchestratorMetrics.TotalStores,
		AvgRetrievalMs:        orchestratorMetrics.AvgRetrievalMs,
		ConsolidationRuns:     consolidationMetrics.TotalConsolidations,
		SessionsProcessed:     consolidationMetrics.SessionsProcessed,
		TTLDeletions:          consolidationMetrics.TTLDeletions,
		ActiveWorkingMemories: i.orchestrator.workingMemoryStore.Count(),
	}
}

// IntegrationMetrics combines metrics from all memory components
type IntegrationMetrics struct {
	TotalQueries          int64   `json:"total_queries"`
	TotalStores           int64   `json:"total_stores"`
	AvgRetrievalMs        float64 `json:"avg_retrieval_ms"`
	ConsolidationRuns     int64   `json:"consolidation_runs"`
	SessionsProcessed     int64   `json:"sessions_processed"`
	TTLDeletions          int64   `json:"ttl_deletions"`
	ActiveWorkingMemories int     `json:"active_working_memories"`
}

// SubscribeToMemoryEvents subscribes to real-time memory events for a user
func (i *AgentMemoryIntegration) SubscribeToMemoryEvents(userID string) chan MemoryEvent {
	return i.orchestrator.Subscribe(userID)
}

// UnsubscribeFromMemoryEvents unsubscribes from memory events
func (i *AgentMemoryIntegration) UnsubscribeFromMemoryEvents(userID string, ch chan MemoryEvent) {
	i.orchestrator.Unsubscribe(userID, ch)
}

// Helper to generate IDs
func generateID(prefix, userID string, t time.Time) string {
	return prefix + ":" + userID + ":" + t.Format("20060102150405.000000000")
}