hfb_sys/backend/internal/modules/chathub/hub.go

package chathub

import (
	"encoding/json"
	"sync"

	"gorm.io/gorm"
)

// ChatEvent 是通过 SSE 推送给客户端的事件。
type ChatEvent struct {
	Type           string       `json:"type"` // "new_message" | "conversation_updated" | "conversation_read"
	ConversationID uint64       `json:"conversation_id"`
	Message        *MessageData `json:"message,omitempty"`
	// 以下字段仅在 conversation_read 事件中使用，标识谁在何时读取了会话。
	ReaderType string `json:"reader_type,omitempty"`
	ReaderID   uint64 `json:"reader_id,omitempty"`
	ReadAt     string `json:"read_at,omitempty"`
}

// MessageData 是事件中携带的消息数据，与 chat.MessageDTO 对齐。
type MessageData struct {
	ID                 uint64   `json:"id"`
	ConversationID     uint64   `json:"conversation_id"`
	SenderType         string   `json:"sender_type"`
	SenderID           uint64   `json:"sender_id"`
	SenderRole         string   `json:"sender_role"`
	SenderName         string   `json:"sender_name"`
	ContentType        string   `json:"content_type"`
	Content            string   `json:"content"`
	AttachmentURLS     []string `json:"attachment_urls"`
	AdminAttentionType string   `json:"admin_attention_type,omitempty"`
	CreatedAt          string   `json:"created_at"`
}

// principal 标识一个连接方。
type principal struct {
	Type string // "user" or "admin"
	ID   uint64
}

// Hub 管理所有 SSE 连接。
type Hub struct {
	mu      sync.RWMutex
	clients map[principal]map[chan *ChatEvent]struct{}
	db      *gorm.DB
}

// NewHub 创建 Hub 实例。
func NewHub(db *gorm.DB) *Hub {
	return &Hub{
		clients: make(map[principal]map[chan *ChatEvent]struct{}),
		db:      db,
	}
}

// Subscribe 注册一个 SSE 连接，返回事件 channel。
func (h *Hub) Subscribe(pType string, pID uint64) <-chan *ChatEvent {
	ch := make(chan *ChatEvent, 16)
	key := principal{Type: pType, ID: pID}
	h.mu.Lock()
	if h.clients[key] == nil {
		h.clients[key] = make(map[chan *ChatEvent]struct{})
	}
	h.clients[key][ch] = struct{}{}
	h.mu.Unlock()
	return ch
}

// Unsubscribe 注销 SSE 连接。
func (h *Hub) Unsubscribe(pType string, pID uint64, ch <-chan *ChatEvent) {
	key := principal{Type: pType, ID: pID}
	h.mu.Lock()
	if clients, ok := h.clients[key]; ok {
		// 找到对应的发送 channel 并删除
		for sendCh := range clients {
			if sendCh == ch {
				delete(clients, sendCh)
				close(sendCh)
				break
			}
		}
		if len(clients) == 0 {
			delete(h.clients, key)
		}
	}
	h.mu.Unlock()
}

// DisconnectUser 关闭指定用户当前进程内的全部实时连接。
func (h *Hub) DisconnectUser(userID uint64) {
	h.disconnect("user", userID)
}

func (h *Hub) disconnect(pType string, pID uint64) {
	key := principal{Type: pType, ID: pID}
	h.mu.Lock()
	clients := h.clients[key]
	delete(h.clients, key)
	for ch := range clients {
		close(ch)
	}
	h.mu.Unlock()
}

// NotifyConversation 查询会话参与者，向所有在线参与者推送事件。
func (h *Hub) NotifyConversation(conversationID uint64, event *ChatEvent) {
	if h.db == nil {
		return
	}
	// 查询会话参与者
	type participantRow struct {
		ParticipantType string
		ParticipantID   uint64
	}
	var rows []participantRow
	h.db.Table("chat_participants").
		Select("participant_type, participant_id").
		Where("conversation_id = ?", conversationID).
		Find(&rows)

	h.mu.RLock()
	defer h.mu.RUnlock()

	seen := make(map[principal]struct{})
	for _, row := range rows {
		key := principal{Type: row.ParticipantType, ID: row.ParticipantID}
		if _, ok := seen[key]; ok {
			continue
		}
		seen[key] = struct{}{}
		if clients, ok := h.clients[key]; ok {
			for ch := range clients {
				select {
				case ch <- event:
				default:
					// channel 满了，丢弃事件避免阻塞
				}
			}
		}
	}
}

// NotifyUser 直接向指定用户/Admin 推送事件。
func (h *Hub) NotifyUser(pType string, pID uint64, event *ChatEvent) {
	key := principal{Type: pType, ID: pID}
	h.mu.RLock()
	defer h.mu.RUnlock()
	if clients, ok := h.clients[key]; ok {
		for ch := range clients {
			select {
			case ch <- event:
			default:
			}
		}
	}
}

// NotifyAllAdmins 向所有在线管理员（客服）推送事件。
// 客服工作台的「未分配 / 全部」视图会展示当前客服并非参与者的会话，
// 仅靠 NotifyConversation（只推参与者）无法实时刷新，故对所有在线客服广播。
func (h *Hub) NotifyAllAdmins(event *ChatEvent) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for key, clients := range h.clients {
		if key.Type != "admin" {
			continue
		}
		for ch := range clients {
			select {
			case ch <- event:
			default:
			}
		}
	}
}

// MarshalEvent 将事件序列化为 SSE data 行。
func MarshalEvent(event *ChatEvent) string {
	raw, _ := json.Marshal(event)
	return string(raw)
}

// OnlineCount 返回当前在线连接数。
func (h *Hub) OnlineCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	count := 0
	for _, clients := range h.clients {
		count += len(clients)
	}
	return count
}

// OnlineCountByType 返回当前进程内指定主体类型的 SSE 连接数。
func (h *Hub) OnlineCountByType(pType string) int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	count := 0
	for key, clients := range h.clients {
		if key.Type == pType {
			count += len(clients)
		}
	}
	return count
}