core/server/tools/memory.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strings"

	"switchboard-core/knowledge"
	"switchboard-core/models"
	"switchboard-core/store"
)

// ── Late Registration ────────────────────────
// Memory tools use late registration because they need stores + embedder.

// RegisterMemoryTools registers the memory_save and memory_recall tools.
// Called from main.go after stores and embedder are initialized.
func RegisterMemoryTools(stores store.Stores, embedder *knowledge.Embedder) {
	if stores.Memories == nil {
		log.Println("⚠ memory tools: MemoryStore not available, skipping registration")
		return
	}

	Register(&memorySaveTool{stores: stores, embedder: embedder})
	Register(&memoryRecallTool{stores: stores, embedder: embedder})

	log.Println("✅ memory tools registered (memory_save, memory_recall)")
}

// ═══════════════════════════════════════════
// memory_save
// ═══════════════════════════════════════════

type memorySaveTool struct {
	visitorDeniedBase
	stores   store.Stores
	embedder *knowledge.Embedder
}

func (t *memorySaveTool) Definition() ToolDef {
	return ToolDef{
		Name:        "memory_save",
		DisplayName: "Save Memory",
		Category:    "memory",
		Description: "Save a fact or preference about the user for future conversations. " +
			"Use this when the user shares something worth remembering long-term " +
			"(preferences, technical stack, project details, personal facts). " +
			"Memories persist across all conversations.",
		Parameters: JSONSchema(map[string]interface{}{
			"key":   Prop("string", "Short label for the memory (2-5 words, e.g. 'preferred language', 'deployment target')"),
			"value": Prop("string", "The fact or detail to remember (1-2 sentences)"),
			"confidence": map[string]interface{}{
				"type":        "number",
				"description": "How confident you are in this fact (0.0-1.0, default 0.9)",
			},
		}, []string{"key", "value"}),
	}
}

func (t *memorySaveTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
	var args struct {
		Key        string  `json:"key"`
		Value      string  `json:"value"`
		Confidence float64 `json:"confidence"`
	}
	if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
		return "", fmt.Errorf("invalid arguments: %w", err)
	}

	if args.Key == "" || args.Value == "" {
		return "", fmt.Errorf("key and value are required")
	}

	confidence := args.Confidence
	if confidence <= 0 || confidence > 1.0 {
		confidence = 0.9
	}

	mem := &models.Memory{
		ID:         store.NewID(),
		Scope:      models.MemoryScopeUser,
		OwnerID:    execCtx.UserID,
		Key:        strings.TrimSpace(args.Key),
		Value:      strings.TrimSpace(args.Value),
		Confidence: confidence,
		Status:     models.MemoryStatusActive,
	}

	if execCtx.ChannelID != "" {
		mem.SourceChannelID = &execCtx.ChannelID
	}

	// If persona is active, use persona_user scope
	if execCtx.PersonaID != "" {
		mem.Scope = models.MemoryScopePersonaUser
		mem.OwnerID = execCtx.PersonaID
		mem.UserID = &execCtx.UserID
	}

	if err := t.stores.Memories.Upsert(ctx, mem); err != nil {
		return "", fmt.Errorf("save memory: %w", err)
	}

	// Embed for semantic recall
	t.embedMemory(ctx, mem, execCtx.UserID)

	return fmt.Sprintf("Remembered: %s = %s (confidence: %.0f%%)",
		args.Key, args.Value, confidence*100), nil
}

func (t *memorySaveTool) embedMemory(ctx context.Context, m *models.Memory, userID string) {
	if t.embedder == nil || !t.embedder.IsConfigured(ctx) {
		return
	}

	text := m.Key + ": " + m.Value

	var teamID *string
	if t.stores.Teams != nil {
		ids, _ := t.stores.Teams.GetUserTeamIDs(ctx, userID)
		if len(ids) > 0 {
			teamID = &ids[0]
		}
	}

	result, err := t.embedder.EmbedChunks(ctx, userID, teamID, []string{text})
	if err != nil || len(result.Vectors) == 0 {
		return
	}

	vecStr := vectorToString(result.Vectors[0])
	_ = t.stores.Memories.SetEmbedding(ctx, m.ID, vecStr)

	log.Printf("🧠 memory %s embedded (%d dims)", m.ID[:8], len(result.Vectors[0]))
}

// ═══════════════════════════════════════════
// memory_recall
// ═══════════════════════════════════════════

type memoryRecallTool struct {
	visitorDeniedBase
	stores   store.Stores
	embedder *knowledge.Embedder
}

func (t *memoryRecallTool) Definition() ToolDef {
	return ToolDef{
		Name:        "memory_recall",
		DisplayName: "Recall Memories",
		Category:    "memory",
		Description: "Search your memories about this user. Use this at the start of conversations " +
			"or when you need context about the user's preferences, projects, or history.",
		Parameters: JSONSchema(map[string]interface{}{
			"query": Prop("string", "Optional search query to filter memories (leave empty for all)"),
		}, nil),
	}
}

func (t *memoryRecallTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
	var args struct {
		Query string `json:"query"`
	}
	if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
		return "", fmt.Errorf("invalid arguments: %w", err)
	}

	var personaID *string
	if execCtx.PersonaID != "" {
		personaID = &execCtx.PersonaID
	}

	// Try hybrid recall if embedder available and query provided
	var memories []models.Memory
	var err error

	if t.embedder != nil && t.embedder.IsConfigured(ctx) && args.Query != "" {
		memories, err = t.recallWithEmbedding(ctx, execCtx.UserID, personaID, args.Query)
	}

	// Fall back to keyword recall
	if err != nil || len(memories) == 0 {
		memories, err = t.stores.Memories.Recall(ctx, execCtx.UserID, personaID, args.Query, 20)
	}

	if err != nil {
		return "", fmt.Errorf("recall memories: %w", err)
	}

	if len(memories) == 0 {
		suffix := ""
		if args.Query != "" {
			suffix = " matching '" + args.Query + "'"
		}
		return "No memories found" + suffix + ".", nil
	}

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Found %d memories:\n\n", len(memories)))
	for _, m := range memories {
		scope := m.Scope
		if scope == "persona_user" {
			scope = "persona"
		}
		sb.WriteString(fmt.Sprintf("• [%s] %s: %s (%.0f%% confidence)\n",
			scope, m.Key, m.Value, m.Confidence*100))
	}
	return sb.String(), nil
}

func (t *memoryRecallTool) recallWithEmbedding(ctx context.Context, userID string, personaID *string, query string) ([]models.Memory, error) {
	text := query
	if len(text) > 2000 {
		text = text[:2000]
	}

	var teamID *string
	if t.stores.Teams != nil {
		ids, _ := t.stores.Teams.GetUserTeamIDs(ctx, userID)
		if len(ids) > 0 {
			teamID = &ids[0]
		}
	}

	result, err := t.embedder.EmbedChunks(ctx, userID, teamID, []string{text})
	if err != nil || len(result.Vectors) == 0 {
		return nil, err
	}

	return t.stores.Memories.RecallHybrid(ctx, userID, personaID, query, result.Vectors[0], 20)
}