tamabox/internal/db/message.go

package db

import (
	"context"
	"fmt"

	"github.com/pkg/errors"
	"gorm.io/gorm"

	"github.com/tamakyi/TamaBox/internal/conf"
	"github.com/tamakyi/TamaBox/internal/dbutil"
)

// minChatMessagesPerPage 是会话消息单次加载的兜底条数。
//
// 正常情况下取 conf.Limits.ChatMaxMessages（默认 200，即一次把整个会话读满）。
// 这个常量只在「配置未初始化」（零值，测试里直接构造 store 时会遇到）时生效，
// 不会覆盖箱主显式配置的更大或更小的值。
const minChatMessagesPerPage = 50

var Messages MessageStore

type MessageStore interface {
	Create(ctx context.Context, opts CreateMessageOptions) (*Message, error)
	GetByConversationID(ctx context.Context, conversationID uint, cursor *dbutil.Cursor) ([]*Message, error)
	// GetFirstMessagesByConversationIDs 批量获取多个会话的第一条消息（解决 N+1 查询）
	GetFirstMessagesByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint]string, error)
	// GetAllByConversationIDs 批量获取多个会话的全部消息（导出用，不分页），按 conversation_id 分组。
	// 将每个会话原本的 1 次查询合并为 1 次 IN 查询，避免 N+1。
	GetAllByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint][]*Message, error)
}

func NewMessageStore(db *gorm.DB) MessageStore {
	return &messages{db}
}

type messages struct {
	*gorm.DB
}

type Message struct {
	dbutil.Model
	ConversationID uint   `gorm:"index:idx_message_conversation" json:"conversation_id"`
	SenderID       uint   `json:"sender_id"`
	Content        string `gorm:"type:text" json:"content"`
	Read           bool   `gorm:"default:false" json:"read"`
}

type CreateMessageOptions struct {
	ConversationID uint
	SenderID       uint
	Content        string
}

func (db *messages) Create(ctx context.Context, opts CreateMessageOptions) (*Message, error) {
	msg := &Message{
		ConversationID: opts.ConversationID,
		SenderID:       opts.SenderID,
		Content:        opts.Content,
	}
	if err := db.WithContext(ctx).Create(msg).Error; err != nil {
		return nil, errors.Wrap(err, "create message")
	}
	return msg, nil
}

func (db *messages) GetByConversationID(ctx context.Context, conversationID uint, cursor *dbutil.Cursor) ([]*Message, error) {
	var messages []*Message
	q := db.WithContext(ctx).Where("conversation_id = ?", conversationID)

	// 默认一次取满整个会话，而不是固定 50 条。
	//
	// 原先写死 limit=50 且按 created_at ASC 取，取到的是「最早的 50 条」；
	// 而会话上限是 conf.Limits.ChatMaxMessages（默认 200）。结果是第 51 条起
	// 写库成功、message_count 在涨、页面显示 51/200，但列表里永远没有它 ——
	// 用户无论等多久、换设备、重新登录都看不到（实测 55 条时只能渲染 #1–#50）。
	limit := conf.Limits.ChatMaxMessages
	if limit <= 0 {
		// 仅兜底「配置未初始化」（零值）这一种情况。
		// 不能写成 limit < 50 → 50：那样会把箱主显式配的小上限抬大，
		// 属于替用户改配置，而不是补默认值。
		limit = minChatMessagesPerPage
	}
	if cursor != nil {
		if cursor.Value != nil && fmt.Sprintf("%v", cursor.Value) != "" {
			// 游标语义是「加载更早的消息」，与下面的 ORDER BY id DESC 配套；
			// 若沿用旧的 ASC 排序，id < cursor 每次都会命中同一批最老的消息。
			q = q.Where("id < ?", cursor.Value)
		}
		limit = cursor.Limit()
	}

	// 按 id 倒序取最新 limit 条，再反转成时间正序交给视图层
	// （气泡从上到下 = 旧 → 新）。
	//
	// 排序键用 id 而非 created_at：同一毫秒内写入的多条消息 created_at 完全相同，
	// 排序不稳定会导致同一页里顺序来回变；id 自增与写入顺序严格一致。
	q = q.Order("id DESC")
	q = q.Limit(limit)

	if err := q.Find(&messages).Error; err != nil {
		return nil, errors.Wrap(err, "get messages")
	}
	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
}

// GetFirstMessagesByConversationIDs 批量获取多个会话的第一条消息内容。
// 使用窗口函数 ROW_NUMBER() 选出每个会话最早的消息，避免 N+1 查询。
func (db *messages) GetFirstMessagesByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint]string, error) {
	if len(conversationIDs) == 0 {
		return make(map[uint]string), nil
	}

	// 使用子查询 + ROW_NUMBER 窗口函数获取每个 conversation 的第一条消息
	// 兼容 MySQL 8+ / PostgreSQL / SQLite 3.25+
	var results []struct {
		ConversationID uint   `gorm:"column:conversation_id"`
		Content        string `gorm:"column:content"`
	}

	err := db.WithContext(ctx).Raw(`
		SELECT conversation_id, content FROM (
			SELECT conversation_id, content,
				ROW_NUMBER() OVER (PARTITION BY conversation_id ORDER BY created_at ASC) AS rn
			FROM messages
			WHERE conversation_id IN ?
		) sub WHERE rn = 1
	`, conversationIDs).Scan(&results).Error

	if err != nil {
		return nil, errors.Wrap(err, "get first messages batch")
	}

	resultMap := make(map[uint]string, len(results))
	for _, r := range results {
		resultMap[r.ConversationID] = r.Content
	}
	return resultMap, nil
}

// GetAllByConversationIDs 批量获取多个会话的全部消息（导出用，不分页），按 conversation_id 分组。
// 用于账号导出：将每个会话原本的 1 次查询合并为 1 次 IN 查询，避免 N+1。
func (db *messages) GetAllByConversationIDs(ctx context.Context, conversationIDs []uint) (map[uint][]*Message, error) {
	result := make(map[uint][]*Message)
	if len(conversationIDs) == 0 {
		return result, nil
	}
	var items []*Message
	if err := db.WithContext(ctx).
		Where("conversation_id IN ?", conversationIDs).
		Order("conversation_id ASC, created_at ASC").
		Find(&items).Error; err != nil {
		return nil, errors.Wrap(err, "get messages by conversation ids")
	}
	for _, m := range items {
		result[m.ConversationID] = append(result[m.ConversationID], m)
	}
	return result, nil
}