yao/agent/store/xun/message.go

package xun

import (
	"fmt"
	"time"

	"github.com/google/uuid"
	jsoniter "github.com/json-iterator/go"
	"github.com/yaoapp/yao/agent/store/types"
)

// =============================================================================
// Message Management
// =============================================================================

// SaveMessages batch saves messages for a chat using a single database call
// This is the primary write method - messages are buffered during execution
// and batch-written at the end of a request
func (store *Xun) SaveMessages(chatID string, messages []*types.Message) error {
	if chatID == "" {
		return fmt.Errorf("chat_id is required")
	}
	if len(messages) == 0 {
		return nil // Nothing to save
	}

	// Prepare batch insert data
	now := time.Now()
	rows := make([]map[string]interface{}, 0, len(messages))

	for _, msg := range messages {
		if msg == nil {
			continue
		}

		// Generate message_id if not provided
		messageID := msg.MessageID
		if messageID == "" {
			messageID = uuid.New().String()
		}

		// Validate required fields
		if msg.Role == "" {
			return fmt.Errorf("message role is required")
		}
		if msg.Type == "" {
			return fmt.Errorf("message type is required")
		}
		if msg.Props == nil {
			return fmt.Errorf("message props is required")
		}

		// Serialize JSON fields
		propsJSON, err := jsoniter.MarshalToString(msg.Props)
		if err != nil {
			return fmt.Errorf("failed to marshal props: %w", err)
		}

		// Build row with all fields (including nullable ones for consistent batch insert)
		row := map[string]interface{}{
			"message_id":   messageID,
			"chat_id":      chatID,
			"role":         msg.Role,
			"type":         msg.Type,
			"props":        propsJSON,
			"sequence":     msg.Sequence,
			"request_id":   nil,
			"block_id":     nil,
			"thread_id":    nil,
			"assistant_id": nil,
			"connector":    nil,
			"mode":         nil,
			"metadata":     nil,
			"created_at":   now,
			"updated_at":   now,
		}

		// Set nullable fields if they have values
		if msg.RequestID != "" {
			row["request_id"] = msg.RequestID
		}
		if msg.BlockID != "" {
			row["block_id"] = msg.BlockID
		}
		if msg.ThreadID != "" {
			row["thread_id"] = msg.ThreadID
		}
		if msg.AssistantID != "" {
			row["assistant_id"] = msg.AssistantID
		}
		if msg.Connector != "" {
			row["connector"] = msg.Connector
		}
		if msg.Mode != "" {
			row["mode"] = msg.Mode
		}
		if msg.Metadata != nil {
			metadataJSON, err := jsoniter.MarshalToString(msg.Metadata)
			if err != nil {
				return fmt.Errorf("failed to marshal metadata: %w", err)
			}
			row["metadata"] = metadataJSON
		}

		rows = append(rows, row)
	}

	if len(rows) == 0 {
		return nil
	}

	// Single batch insert - one database call for all messages
	return store.newQueryMessage().Insert(rows)
}

// GetMessages retrieves messages for a chat with filtering
func (store *Xun) GetMessages(chatID string, filter types.MessageFilter) ([]*types.Message, error) {
	if chatID == "" {
		return nil, fmt.Errorf("chat_id is required")
	}

	qb := store.newQueryMessage().
		Where("chat_id", chatID).
		WhereNull("deleted_at")

	// Apply filters
	if filter.RequestID != "" {
		qb.Where("request_id", filter.RequestID)
	}
	if filter.Role != "" {
		qb.Where("role", filter.Role)
	}
	if filter.BlockID != "" {
		qb.Where("block_id", filter.BlockID)
	}
	if filter.ThreadID != "" {
		qb.Where("thread_id", filter.ThreadID)
	}
	if filter.Type != "" {
		qb.Where("type", filter.Type)
	}

	// When Limit is specified WITHOUT Offset, we want the N most-recent
	// messages.  Strategy: query DESC to get the latest rows, then reverse
	// the slice so the caller receives them in chronological (ASC) order.
	// When Offset is also present, the caller is doing forward pagination,
	// so we keep ASC order and apply Limit+Offset normally.
	needReverse := false
	if filter.Limit > 0 && filter.Offset <= 0 {
		qb.Limit(filter.Limit)
		qb.OrderBy("id", "desc")
		needReverse = true
	} else if filter.Limit > 0 && filter.Offset > 0 {
		qb.Limit(filter.Limit).Offset(filter.Offset)
		qb.OrderBy("id", "asc")
	} else {
		if filter.Offset > 0 {
			qb.Limit(1000000).Offset(filter.Offset)
		}
		qb.OrderBy("id", "asc")
	}

	rows, err := qb.Get()
	if err != nil {
		return nil, err
	}

	messages := make([]*types.Message, 0, len(rows))
	for _, row := range rows {
		data := row.ToMap()
		if data == nil || data["message_id"] == nil {
			continue
		}

		msg, err := store.rowToMessage(data)
		if err != nil {
			continue
		}
		messages = append(messages, msg)
	}

	if needReverse {
		for i, j := 0, len(messages)-1; i < j; i, j = i+1, j-1 {
			messages[i], messages[j] = messages[j], messages[i]
		}
	}

	return messages, nil
}

// UpdateMessage updates a single message
func (store *Xun) UpdateMessage(messageID string, updates map[string]interface{}) error {
	if messageID == "" {
		return fmt.Errorf("message_id is required")
	}
	if len(updates) == 0 {
		return fmt.Errorf("no fields to update")
	}

	// Check if message exists
	exists, err := store.newQueryMessage().
		Where("message_id", messageID).
		WhereNull("deleted_at").
		Exists()
	if err != nil {
		return err
	}
	if !exists {
		return fmt.Errorf("message %s not found", messageID)
	}

	// Prepare update data
	data := make(map[string]interface{})

	for key, value := range updates {
		// Skip system fields
		if key == "message_id" || key == "chat_id" || key == "created_at" {
			continue
		}

		// Handle JSON fields
		if key == "props" || key == "metadata" {
			if value != nil {
				jsonStr, err := jsoniter.MarshalToString(value)
				if err != nil {
					return fmt.Errorf("failed to marshal %s: %w", key, err)
				}
				data[key] = jsonStr
			} else {
				data[key] = nil
			}
			continue
		}

		data[key] = value
	}

	// Always update updated_at
	data["updated_at"] = time.Now()

	if len(data) == 0 {
		return fmt.Errorf("no valid fields to update")
	}

	_, err = store.newQueryMessage().
		Where("message_id", messageID).
		Update(data)

	return err
}

// DeleteMessages soft deletes specific messages from a chat
func (store *Xun) DeleteMessages(chatID string, messageIDs []string) error {
	if chatID == "" {
		return fmt.Errorf("chat_id is required")
	}
	if len(messageIDs) == 0 {
		return nil // Nothing to delete
	}

	// Soft delete all specified messages in one query
	_, err := store.newQueryMessage().
		Where("chat_id", chatID).
		WhereIn("message_id", messageIDs).
		WhereNull("deleted_at").
		Update(map[string]interface{}{
			"deleted_at": time.Now(),
			"updated_at": time.Now(),
		})

	return err
}

// GetMessageByID retrieves a single message by ID
func (store *Xun) GetMessageByID(messageID string) (*types.Message, error) {
	if messageID == "" {
		return nil, fmt.Errorf("message_id is required")
	}

	row, err := store.newQueryMessage().
		Where("message_id", messageID).
		WhereNull("deleted_at").
		First()
	if err != nil {
		return nil, err
	}

	if row == nil {
		return nil, fmt.Errorf("message %s not found", messageID)
	}

	data := row.ToMap()
	if len(data) == 0 || data["message_id"] == nil {
		return nil, fmt.Errorf("message %s not found", messageID)
	}

	return store.rowToMessage(data)
}

// GetMessageCount returns the count of messages for a chat
func (store *Xun) GetMessageCount(chatID string) (int64, error) {
	if chatID == "" {
		return 0, fmt.Errorf("chat_id is required")
	}

	return store.newQueryMessage().
		Where("chat_id", chatID).
		WhereNull("deleted_at").
		Count()
}

// GetLastSequence returns the last sequence number for a chat
func (store *Xun) GetLastSequence(chatID string) (int, error) {
	if chatID == "" {
		return 0, fmt.Errorf("chat_id is required")
	}

	row, err := store.newQueryMessage().
		Where("chat_id", chatID).
		WhereNull("deleted_at").
		OrderBy("sequence", "desc").
		First()
	if err != nil {
		return 0, err
	}

	if row == nil {
		return 0, nil
	}

	data := row.ToMap()
	return getInt(data, "sequence"), nil
}

// =============================================================================
// Helper Functions
// =============================================================================

// rowToMessage converts a database row to a Message struct
func (store *Xun) rowToMessage(data map[string]interface{}) (*types.Message, error) {
	msg := &types.Message{
		MessageID:   getString(data, "message_id"),
		ChatID:      getString(data, "chat_id"),
		RequestID:   getString(data, "request_id"),
		Role:        getString(data, "role"),
		Type:        getString(data, "type"),
		BlockID:     getString(data, "block_id"),
		ThreadID:    getString(data, "thread_id"),
		AssistantID: getString(data, "assistant_id"),
		Connector:   getString(data, "connector"),
		Mode:        getString(data, "mode"),
		Sequence:    getInt(data, "sequence"),
	}

	// Handle timestamps
	if createdAt := getTime(data, "created_at"); createdAt != nil {
		msg.CreatedAt = *createdAt
	}
	if updatedAt := getTime(data, "updated_at"); updatedAt != nil {
		msg.UpdatedAt = *updatedAt
	}

	// Handle props (required)
	if props := data["props"]; props != nil {
		if propsStr, ok := props.(string); ok && propsStr != "" {
			var propsMap map[string]interface{}
			if err := jsoniter.UnmarshalFromString(propsStr, &propsMap); err == nil {
				msg.Props = propsMap
			}
		} else if propsMap, ok := props.(map[string]interface{}); ok {
			msg.Props = propsMap
		}
	}

	// Handle metadata (optional)
	if metadata := data["metadata"]; metadata != nil {
		if metaStr, ok := metadata.(string); ok && metaStr != "" {
			var metaMap map[string]interface{}
			if err := jsoniter.UnmarshalFromString(metaStr, &metaMap); err == nil {
				msg.Metadata = metaMap
			}
		} else if metaMap, ok := metadata.(map[string]interface{}); ok {
			msg.Metadata = metaMap
		}
	}

	return msg, nil
}