backend/internal/pkg/ws/bus.go

package ws

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
)

// busMsg Redis Pub/Sub 消息格式
type busMsg struct {
	InstanceID  string   `json:"instance_id"`
	MsgType     string   `json:"msg_type"` // "user_msg" | "broadcast"
	TargetUsers []string `json:"target_users"`
	EventType   string   `json:"event_type"`
	EventID     uint64   `json:"event_id"`
	TS          int64    `json:"ts"`
	Payload     json.RawMessage `json:"payload"`
}

// WSClusterConfig WebSocket 集群配置
type WSClusterConfig struct {
	InstanceID        string
	MsgChannel        string
	OnlineTTL         int
	HeartbeatInterval int
}

// Bus 组合 Hub + Redis Pub/Sub + OnlineTracker，实现跨实例消息路由
type Bus struct {
	hub       *Hub
	rdb       *redis.Client
	tracker   *OnlineTracker

	instanceID string
	channel    string
	sub        *redis.PubSub
	cancel     context.CancelFunc

	mu     sync.Mutex
	closed bool
}

// NewBus 创建 Bus 实例
// 如果 rdb 为 nil（standalone 模式或 Redis 不可用），Bus 退化为纯本地 Hub 代理
func NewBus(hub *Hub, rdb *redis.Client, cfg *WSClusterConfig) *Bus {
	instanceID := cfg.InstanceID
	if instanceID == "" {
		instanceID = uuid.New().String()
	}

	onlineTTL := time.Duration(cfg.OnlineTTL) * time.Second
	heartbeatInterval := time.Duration(cfg.HeartbeatInterval) * time.Second
	if heartbeatInterval <= 0 {
		heartbeatInterval = 20 * time.Second
	}
	if onlineTTL <= 0 {
		onlineTTL = 60 * time.Second
	}

	tracker := NewOnlineTracker(rdb, instanceID, onlineTTL, heartbeatInterval)

	return &Bus{
		hub:        hub,
		rdb:        rdb,
		tracker:    tracker,
		instanceID: instanceID,
		channel:    cfg.MsgChannel,
	}
}

// Hub 获取底层 Hub（供 WebSocket handler 使用 Register/Unregister 等本地操作）
func (b *Bus) Hub() *Hub {
	return b.hub
}

// InstanceID 获取实例 ID
func (b *Bus) InstanceID() string {
	return b.instanceID
}

// Start 启动 Bus 的 Redis 订阅和心跳
func (b *Bus) Start() error {
	if b.rdb == nil {
		zap.L().Info("WebSocket Bus: no Redis client, running in standalone mode")
		return nil
	}

	// 启动心跳
	b.tracker.StartHeartbeat()

	// 订阅 Redis channel
	ctx, cancel := context.WithCancel(context.Background())
	b.cancel = cancel

	sub := b.rdb.Subscribe(ctx, b.channel)
	b.sub = sub

	// 确保订阅已建立
	if _, err := sub.Receive(ctx); err != nil {
		cancel()
		return err
	}

	go b.subscribeLoop(ctx, sub)

	zap.L().Info("WebSocket Bus started",
		zap.String("instance_id", b.instanceID),
		zap.String("channel", b.channel),
	)
	return nil
}

// Register 注册客户端连接（代理 Hub.Register + tracker）
func (b *Bus) Register(client *Client) error {
	if err := b.hub.Register(client); err != nil {
		return err
	}
	b.tracker.UserConnected(client.UserID)
	return nil
}

// Unregister 注销客户端连接（代理 Hub.Unregister + tracker）
func (b *Bus) Unregister(client *Client) {
	b.hub.Unregister(client)
	b.tracker.UserDisconnected(client.UserID)
}

// Close 关闭 Bus：关闭订阅，清理在线状态
func (b *Bus) Close() error {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.closed {
		return nil
	}
	b.closed = true

	if b.cancel != nil {
		b.cancel()
	}
	if b.sub != nil {
		b.sub.Close()
	}
	b.tracker.Stop()
	b.hub.CloseAllConnections()

	zap.L().Info("WebSocket Bus closed", zap.String("instance_id", b.instanceID))
	return nil
}

// --- MessagePublisher 接口实现 ---

// PublishToUser 向指定用户发送事件（本地 + 远程）
func (b *Bus) PublishToUser(userID string, eventType string, payload any) Event {
	ev := b.hub.PublishToUser(userID, eventType, payload)
	b.publishToRemote([]string{userID}, "user_msg", ev, payload)
	return ev
}

// PublishToUsers 向多个用户发送事件
func (b *Bus) PublishToUsers(userIDs []string, eventType string, payload any) {
	b.hub.PublishToUsers(userIDs, eventType, payload)
	b.publishToRemote(userIDs, "user_msg", Event{
		Type: eventType,
		TS:   time.Now().UnixMilli(),
	}, payload)
}

// PublishToUsersConcurrent 并发推送消息给多个用户
func (b *Bus) PublishToUsersConcurrent(userIDs []string, eventType string, payload any) {
	b.hub.PublishToUsersConcurrent(userIDs, eventType, payload)
	b.publishToRemote(userIDs, "user_msg", Event{
		Type: eventType,
		TS:   time.Now().UnixMilli(),
	}, payload)
}

// PublishToUserOnline 仅在用户在线时发送事件
func (b *Bus) PublishToUserOnline(userID string, eventType string, payload any) bool {
	delivered := b.hub.PublishToUserOnline(userID, eventType, payload)
	if !delivered {
		// 本地不在线，尝试远程投递
		ev := Event{
			Type: eventType,
			TS:   time.Now().UnixMilli(),
		}
		b.publishToRemote([]string{userID}, "user_msg", ev, payload)
	}
	return delivered
}

// PublishToUserOnlineReliable 可靠地发送事件给在线用户
func (b *Bus) PublishToUserOnlineReliable(userID string, eventType string, payload any) bool {
	delivered := b.hub.PublishToUserOnlineReliable(userID, eventType, payload)
	if !delivered {
		ev := Event{
			Type: eventType,
			TS:   time.Now().UnixMilli(),
		}
		b.publishToRemote([]string{userID}, "user_msg", ev, payload)
	}
	return delivered
}

// Broadcast 广播消息给所有连接
func (b *Bus) Broadcast(eventType string, payload any) {
	b.hub.Broadcast(eventType, payload)
	b.publishToRemote(nil, "broadcast", Event{
		Type: eventType,
		TS:   time.Now().UnixMilli(),
	}, payload)
}

// HasClients 检查用户是否在线（本地优先，然后查 Redis）
func (b *Bus) HasClients(userID string) bool {
	if b.hub.HasClients(userID) {
		return true
	}
	return b.tracker.IsOnline(userID)
}

// FilterOnline 批量检查用户在线状态
func (b *Bus) FilterOnline(userIDs []string) (online, offline []string) {
	localOnline, offline := b.hub.FilterOnline(userIDs)
	if b.rdb == nil {
		return localOnline, offline
	}
	// 对 offline 的列表，查 Redis
	remoteOnline, stillOffline := b.tracker.FilterOnline(offline)
	online = append(localOnline, remoteOnline...)
	return online, stillOffline
}

// SendError 向客户端发送错误消息（本地操作）
func (b *Bus) SendError(client *Client, code string, message string) {
	b.hub.SendError(client, code, message)
}

// NextID 获取下一个事件 ID
func (b *Bus) NextID() uint64 {
	return b.hub.NextID()
}

// OnDisconnect 注册断开连接事件处理器
func (b *Bus) OnDisconnect(handler DisconnectHandler) {
	b.hub.OnDisconnect(handler)
}

// OnConnect 注册连接事件处理器
func (b *Bus) OnConnect(handler ConnectHandler) {
	b.hub.OnConnect(handler)
}

// --- 内部方法 ---

// publishToRemote 发布消息到 Redis channel
func (b *Bus) publishToRemote(targetUsers []string, msgType string, ev Event, payload any) {
	if b.rdb == nil {
		return
	}

	payloadBytes, err := json.Marshal(payload)
	if err != nil {
		zap.L().Error("Bus: failed to marshal payload for remote publish",
			zap.String("event_type", ev.Type),
			zap.Error(err),
		)
		return
	}

	msg := busMsg{
		InstanceID:  b.instanceID,
		MsgType:     msgType,
		TargetUsers: targetUsers,
		EventType:   ev.Type,
		EventID:     ev.ID,
		TS:          ev.TS,
		Payload:     payloadBytes,
	}

	data, err := json.Marshal(msg)
	if err != nil {
		zap.L().Error("Bus: failed to marshal busMsg", zap.Error(err))
		return
	}

	if err := b.rdb.Publish(context.Background(), b.channel, data).Err(); err != nil {
		zap.L().Error("Bus: failed to publish to Redis",
			zap.String("channel", b.channel),
			zap.Error(err),
		)
	}
}

// subscribeLoop 处理 Redis 订阅消息
func (b *Bus) subscribeLoop(ctx context.Context, sub *redis.PubSub) {
	ch := sub.Channel()
	for {
		select {
		case <-ctx.Done():
			return
		case msg, ok := <-ch:
			if !ok {
				return
			}
			b.handleRemoteMessage(msg)
		}
	}
}

// handleRemoteMessage 处理从 Redis 收到的远程消息
func (b *Bus) handleRemoteMessage(msg *redis.Message) {
	var bm busMsg
	if err := json.Unmarshal([]byte(msg.Payload), &bm); err != nil {
		zap.L().Error("Bus: failed to unmarshal remote message", zap.Error(err))
		return
	}

	// 过滤自身发送的消息
	if bm.InstanceID == b.instanceID {
		return
	}

	switch bm.MsgType {
	case "user_msg":
		b.handleRemoteUserMsg(bm)
	case "broadcast":
		b.handleRemoteBroadcast(bm)
	default:
		zap.L().Warn("Bus: unknown msg_type from remote",
			zap.String("msg_type", bm.MsgType),
		)
	}
}

// handleRemoteUserMsg 处理远程定向消息
func (b *Bus) handleRemoteUserMsg(bm busMsg) {
	// 构造 ResponseMessage 用于本地投递
	respMsg := ResponseMessage{
		EventID: bm.EventID,
		Type:    bm.EventType,
		TS:      bm.TS,
	}

	// 反序列化 payload
	var payload any
	if err := json.Unmarshal(bm.Payload, &payload); err != nil {
		zap.L().Error("Bus: failed to unmarshal remote payload", zap.Error(err))
		return
	}
	respMsg.Payload = payload

	data, err := json.Marshal(respMsg)
	if err != nil {
		zap.L().Error("Bus: failed to marshal response message", zap.Error(err))
		return
	}

	b.hub.deliverToUsers(bm.TargetUsers, data)
}

// handleRemoteBroadcast 处理远程广播消息
func (b *Bus) handleRemoteBroadcast(bm busMsg) {
	b.hub.Broadcast(bm.EventType, bm.Payload)
}

// deliverToUsers 向指定用户列表投递预序列化消息
// 这是 Hub 批量投递的内部方法，Bus 需要调用它来投递远程消息
func (h *Hub) deliverToUsers(userIDs []string, data []byte) {
	for _, uid := range userIDs {
		h.mu.RLock()
		targets := make([]*Client, 0, len(h.clients[uid]))
		for _, c := range h.clients[uid] {
			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", uid),
					zap.Uint64("client_id", c.ID),
				)
				close(c.Quit)
			}
		}
	}
}