backend/internal/pkg/ws/hub.go

package ws

import (
	"encoding/json"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"go.uber.org/zap"
)

const (
	defaultUserBufferSize    = 128
	groupPushConcurrency     = 64
	groupPushQueueSize       = 512
	maxTotalConnections      = 100000
	maxConnectionsPerUser    = 10
	ackTimeout               = 30 * time.Second
	ackCheckInterval         = 10 * time.Second
)

// Event 表示一个WebSocket事件
type Event struct {
	ID      uint64 `json:"event_id"`
	Type    string `json:"type"`
	Event   string `json:"event"` // 事件名称（兼容多种协议）
	TS      int64  `json:"ts"`
	Payload any    `json:"payload"`
}

// Client 表示一个WebSocket客户端连接
type Client struct {
	ID         uint64
	UserID     string
	Send       chan []byte
	Quit       chan struct{}
	PendingAcks map[string]time.Time
}

// ErrConnectionLimit 连接数限制错误
var ErrConnectionLimit = errors.New("connection limit reached")

// Message 表示WebSocket消息
type Message struct {
	Type    string          `json:"type"`
	Payload json.RawMessage `json:"payload"`
}

// ResponseMessage 表示发送给客户端的消息
type ResponseMessage struct {
	EventID uint64 `json:"event_id"`
	Type    string `json:"type"`
	TS      int64  `json:"ts"`
	Payload any    `json:"payload,omitempty"`
}

// ErrorMessage 表示错误消息
type ErrorMessage struct {
	Type    string `json:"type"`
	Payload struct {
		Code    string `json:"code"`
		Message string `json:"message"`
	} `json:"payload"`
}

// subscriber 表示一个事件订阅者
type subscriber struct {
	id   uint64
	ch   chan Event
	quit chan struct{}
}

// DisconnectHandler 用户断开连接时的回调函数
// userID: 断开的用户ID
// remainingCount: 该用户剩余的连接数（0表示用户完全离线）
type DisconnectHandler func(userID string, remainingCount int)

// ConnectHandler 用户连接时的回调函数
// userID: 连接的用户ID
// isFirstConnection: 是否是该用户的第一个连接（之前完全离线）
type ConnectHandler func(userID string, isFirstConnection bool)

// Hub WebSocket连接管理器
type Hub struct {
	seq       uint64
	connCount atomic.Int64

	mu                 sync.RWMutex
	clients            map[string]map[uint64]*Client
	subscribers        map[string]map[uint64]*subscriber
	disconnectHandlers []DisconnectHandler
	connectHandlers    []ConnectHandler
}

// NewHub 创建WebSocket Hub
func NewHub() *Hub {
	return &Hub{
		clients:            make(map[string]map[uint64]*Client),
		subscribers:        make(map[string]map[uint64]*subscriber),
		disconnectHandlers: make([]DisconnectHandler, 0),
		connectHandlers:    make([]ConnectHandler, 0),
	}
}

// NextID 获取下一个事件ID
func (h *Hub) NextID() uint64 {
	return atomic.AddUint64(&h.seq, 1)
}

// Register 注册客户端连接，返回可能的错误（连接数超限）
func (h *Hub) Register(client *Client) error {
	if h.connCount.Load() >= int64(maxTotalConnections) {
		return ErrConnectionLimit
	}

	h.mu.Lock()

	if len(h.clients[client.UserID]) >= maxConnectionsPerUser {
		h.mu.Unlock()
		return ErrConnectionLimit
	}

	_, existed := h.clients[client.UserID]
	isFirstConnection := !existed

	if !existed {
		h.clients[client.UserID] = make(map[uint64]*Client)
	}
	h.clients[client.UserID][client.ID] = client
	h.connCount.Add(1)

	zap.L().Debug("WebSocket client registered",
		zap.String("user_id", client.UserID),
		zap.Uint64("client_id", client.ID),
		zap.Int("total_clients", len(h.clients[client.UserID])),
		zap.Int64("total_connections", h.connCount.Load()),
		zap.Bool("is_first_connection", isFirstConnection),
	)

	handlers := make([]ConnectHandler, len(h.connectHandlers))
	copy(handlers, h.connectHandlers)
	h.mu.Unlock()

	if len(handlers) > 0 {
		go func() {
			for _, handler := range handlers {
				handler(client.UserID, isFirstConnection)
			}
		}()
	}

	return nil
}

// Unregister 注销客户端连接
func (h *Hub) Unregister(client *Client) {
	h.mu.Lock()

	var remainingCount int
	if userClients, ok := h.clients[client.UserID]; ok {
		if _, exists := userClients[client.ID]; exists {
			delete(userClients, client.ID)
			close(client.Quit)
			h.connCount.Add(-1)
			remainingCount = len(userClients)

			if remainingCount == 0 {
				delete(h.clients, client.UserID)
			}

			zap.L().Debug("WebSocket client unregistered",
				zap.String("user_id", client.UserID),
				zap.Uint64("client_id", client.ID),
				zap.Int("remaining_connections", remainingCount),
			)
		}
	}

	// 复制处理器列表避免锁竞争
	handlers := make([]DisconnectHandler, len(h.disconnectHandlers))
	copy(handlers, h.disconnectHandlers)
	h.mu.Unlock()

	// 异步调用断开处理器，避免阻塞 Unregister
	if len(handlers) > 0 {
		go func() {
			for _, handler := range handlers {
				handler(client.UserID, remainingCount)
			}
		}()
	}
}

// HasClients 检查用户是否有活跃连接
func (h *Hub) HasClients(userID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients[userID]) > 0
}

// FilterOnline 批量检查用户在线状态，一次加锁返回在线/离线列表
func (h *Hub) FilterOnline(userIDs []string) (online []string, offline []string) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, uid := range userIDs {
		if len(h.clients[uid]) > 0 {
			online = append(online, uid)
		} else {
			offline = append(offline, uid)
		}
	}
	return
}

// OnDisconnect 注册断开连接事件处理器
// 当用户断开 WebSocket 连接时，会调用所有注册的处理器
// 处理器会收到用户ID和该用户剩余的连接数
func (h *Hub) OnDisconnect(handler DisconnectHandler) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.disconnectHandlers = append(h.disconnectHandlers, handler)
}

// OnConnect 注册连接事件处理器
// 当用户建立 WebSocket 连接时，会调用所有注册的处理器
// 处理器会收到用户ID和是否是该用户的第一个连接
func (h *Hub) OnConnect(handler ConnectHandler) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.connectHandlers = append(h.connectHandlers, handler)
}

// GetClientCount 获取用户的连接数
func (h *Hub) GetClientCount(userID string) int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients[userID])
}

// PublishToUser 向指定用户发送事件
func (h *Hub) PublishToUser(userID string, eventType string, payload any) Event {
	ev := Event{
		ID:      h.NextID(),
		Type:    eventType,
		TS:      time.Now().UnixMilli(),
		Payload: payload,
	}
	h.publish(userID, ev)
	return ev
}

// PublishToUsers 向多个用户发送事件（共享一次序列化）
func (h *Hub) PublishToUsers(userIDs []string, eventType string, payload any) {
	if len(userIDs) == 0 {
		return
	}

	ev := Event{
		ID:      h.NextID(),
		Type:    eventType,
		TS:      time.Now().UnixMilli(),
		Payload: payload,
	}

	msg := ResponseMessage{
		EventID: ev.ID,
		Type:    eventType,
		TS:      ev.TS,
		Payload: payload,
	}
	data, err := json.Marshal(msg)
	if err != nil {
		zap.L().Error("Failed to marshal WebSocket message",
			zap.String("event_type", eventType),
			zap.Error(err),
		)
		return
	}

	for _, uid := range userIDs {
		h.publishPreMarshaled(uid, ev, data)
	}
}

// publishPreMarshaled 使用预序列化的数据发送事件
func (h *Hub) publishPreMarshaled(userID string, ev Event, data []byte) {
	h.mu.RLock()
	targets := make([]*Client, 0, len(h.clients[userID]))
	for _, c := range h.clients[userID] {
		targets = append(targets, c)
	}
	h.mu.RUnlock()

	for _, c := range targets {
		select {
		case <-c.Quit:
		case c.Send <- data:
		default:
			zap.L().Warn("WebSocket client buffer full, closing slow client",
				zap.String("user_id", userID),
				zap.Uint64("client_id", c.ID),
			)
			close(c.Quit)
		}
	}
}

// PublishToUserOnline 仅在用户在线时发送事件，不存入历史回放
// 适用于通话信令等实时性消息，避免断线重连时重放过期消息
func (h *Hub) PublishToUserOnline(userID string, eventType string, payload any) bool {
	h.mu.RLock()
	targets := make([]*Client, 0, len(h.clients[userID]))
	for _, c := range h.clients[userID] {
		targets = append(targets, c)
	}
	h.mu.RUnlock()

	if len(targets) == 0 {
		return false
	}

	msg := ResponseMessage{
		EventID: h.NextID(),
		Type:    eventType,
		TS:      time.Now().UnixMilli(),
		Payload: payload,
	}
	data, err := json.Marshal(msg)
	if err != nil {
		return false
	}

	for _, c := range targets {
		select {
		case <-c.Quit:
		case c.Send <- data:
		default:
		}
	}
	return true
}

// PublishToUserOnlineReliable 可靠地发送事件给在线用户（阻塞发送）
// 用于重要的信令消息（如 SDP, ICE candidate），确保消息送达
// 如果用户不在线，返回 false
func (h *Hub) PublishToUserOnlineReliable(userID string, eventType string, payload any) bool {
	h.mu.RLock()
	targets := make([]*Client, 0, len(h.clients[userID]))
	for _, c := range h.clients[userID] {
		targets = append(targets, c)
	}
	h.mu.RUnlock()

	if len(targets) == 0 {
		return false
	}

	msg := ResponseMessage{
		EventID: h.NextID(),
		Type:    eventType,
		TS:      time.Now().UnixMilli(),
		Payload: payload,
	}
	data, err := json.Marshal(msg)
	if err != nil {
		return false
	}

	// 阻塞发送，确保消息送达
	for _, c := range targets {
		select {
		case <-c.Quit:
			// 客户端已断开，跳过
		case c.Send <- data:
			// 消息已发送
		}
	}
	return true
}

// publish 内部发布方法
func (h *Hub) publish(userID string, ev Event) {
	h.mu.RLock()
	targets := make([]*Client, 0, len(h.clients[userID]))
	for _, c := range h.clients[userID] {
		targets = append(targets, c)
	}
	h.mu.RUnlock()

	// 构建消息
	msg := ResponseMessage{
		EventID: ev.ID,
		Type:    ev.Type,
		TS:      ev.TS,
		Payload: ev.Payload,
	}
	data, err := json.Marshal(msg)
	if err != nil {
		zap.L().Error("Failed to marshal WebSocket message",
			zap.String("user_id", userID),
			zap.String("event_type", ev.Type),
			zap.Error(err),
		)
		return
	}

	for _, c := range targets {
		select {
		case <-c.Quit:
		case c.Send <- data:
		default:
			zap.L().Warn("WebSocket client buffer full, closing slow client",
				zap.String("user_id", userID),
				zap.Uint64("client_id", c.ID),
			)
			close(c.Quit)
		}
	}
}

// SendError 向客户端发送错误消息
func (h *Hub) SendError(client *Client, code string, message string) {
	errMsg := ErrorMessage{
		Type: "error",
	}
	errMsg.Payload.Code = code
	errMsg.Payload.Message = message

	data, err := json.Marshal(errMsg)
	if err != nil {
		return
	}

	select {
	case <-client.Quit:
	case client.Send <- data:
	default:
	}
}

// Broadcast 广播消息给所有连接
func (h *Hub) Broadcast(eventType string, payload any) {
	h.mu.RLock()
	userIDs := make([]string, 0, len(h.clients))
	for uid := range h.clients {
		userIDs = append(userIDs, uid)
	}
	h.mu.RUnlock()

	h.PublishToUsers(userIDs, eventType, payload)
}

// PublishToUsersConcurrent 并发推送消息给多个用户（适用于群聊等大列表场景）
// 使用有界并发控制，避免 goroutine 爆炸
func (h *Hub) PublishToUsersConcurrent(userIDs []string, eventType string, payload any) {
	if len(userIDs) == 0 {
		return
	}

	ev := Event{
		ID:      h.NextID(),
		Type:    eventType,
		TS:      time.Now().UnixMilli(),
		Payload: payload,
	}

	msg := ResponseMessage{
		EventID: ev.ID,
		Type:    eventType,
		TS:      ev.TS,
		Payload: payload,
	}
	data, err := json.Marshal(msg)
	if err != nil {
		zap.L().Error("Failed to marshal WebSocket message for batch push",
			zap.String("event_type", eventType),
			zap.Error(err),
		)
		return
	}

	sem := make(chan struct{}, groupPushConcurrency)
	var wg sync.WaitGroup
	wg.Add(len(userIDs))

	for _, uid := range userIDs {
		sem <- struct{}{}
		go func(userID string) {
			defer wg.Done()
			defer func() { <-sem }()

			h.mu.RLock()
			targets := make([]*Client, 0, len(h.clients[userID]))
			for _, c := range h.clients[userID] {
				targets = append(targets, c)
			}
			h.mu.RUnlock()

			for _, c := range targets {
				select {
				case <-c.Quit:
				case c.Send <- data:
				default:
					zap.L().Warn("WebSocket client buffer full, closing slow client",
						zap.String("user_id", userID),
						zap.Uint64("client_id", c.ID),
					)
					close(c.Quit)
				}
			}
		}(uid)
	}

	wg.Wait()
}

// GetOnlineUsers 获取在线用户列表
func (h *Hub) GetOnlineUsers() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()

	users := make([]string, 0, len(h.clients))
	for uid := range h.clients {
		users = append(users, uid)
	}
	return users
}

// GetOnlineCount 获取在线用户数
func (h *Hub) GetOnlineCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}

// GetConnectionCount 获取总连接数
func (h *Hub) GetConnectionCount() int64 {
	return h.connCount.Load()
}

// CloseAllConnections 关闭所有WebSocket连接（用于优雅关闭）
func (h *Hub) CloseAllConnections() {
	h.mu.Lock()
	defer h.mu.Unlock()

	for _, userClients := range h.clients {
		for _, c := range userClients {
			select {
			case <-c.Quit:
			default:
				close(c.Quit)
			}
		}
	}
	h.clients = make(map[string]map[uint64]*Client)
	h.connCount.Store(0)
}

// Subscribe 订阅事件
func (h *Hub) Subscribe(userID string, afterID uint64) (chan Event, func()) {
	subID := h.NextID()
	sub := &subscriber{
		id:   subID,
		ch:   make(chan Event, defaultUserBufferSize),
		quit: make(chan struct{}),
	}

	h.mu.Lock()
	if _, ok := h.subscribers[userID]; !ok {
		h.subscribers[userID] = make(map[uint64]*subscriber)
	}
	h.subscribers[userID][subID] = sub
	h.mu.Unlock()

	cancel := func() {
		h.mu.Lock()
		defer h.mu.Unlock()
		if userSubs, ok := h.subscribers[userID]; ok {
			if s, exists := userSubs[subID]; exists {
				close(s.quit)
				delete(userSubs, subID)
				close(s.ch)
			}
			if len(userSubs) == 0 {
				delete(h.subscribers, userID)
			}
		}
	}

	return sub.ch, cancel
}

// AckMessage 确认消息已送达，从 pendingAcks 中移除
func (h *Hub) AckMessage(userID, messageID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for _, client := range h.clients[userID] {
		if _, ok := client.PendingAcks[messageID]; ok {
			delete(client.PendingAcks, messageID)
			return true
		}
	}
	return false
}

// startAckChecker 启动超时 ACK 检查器
func (h *Hub) startAckChecker(onTimeout func(userID, messageID string)) {
	ticker := time.NewTicker(ackCheckInterval)
	defer ticker.Stop()

	for range ticker.C {
		h.mu.RLock()
		// 收集超时的 ack
		type timeoutEntry struct {
			userID    string
			messageID string
		}
		var timeouts []timeoutEntry

		now := time.Now()
		for uid, userClients := range h.clients {
			for _, client := range userClients {
				for msgID, sentAt := range client.PendingAcks {
					if now.Sub(sentAt) > ackTimeout {
						timeouts = append(timeouts, timeoutEntry{uid, msgID})
						delete(client.PendingAcks, msgID)
					}
				}
			}
		}
		h.mu.RUnlock()

		for _, entry := range timeouts {
			onTimeout(entry.userID, entry.messageID)
		}
	}
}

// StartAckChecker 启动 ACK 检查器（公开方法）
func (h *Hub) StartAckChecker(onTimeout func(userID, messageID string)) {
	go h.startAckChecker(onTimeout)
}

// EncodeData 编码事件数据
func EncodeData(ev Event) (string, error) {
	body, err := json.Marshal(ev.Payload)
	if err != nil {
		return "", err
	}
	return string(body), nil
}