xinmi/internal/application/repository/message.go

package repository

import (
	"context"
	"slices"
	"strconv"
	"strings"
	"time"

	"gorm.io/gorm"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
)

// messageRepository implements the message repository interface
type messageRepository struct {
	db *gorm.DB
}

// NewMessageRepository creates a new message repository
func NewMessageRepository(db *gorm.DB) interfaces.MessageRepository {
	return &messageRepository{
		db: db,
	}
}

// CreateMessage creates a new message
func (r *messageRepository) CreateMessage(
	ctx context.Context, message *types.Message,
) (*types.Message, error) {
	if err := r.db.WithContext(ctx).Create(message).Error; err != nil {
		return nil, err
	}
	return message, nil
}

// GetMessage retrieves a message
func (r *messageRepository) GetMessage(
	ctx context.Context, sessionID string, messageID string,
) (*types.Message, error) {
	var message types.Message
	if err := r.db.WithContext(ctx).Where(
		"id = ? AND session_id = ?", messageID, sessionID,
	).First(&message).Error; err != nil {
		return nil, err
	}
	return &message, nil
}

// GetMessagesBySession retrieves all messages for a session with pagination
func (r *messageRepository) GetMessagesBySession(
	ctx context.Context, sessionID string, page int, pageSize int,
) ([]*types.Message, error) {
	var messages []*types.Message
	if err := r.db.WithContext(ctx).Where("session_id = ?", sessionID).Order("created_at ASC").
		Offset((page - 1) * pageSize).Limit(pageSize).Find(&messages).Error; err != nil {
		return nil, err
	}
	return messages, nil
}

// GetRecentMessagesBySession retrieves recent messages for a session
func (r *messageRepository) GetRecentMessagesBySession(
	ctx context.Context, sessionID string, limit int,
) ([]*types.Message, error) {
	var messages []*types.Message
	if err := r.db.WithContext(ctx).Where(
		"session_id = ?", sessionID,
	).Order("created_at DESC").Limit(limit).Find(&messages).Error; err != nil {
		if err == gorm.ErrRecordNotFound {
			return nil, nil
		}
		return nil, err
	}
	slices.SortFunc(messages, func(a, b *types.Message) int {
		cmp := a.CreatedAt.Compare(b.CreatedAt)
		if cmp == 0 {
			if a.Role == "user" { // User messages come first
				return -1
			}
			return 1 // Assistant messages come last
		}
		return cmp
	})
	return messages, nil
}

// GetMessagesBySessionBeforeTime retrieves messages from a session created before a specific time
func (r *messageRepository) GetMessagesBySessionBeforeTime(
	ctx context.Context, sessionID string, beforeTime time.Time, limit int,
) ([]*types.Message, error) {
	var messages []*types.Message
	if err := r.db.WithContext(ctx).Where(
		"session_id = ? AND created_at < ?", sessionID, beforeTime,
	).Order("created_at DESC").Limit(limit).Find(&messages).Error; err != nil {
		return nil, err
	}
	slices.SortFunc(messages, func(a, b *types.Message) int {
		cmp := a.CreatedAt.Compare(b.CreatedAt)
		if cmp == 0 {
			if a.Role == "user" { // User messages come first
				return -1
			}
			return 1 // Assistant messages come last
		}
		return cmp
	})
	return messages, nil
}

// ListMessagesBySessionAfterTime returns the oldest messages created after
// afterTime, so a caller holding a watermark can walk a session forward
// without skipping anything when it has more new messages than one page.
func (r *messageRepository) ListMessagesBySessionAfterTime(
	ctx context.Context, sessionID string, afterTime time.Time, limit int,
) ([]*types.Message, error) {
	var messages []*types.Message
	query := r.db.WithContext(ctx).Where("session_id = ?", sessionID)
	if !afterTime.IsZero() {
		query = query.Where("created_at > ?", afterTime)
	}
	if err := query.Order("created_at ASC, id ASC").Limit(limit).Find(&messages).Error; err != nil {
		return nil, err
	}
	// Memory extraction checkpoints by timestamp. Include the rest of the
	// boundary timestamp so advancing that checkpoint cannot skip tied rows.
	if limit > 0 && len(messages) == limit {
		last := messages[len(messages)-1]
		var tied []*types.Message
		if err := r.db.WithContext(ctx).Where("session_id = ? AND created_at = ? AND id > ?", sessionID, last.CreatedAt, last.ID).
			Order("id ASC").Find(&tied).Error; err != nil {
			return nil, err
		}
		messages = append(messages, tied...)
	}
	return messages, nil
}

func (r *messageRepository) MemoryOwnedSessionIDs(
	ctx context.Context, scope interfaces.MemoryScope, sessionIDs []string,
) (map[string]bool, error) {
	owned := map[string]bool{}
	if !scope.Valid() || len(sessionIDs) == 0 {
		return owned, nil
	}
	query := r.db.WithContext(ctx).Table("sessions").Select("sessions.id").
		Where("sessions.tenant_id = ? AND sessions.deleted_at IS NULL AND sessions.id IN ?", scope.TenantID, sessionIDs)
	parts := strings.SplitN(scope.SubjectID, ":", 2)
	if len(parts) != 2 || parts[1] == "" {
		return owned, nil
	}
	switch parts[0] {
	case types.PrincipalWebUser:
		query = query.Where("sessions.user_id = ?", parts[1]).Where(`EXISTS (
			SELECT 1 FROM users AS memory_user
			WHERE memory_user.id = sessions.user_id AND memory_user.is_active = ?
			AND memory_user.deleted_at IS NULL
			AND (memory_user.is_system_admin = ? OR EXISTS (
				SELECT 1 FROM tenant_members AS memory_member
				WHERE memory_member.user_id = memory_user.id AND memory_member.tenant_id = sessions.tenant_id
				AND memory_member.status = ? AND memory_member.deleted_at IS NULL)))`,
			true, true, types.TenantMemberStatusActive)
	case types.PrincipalAPIExternalUser, types.PrincipalEmbedSession:
		query = query.Where("sessions.user_id = ?", scope.SubjectID)
	case types.PrincipalIMUser:
		identity := strings.SplitN(parts[1], ":", 4)
		if len(identity) != 4 || identity[0] != strconv.FormatUint(scope.TenantID, 10) {
			return owned, nil
		}
		// Shared IM threads do not retain the author on individual messages.
		// They cannot safely become one participant's personal memory.
		query = query.Where(`EXISTS (SELECT 1 FROM im_channel_sessions AS im
			WHERE im.session_id = sessions.id AND im.tenant_id = ? AND im.im_channel_id = ?
			AND im.platform = ? AND im.user_id = ? AND im.thread_id = '' AND im.deleted_at IS NULL)`,
			scope.TenantID, identity[1], identity[2], identity[3])
	default:
		return owned, nil
	}
	var ids []string
	if err := query.Pluck("sessions.id", &ids).Error; err != nil {
		return nil, err
	}
	for _, id := range ids {
		owned[id] = true
	}
	return owned, nil
}

func (r *messageRepository) MemoryTurnStart(
	ctx context.Context, scope interfaces.MemoryScope, sessionID, messageID string,
) (time.Time, error) {
	owned, err := r.MemoryOwnedSessionIDs(ctx, scope, []string{sessionID})
	if err != nil {
		return time.Time{}, err
	}
	if !owned[sessionID] {
		return time.Time{}, gorm.ErrRecordNotFound
	}
	message, err := r.GetMessage(ctx, sessionID, messageID)
	if err != nil {
		return time.Time{}, err
	}
	if message.Role == "user" {
		return message.CreatedAt, nil
	}
	var input types.Message
	query := r.db.WithContext(ctx).Where("session_id = ? AND role = ?", sessionID, "user")
	if message.RequestID != "" {
		query = query.Where("request_id = ?", message.RequestID)
	} else {
		query = query.Where("created_at <= ?", message.CreatedAt)
	}
	if err := query.Order("created_at DESC, id DESC").First(&input).Error; err != nil {
		return time.Time{}, err
	}
	return input.CreatedAt, nil
}

// UpdateMessage updates an existing message
func (r *messageRepository) UpdateMessage(ctx context.Context, message *types.Message) error {
	return r.db.WithContext(ctx).Model(&types.Message{}).Where(
		"id = ? AND session_id = ?", message.ID, message.SessionID,
	).Updates(message).Error
}

// DeleteMessage deletes a message
func (r *messageRepository) DeleteMessage(ctx context.Context, sessionID string, messageID string) error {
	return r.db.WithContext(ctx).Where(
		"id = ? AND session_id = ?", messageID, sessionID,
	).Delete(&types.Message{}).Error
}

// GetFirstMessageOfUser retrieves the first message from a user in a session
func (r *messageRepository) GetFirstMessageOfUser(ctx context.Context, sessionID string) (*types.Message, error) {
	var message types.Message
	if err := r.db.WithContext(ctx).Where(
		"session_id = ? and role = ?", sessionID, "user",
	).Order("created_at ASC").First(&message).Error; err != nil {
		return nil, err
	}
	return &message, nil
}

// GetMessageByRequestID retrieves a message by request ID
func (r *messageRepository) GetMessageByRequestID(
	ctx context.Context, sessionID string, requestID string,
) (*types.Message, error) {
	var message types.Message

	result := r.db.WithContext(ctx).
		Where("session_id = ? AND request_id = ?", sessionID, requestID).
		First(&message)

	if result.Error != nil {
		if result.Error == gorm.ErrRecordNotFound {
			return nil, nil
		}
		return nil, result.Error
	}

	return &message, nil
}

// SearchMessagesByKeyword searches messages by keyword (ILIKE) across sessions for a tenant
func (r *messageRepository) SearchMessagesByKeyword(
	ctx context.Context, tenantID uint64, ownerID, keyword string, sessionIDs []string, limit int,
) ([]*types.MessageWithSession, error) {
	if limit <= 0 {
		limit = 20
	}

	var results []*types.MessageWithSession

	query := r.db.WithContext(ctx).
		Table("messages").
		Select("messages.*, sessions.title as session_title").
		Joins("INNER JOIN sessions ON sessions.id = messages.session_id AND sessions.deleted_at IS NULL").
		Where("sessions.tenant_id = ?", tenantID).
		Where("messages.deleted_at IS NULL").
		Where("messages.content ILIKE ?", "%"+escapeLikeKeyword(keyword)+"%")

	// Matches the scoping used when listing sessions, including the legacy
	// allowance for tenant-level sessions created before per-user ownership.
	if ownerID != "" {
		query = query.Where("(sessions.user_id = ? OR sessions.user_id IS NULL OR sessions.user_id = '')", ownerID)
	}

	if len(sessionIDs) > 0 {
		query = query.Where("messages.session_id IN ?", sessionIDs)
	}

	if err := query.Order("messages.created_at DESC").Limit(limit).Find(&results).Error; err != nil {
		return nil, err
	}

	return results, nil
}

// OwnedSessionIDs narrows a set of session ids to the ones this person owns.
//
// The vector path finds messages through a shared knowledge base that has no
// notion of who wrote them, so ownership has to be re-established here before
// anything is returned.
func (r *messageRepository) OwnedSessionIDs(
	ctx context.Context, tenantID uint64, ownerID string, sessionIDs []string,
) (map[string]bool, error) {
	owned := make(map[string]bool, len(sessionIDs))
	if len(sessionIDs) == 0 {
		return owned, nil
	}
	var ids []string
	query := r.db.WithContext(ctx).
		Table("sessions").
		Select("id").
		Where("tenant_id = ?", tenantID).
		Where("deleted_at IS NULL").
		Where("id IN ?", sessionIDs)
	if ownerID != "" {
		query = query.Where("(user_id = ? OR user_id IS NULL OR user_id = '')", ownerID)
	}
	if err := query.Pluck("id", &ids).Error; err != nil {
		return nil, err
	}
	for _, id := range ids {
		owned[id] = true
	}
	return owned, nil
}

// GetMessagesByKnowledgeIDs retrieves messages by their associated Knowledge IDs
func (r *messageRepository) GetMessagesByKnowledgeIDs(
	ctx context.Context, knowledgeIDs []string,
) ([]*types.MessageWithSession, error) {
	if len(knowledgeIDs) == 0 {
		return nil, nil
	}
	var results []*types.MessageWithSession
	if err := r.db.WithContext(ctx).
		Table("messages").
		Select("messages.*, sessions.title as session_title").
		Joins("INNER JOIN sessions ON sessions.id = messages.session_id AND sessions.deleted_at IS NULL").
		Where("messages.deleted_at IS NULL").
		Where("messages.knowledge_id IN ?", knowledgeIDs).
		Find(&results).Error; err != nil {
		return nil, err
	}
	return results, nil
}

// GetMessagesByRequestIDs retrieves messages by their request IDs (used to fetch Q&A pair partners)
func (r *messageRepository) GetMessagesByRequestIDs(
	ctx context.Context, requestIDs []string,
) ([]*types.MessageWithSession, error) {
	if len(requestIDs) == 0 {
		return nil, nil
	}
	var results []*types.MessageWithSession
	if err := r.db.WithContext(ctx).
		Table("messages").
		Select("messages.*, sessions.title as session_title").
		Joins("INNER JOIN sessions ON sessions.id = messages.session_id AND sessions.deleted_at IS NULL").
		Where("messages.deleted_at IS NULL").
		Where("messages.request_id IN ?", requestIDs).
		Find(&results).Error; err != nil {
		return nil, err
	}
	return results, nil
}

// GetKnowledgeIDsBySessionID retrieves all knowledge IDs for messages in a session
func (r *messageRepository) GetKnowledgeIDsBySessionID(
	ctx context.Context, sessionID string,
) ([]string, error) {
	var knowledgeIDs []string
	if err := r.db.WithContext(ctx).
		Model(&types.Message{}).
		Where("session_id = ? AND knowledge_id != '' AND knowledge_id IS NOT NULL AND deleted_at IS NULL", sessionID).
		Pluck("knowledge_id", &knowledgeIDs).Error; err != nil {
		return nil, err
	}
	return knowledgeIDs, nil
}

// UpdateMessageImages updates only the images JSONB column for a message.
// Uses Select to force GORM to include the column even when struct-based
// Updates would otherwise skip custom Valuer types.
func (r *messageRepository) UpdateMessageImages(ctx context.Context, sessionID, messageID string, images types.MessageImages) error {
	return r.db.WithContext(ctx).
		Model(&types.Message{}).
		Where("id = ? AND session_id = ?", messageID, sessionID).
		Update("images", images).Error
}

// UpdateMessageRenderedContent updates only the rendered_content column for a message.
func (r *messageRepository) UpdateMessageRenderedContent(ctx context.Context, sessionID, messageID string, renderedContent string) error {
	return r.db.WithContext(ctx).
		Model(&types.Message{}).
		Where("id = ? AND session_id = ?", messageID, sessionID).
		Update("rendered_content", renderedContent).Error
}

// DeleteMessagesBySessionID deletes all messages belonging to a session (soft delete)
func (r *messageRepository) DeleteMessagesBySessionID(ctx context.Context, sessionID string) error {
	return r.db.WithContext(ctx).Where("session_id = ?", sessionID).Delete(&types.Message{}).Error
}

// UpdateMessageKnowledgeID updates the knowledge_id field for a message
func (r *messageRepository) UpdateMessageKnowledgeID(
	ctx context.Context, messageID string, knowledgeID string,
) error {
	return r.db.WithContext(ctx).
		Model(&types.Message{}).
		Where("id = ?", messageID).
		Update("knowledge_id", knowledgeID).Error
}

// GetSessionArtifacts returns every skill-produced MessageArtifact recorded
// against any assistant message of the session, in creation order.
//
// Projection is scoped to the artifacts JSONB column plus created_at (used
// to order the flattened output). Assistant messages without artifacts (the
// common case) contribute an empty slice and cost nothing extra.
func (r *messageRepository) GetSessionArtifacts(
	ctx context.Context, sessionID string,
) (types.MessageArtifacts, error) {
	if sessionID == "" {
		return nil, nil
	}
	var rows []struct {
		Artifacts types.MessageArtifacts `gorm:"column:artifacts"`
		CreatedAt time.Time              `gorm:"column:created_at"`
	}
	if err := r.db.WithContext(ctx).
		Model(&types.Message{}).
		Select("artifacts", "created_at").
		Where("session_id = ? AND deleted_at IS NULL", sessionID).
		Order("created_at ASC").
		Find(&rows).Error; err != nil {
		return nil, err
	}
	if len(rows) == 0 {
		return types.MessageArtifacts{}, nil
	}
	result := make(types.MessageArtifacts, 0, len(rows))
	for _, row := range rows {
		if len(row.Artifacts) == 0 {
			continue
		}
		result = append(result, row.Artifacts...)
	}
	return result, nil
}

// GetSessionAttachments returns every user-uploaded attachment in creation
// order while projecting only the attachments JSON column.
func (r *messageRepository) GetSessionAttachments(
	ctx context.Context, sessionID string,
) (types.MessageAttachments, error) {
	if sessionID == "" {
		return nil, nil
	}
	var rows []struct {
		Attachments types.MessageAttachments `gorm:"column:attachments"`
		CreatedAt   time.Time                `gorm:"column:created_at"`
	}
	if err := r.db.WithContext(ctx).
		Model(&types.Message{}).
		Select("attachments", "created_at").
		Where("session_id = ? AND deleted_at IS NULL", sessionID).
		Order("created_at ASC").
		Find(&rows).Error; err != nil {
		return nil, err
	}
	result := make(types.MessageAttachments, 0, len(rows))
	for _, row := range rows {
		result = append(result, row.Attachments...)
	}
	return result, nil
}