Agent-System/internal/ctxmgr/chat_unread.go

package ctxmgr

import (
	"context"
	"fmt"
	"sort"
	"strconv"
	"time"
)

// ChatReadPosition 双轨制未读标记位点
type ChatReadPosition struct {
	SessionID          string    `json:"session_id"`
	UserID             int64     `json:"user_id"`
	LastReadMessageID  int64     `json:"last_read_message_id"` // 连续已读基线
	DismissedUnreadIDs []int64   `json:"dismissed_unread_ids"` // 离散跳读/倒序消费的消息 ID
	UpdatedAt          time.Time `json:"updated_at"`
}

// UnreadDayItem 某日的未读消息聚合项
type UnreadDayItem struct {
	Date          string   `json:"date"`
	Count         int      `json:"count"`
	LastMessageID int64    `json:"last_message_id"`
	Senders       []string `json:"senders,omitempty"`
}

// ChatUnreadSummary 未读消息汇总结构体
type ChatUnreadSummary struct {
	Total            int             `json:"total"`
	NextMessageID    int64           `json:"next_message_id"`
	UnreadMessageIDs []int64         `json:"unread_message_ids"`
	Items            []UnreadDayItem `json:"items"`
}

// MarkChatReadRequest 标记已读请求
type MarkChatReadRequest struct {
	SessionID      string  `json:"session_id"`
	UserID         int64   `json:"user_id"`
	ReadMessageIDs []int64 `json:"read_message_ids,omitempty"`
	AdvanceTo      int64   `json:"advance_to,omitempty"`
}

// GetChatReadPosition 获取会话与用户的已读位点
func (m *ContextManager) GetChatReadPosition(ctx context.Context, sessionID string, userID int64) (*ChatReadPosition, error) {
	if m.pool != nil {
		row := m.pool.QueryRow(ctx, `
			SELECT session_id, user_id, last_read_message_id, dismissed_unread_ids, updated_at
			FROM chat_read_positions
			WHERE session_id = $1 AND user_id = $2
		`, sessionID, userID)

		var pos ChatReadPosition
		if err := row.Scan(&pos.SessionID, &pos.UserID, &pos.LastReadMessageID, &pos.DismissedUnreadIDs, &pos.UpdatedAt); err != nil {
			return nil, nil
		}
		return &pos, nil
	}

	m.mu.RLock()
	defer m.mu.RUnlock()

	key := sessionID + ":" + strconv.FormatInt(userID, 10)
	if pos, ok := m.readPositions[key]; ok {
		cp := *pos
		cp.DismissedUnreadIDs = append([]int64(nil), pos.DismissedUnreadIDs...)
		return &cp, nil
	}
	return nil, nil
}

// GetChatUnreadSummary 计算当前用户在指定会话的未读汇总
func (m *ContextManager) GetChatUnreadSummary(ctx context.Context, sessionID string, userID int64, sinceMessageID int64) (*ChatUnreadSummary, error) {
	pos, err := m.GetChatReadPosition(ctx, sessionID, userID)
	if err != nil {
		return nil, fmt.Errorf("get read position failed: %w", err)
	}

	// 取出该会话所有有效消息 (ID 正序)
	msgs, err := m.GetChatHistory(ctx, sessionID, 5000, 0)
	if err != nil {
		return nil, fmt.Errorf("get chat history failed: %w", err)
	}

	// 若尚未有位点记录：
	// 若外部传入了 sinceMessageID > 0，以该 sinceMessageID 作为初始基线；
	// 否则作为初次打开会话，默认视作全部已读（基线初始化至最新消息 ID）
	if pos == nil {
		if sinceMessageID > 0 {
			pos = &ChatReadPosition{
				SessionID:          sessionID,
				UserID:             userID,
				LastReadMessageID:  sinceMessageID,
				DismissedUnreadIDs: []int64{},
				UpdatedAt:          time.Now(),
			}
			_, _ = m.MarkChatRead(ctx, MarkChatReadRequest{
				SessionID: sessionID,
				UserID:    userID,
				AdvanceTo: sinceMessageID,
			})
		} else {
			var maxID int64 = 0
			for _, msg := range msgs {
				if msg.ID > maxID {
					maxID = msg.ID
				}
			}
			if maxID > 0 {
				_, _ = m.MarkChatRead(ctx, MarkChatReadRequest{
					SessionID: sessionID,
					UserID:    userID,
					AdvanceTo: maxID,
				})
			}
			return &ChatUnreadSummary{
				Total:            0,
				NextMessageID:    0,
				UnreadMessageIDs: []int64{},
				Items:            []UnreadDayItem{},
			}, nil
		}
	}

	effectiveBaseline := pos.LastReadMessageID
	if sinceMessageID > effectiveBaseline {
		effectiveBaseline = sinceMessageID
	}

	dismissedMap := make(map[int64]bool, len(pos.DismissedUnreadIDs))
	for _, id := range pos.DismissedUnreadIDs {
		dismissedMap[id] = true
	}

	type dayAgg struct {
		date          string
		count         int
		lastMessageID int64
		senderSet     map[string]bool
	}
	dayMap := make(map[string]*dayAgg)
	var dayOrder []string
	var unreadIDs []int64

	for _, msg := range msgs {
		if msg.DeletedAt != nil {
			continue
		}
		if msg.ID <= effectiveBaseline {
			continue
		}
		// 自己发送的 user 消息不计入未读
		if msg.UserID == userID && msg.Role == "user" {
			continue
		}
		// 已离散消费的消息不计入未读
		if dismissedMap[msg.ID] {
			continue
		}

		unreadIDs = append(unreadIDs, msg.ID)

		dayKey := msg.DayKey
		if dayKey == "" && !msg.CreatedAt.IsZero() {
			dayKey = msg.CreatedAt.Format("2006-01-02")
		}
		if dayKey == "" {
			dayKey = "unknown"
		}

		agg, ok := dayMap[dayKey]
		if !ok {
			agg = &dayAgg{
				date:      dayKey,
				senderSet: make(map[string]bool),
			}
			dayMap[dayKey] = agg
			dayOrder = append(dayOrder, dayKey)
		}
		agg.count++
		if msg.ID > agg.lastMessageID {
			agg.lastMessageID = msg.ID
		}
		if msg.Role != "" {
			agg.senderSet[msg.Role] = true
		}
	}

	// 倒序排序未读 ID，最新未读排在前面
	sort.Slice(unreadIDs, func(i, j int) bool {
		return unreadIDs[i] > unreadIDs[j]
	})

	var nextID int64 = 0
	if len(unreadIDs) > 0 {
		nextID = unreadIDs[0] // 离最新位置最近的一条未读
	}

	items := make([]UnreadDayItem, 0, len(dayOrder))
	for _, dk := range dayOrder {
		agg := dayMap[dk]
		var senders []string
		for s := range agg.senderSet {
			senders = append(senders, s)
		}
		sort.Strings(senders)
		items = append(items, UnreadDayItem{
			Date:          agg.date,
			Count:         agg.count,
			LastMessageID: agg.lastMessageID,
			Senders:       senders,
		})
	}

	return &ChatUnreadSummary{
		Total:            len(unreadIDs),
		NextMessageID:    nextID,
		UnreadMessageIDs: unreadIDs,
		Items:            items,
	}, nil
}

// MarkChatRead 推进已读位点或添加离散消费消息 ID，并执行自动压缩
func (m *ContextManager) MarkChatRead(ctx context.Context, req MarkChatReadRequest) (*ChatReadPosition, error) {
	if req.SessionID == "" {
		return nil, fmt.Errorf("session_id required")
	}

	if m.pool != nil {
		return m.markChatReadPG(ctx, req)
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	key := req.SessionID + ":" + strconv.FormatInt(req.UserID, 10)
	pos, ok := m.readPositions[key]
	if !ok {
		pos = &ChatReadPosition{
			SessionID:          req.SessionID,
			UserID:             req.UserID,
			LastReadMessageID:  0,
			DismissedUnreadIDs: []int64{},
			UpdatedAt:          time.Now(),
		}
		m.readPositions[key] = pos
	}

	if req.AdvanceTo > pos.LastReadMessageID {
		pos.LastReadMessageID = req.AdvanceTo
	}

	for _, id := range req.ReadMessageIDs {
		if id > pos.LastReadMessageID {
			pos.DismissedUnreadIDs = append(pos.DismissedUnreadIDs, id)
		}
	}

	// 获取此会话中大于 baseline 的现有有效消息 ID 列表
	var futureIDs []int64
	for _, msg := range m.messages[req.SessionID] {
		if msg.DeletedAt == nil && msg.ID > pos.LastReadMessageID {
			futureIDs = append(futureIDs, msg.ID)
		}
	}
	sort.Slice(futureIDs, func(i, j int) bool { return futureIDs[i] < futureIDs[j] })

	// 执行自动压缩
	pos.LastReadMessageID, pos.DismissedUnreadIDs = compressDismissed(pos.LastReadMessageID, pos.DismissedUnreadIDs, futureIDs)
	pos.UpdatedAt = time.Now()

	cp := *pos
	cp.DismissedUnreadIDs = append([]int64(nil), pos.DismissedUnreadIDs...)
	return &cp, nil
}

func (m *ContextManager) markChatReadPG(ctx context.Context, req MarkChatReadRequest) (*ChatReadPosition, error) {
	tx, err := m.pool.Begin(ctx)
	if err != nil {
		return nil, fmt.Errorf("begin tx failed: %w", err)
	}
	defer tx.Rollback(ctx)

	var pos ChatReadPosition
	row := tx.QueryRow(ctx, `
		SELECT session_id, user_id, last_read_message_id, dismissed_unread_ids, updated_at
		FROM chat_read_positions
		WHERE session_id = $1 AND user_id = $2
		FOR UPDATE
	`, req.SessionID, req.UserID)

	if err := row.Scan(&pos.SessionID, &pos.UserID, &pos.LastReadMessageID, &pos.DismissedUnreadIDs, &pos.UpdatedAt); err != nil {
		pos = ChatReadPosition{
			SessionID:          req.SessionID,
			UserID:             req.UserID,
			LastReadMessageID:  0,
			DismissedUnreadIDs: []int64{},
		}
	}

	if req.AdvanceTo > pos.LastReadMessageID {
		pos.LastReadMessageID = req.AdvanceTo
	}

	for _, id := range req.ReadMessageIDs {
		if id > pos.LastReadMessageID {
			pos.DismissedUnreadIDs = append(pos.DismissedUnreadIDs, id)
		}
	}

	// 查询 DB 中大于当前 baseline 的消息 ID
	rows, err := tx.Query(ctx, `
		SELECT id FROM chat_messages
		WHERE session_id = $1 AND deleted_at IS NULL AND id > $2
		ORDER BY id ASC
	`, req.SessionID, pos.LastReadMessageID)
	var futureIDs []int64
	if err == nil {
		for rows.Next() {
			var mid int64
			if err := rows.Scan(&mid); err == nil {
				futureIDs = append(futureIDs, mid)
			}
		}
		if err := rows.Err(); err != nil {
			rows.Close()
			return nil, fmt.Errorf("scan future unread ids failed: %w", err)
		}
		rows.Close()
	}

	pos.LastReadMessageID, pos.DismissedUnreadIDs = compressDismissed(pos.LastReadMessageID, pos.DismissedUnreadIDs, futureIDs)
	pos.UpdatedAt = time.Now()

	if pos.DismissedUnreadIDs == nil {
		pos.DismissedUnreadIDs = []int64{}
	}

	_, err = tx.Exec(ctx, `
		INSERT INTO chat_read_positions (session_id, user_id, last_read_message_id, dismissed_unread_ids, updated_at)
		VALUES ($1, $2, $3, COALESCE($4, '{}'::bigint[]), $5)
		ON CONFLICT (session_id, user_id)
		DO UPDATE SET
			last_read_message_id = EXCLUDED.last_read_message_id,
			dismissed_unread_ids = COALESCE(EXCLUDED.dismissed_unread_ids, '{}'::bigint[]),
			updated_at = EXCLUDED.updated_at
	`, pos.SessionID, pos.UserID, pos.LastReadMessageID, pos.DismissedUnreadIDs, pos.UpdatedAt)
	if err != nil {
		return nil, fmt.Errorf("upsert chat_read_positions failed: %w", err)
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, fmt.Errorf("commit tx failed: %w", err)
	}

	return &pos, nil
}

// compressDismissed 自动压缩算法：当离散消费集合与 baseline 连续时，推进 baseline 并裁切数组
func compressDismissed(baseline int64, dismissed []int64, futureMsgIDs []int64) (int64, []int64) {
	if len(dismissed) == 0 {
		return baseline, []int64{}
	}

	dMap := make(map[int64]bool, len(dismissed))
	for _, id := range dismissed {
		if id > baseline {
			dMap[id] = true
		}
	}

	// 沿现有消息列表升序尝试推进连续基线
	for _, mid := range futureMsgIDs {
		if mid <= baseline {
			continue
		}
		if dMap[mid] {
			baseline = mid
			delete(dMap, mid)
		} else {
			// 遇到第一条未消费的消息断点，停止推进连续基线
			break
		}
	}

	newDismissed := make([]int64, 0, len(dMap))
	for id := range dMap {
		if id > baseline {
			newDismissed = append(newDismissed, id)
		}
	}
	sort.Slice(newDismissed, func(i, j int) bool { return newDismissed[i] < newDismissed[j] })

	return baseline, newDismissed
}