ai-cs/backend/websocket/hub.go

package websocket

import (
	"log"
	"sync"
)

// OnClientConnectCallback 客户端连接时的回调函数。
// conversationID: 对话ID
// isVisitor: 是否是访客
// visitorCount: 该对话当前的访客连接数
type OnClientConnectCallback func(conversationID uint, isVisitor bool, visitorCount int)

// OnClientDisconnectCallback 客户端断开连接时的回调函数。
// conversationID: 对话ID
// isVisitor: 是否是访客
// visitorCount: 该对话当前的访客连接数（断开后）
type OnClientDisconnectCallback func(conversationID uint, isVisitor bool, visitorCount int)

// Hub 管理所有 WebSocket 连接
// 每个对话（conversation）可以有多个人连接（访客和客服）
type Hub struct {
	// 每个对话ID对应的客户端连接列表
	// conversationID -> []*Client
	conversations map[uint]map[*Client]bool

	// 注册新客户端（当有人连接时）
	register chan *Client

	// 注销客户端（当有人断开连接时）
	unregister chan *Client

	// 广播消息（当有新消息时，推送给所有相关的客户端）
	broadcast chan *Message

	// 互斥锁（保护并发访问）
	mu sync.RWMutex

	// 回调函数
	onConnect    OnClientConnectCallback
	onDisconnect OnClientDisconnectCallback
}

// Message 是要广播的消息
type Message struct {
	ConversationID uint        `json:"conversation_id"`
	Data           interface{} `json:"data"` // 消息内容（可以是 Message 对象）
	Type           string      `json:"type"` // 消息类型：new_message, conversation_update 等
}

// NewHub 创建一个新的 Hub
func NewHub(onConnect OnClientConnectCallback, onDisconnect OnClientDisconnectCallback) *Hub {
	return &Hub{
		conversations: make(map[uint]map[*Client]bool),
		register:      make(chan *Client),
		unregister:    make(chan *Client),
		broadcast:     make(chan *Message, 256),
		onConnect:     onConnect,
		onDisconnect:  onDisconnect,
	}
}

// Run 启动 Hub，处理所有事件
func (h *Hub) Run() {
	for {
		select {
		// 新客户端连接
		case client := <-h.register:
			h.mu.Lock()
			// 如果这个对话还没有客户端，创建一个新的 map
			if h.conversations[client.conversationID] == nil {
				h.conversations[client.conversationID] = make(map[*Client]bool)
			}
			// 把这个客户端加入到对话中
			h.conversations[client.conversationID][client] = true

			// 统计该对话的访客连接数
			visitorCount := 0
			for c := range h.conversations[client.conversationID] {
				if c.isVisitor {
					visitorCount++
				}
			}
			h.mu.Unlock()

			// 调用连接回调函数
			if h.onConnect != nil {
				h.onConnect(client.conversationID, client.isVisitor, visitorCount)
			}

		// 客户端断开连接
		case client := <-h.unregister:
			h.mu.Lock()
			// 从对话中移除这个客户端
			wasVisitor := client.isVisitor
			if clients, ok := h.conversations[client.conversationID]; ok {
				if _, ok := clients[client]; ok {
					delete(clients, client)
					// 关闭发送通道（避免重复关闭导致 panic）
					select {
					case _, ok := <-client.send:
						if !ok {
							// 通道已经关闭，不需要再次关闭
						}
					default:
						// 通道未关闭，关闭它
						close(client.send)
					}

					// 统计该对话的访客连接数（断开后）
					visitorCount := 0
					for c := range clients {
						if c.isVisitor {
							visitorCount++
						}
					}

					// 如果这个对话没有客户端了，删除对话
					if len(clients) == 0 {
						delete(h.conversations, client.conversationID)
					}
					h.mu.Unlock()

					// 调用断开回调函数
					if h.onDisconnect != nil {
						h.onDisconnect(client.conversationID, wasVisitor, visitorCount)
					}
				} else {
					h.mu.Unlock()
					log.Printf("⚠️ 客户端断开时未找到: 对话ID=%d", client.conversationID)
				}
			} else {
				h.mu.Unlock()
				log.Printf("⚠️ 客户端断开时对话不存在: 对话ID=%d", client.conversationID)
			}

		// 广播消息
		case message := <-h.broadcast:
			h.mu.RLock()
			// 找到这个对话的所有客户端
			clients, ok := h.conversations[message.ConversationID]
			if !ok {
				h.mu.RUnlock()
				log.Printf("⚠️ 广播消息失败: 对话ID=%d 没有客户端连接", message.ConversationID)
				continue
			}
			// 创建一个客户端列表的副本（避免在遍历时修改）
			clientList := make([]*Client, 0, len(clients))
			for client := range clients {
				clientList = append(clientList, client)
			}
			h.mu.RUnlock()

			// 给所有客户端发送消息
			for _, client := range clientList {
				select {
				case client.send <- message:
				default:
					// 如果发送失败（客户端可能已经断开），关闭连接
					log.Printf("⚠️ 发送消息失败: 对话ID=%d, 客户端断开", client.conversationID)
					close(client.send)
					h.mu.Lock()
					delete(h.conversations[client.conversationID], client)
					h.mu.Unlock()
				}
			}
		}
	}
}

// BroadcastMessage 广播消息到指定对话的所有客户端
func (h *Hub) BroadcastMessage(conversationID uint, messageType string, data interface{}) {
	h.broadcast <- &Message{
		ConversationID: conversationID,
		Type:           messageType,
		Data:           data,
	}
}

// BroadcastToAllAgents 广播消息到所有客服客户端（不管连接到哪个对话）
// 用于 visitor_status_update 等需要所有客服都收到的事件
func (h *Hub) BroadcastToAllAgents(messageType string, data interface{}) {
	h.mu.RLock()
	// 收集所有客服客户端（isVisitor == false）
	allAgents := make([]*Client, 0)
	for _, clients := range h.conversations {
		for client := range clients {
			if !client.isVisitor {
				allAgents = append(allAgents, client)
			}
		}
	}
	h.mu.RUnlock()

	// 为每个客服客户端创建消息并发送
	for _, client := range allAgents {
		message := &Message{
			ConversationID: client.conversationID, // 使用客户端连接的对话ID
			Type:           messageType,
			Data:           data,
		}
		select {
		case client.send <- message:
		default:
			// 如果发送失败（客户端可能已经断开），关闭连接
			log.Printf("⚠️ 发送消息到客服失败: 对话ID=%d, 客户端断开", client.conversationID)
			close(client.send)
			h.mu.Lock()
			if clients, ok := h.conversations[client.conversationID]; ok {
				delete(clients, client)
				if len(clients) == 0 {
					delete(h.conversations, client.conversationID)
				}
			}
			h.mu.Unlock()
		}
	}
}