gochat/internal/ws/broadcast.go

// Package ws provides the Redis Pub/Sub broadcast relay for cross-instance
// WebSocket message delivery. Mirrors Chatwoot's ActionCableBroadcastJob:
// when an event is published via the EventBus, this relay picks it up and
// broadcasts to all WebSocket clients subscribed to the relevant room.
package ws

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

	"github.com/redis/go-redis/v9"

	"github.com/gochat/gochat/pkg/logger"
)

// MessageHandler is the interface that the Hub implements to deliver
// messages to locally connected WebSocket clients. This decouples
// BroadcastRelay from the handler/ws package, avoiding circular imports.
type MessageHandler interface {
	// SendToAccount sends a message to all clients subscribed to an account room.
	SendToAccount(accountID uint, data []byte)
	// SendToRoom sends a message to all clients in a named room.
	SendToRoom(room string, data []byte)
}

// BroadcastRelay subscribes to Redis Pub/Sub channels and forwards
// messages to a MessageHandler for local WebSocket delivery.
// This enables cross-instance broadcasting: when one gochat server
// publishes an event, all server instances receive it via Redis
// and push to their locally connected WebSocket clients.
//
// Reference: Chatwoot ActionCableBroadcastJob — publishes to Redis,
// then ActionCableListener picks up the Redis message and delivers
// to subscribed WebSocket connections.
type BroadcastRelay struct {
	rdb   *redis.Client
	hub   MessageHandler // interface — not concrete Hub type
	mu    sync.Mutex
	subs  map[string]*redis.PubSub // active Redis subscriptions per channel
	ctx   context.Context
}

// NewBroadcastRelay creates a new broadcast relay instance.
func NewBroadcastRelay(rdb *redis.Client, hub MessageHandler) *BroadcastRelay {
	return &BroadcastRelay{
		rdb:  rdb,
		hub:  hub,
		subs: make(map[string]*redis.PubSub),
	}
}

// Start begins listening for Redis Pub/Sub messages.
// It subscribes to the global gochat broadcast channels.
func (r *BroadcastRelay) Start(ctx context.Context) error {
	r.ctx = ctx
	r.mu.Lock()
	defer r.mu.Unlock()

	// Subscribe to the wildcard pattern for all room channels
	// Pattern: gochat:ws:room:* and gochat:ws:account:*
	patterns := []string{
		RedisPrefixRoom + "*",
		RedisPrefixAccount + "*",
	}

	for _, pattern := range patterns {
		sub := r.rdb.PSubscribe(ctx, pattern)
		if err := sub.Ping(ctx); err != nil {
			return fmt.Errorf("failed to subscribe to Redis pattern %s: %w", pattern, err)
		}
		r.subs[pattern] = sub

		// Start receiving messages for this pattern
		go r.receiveLoop(sub, pattern)
	}

	logger.L().Infof("ws: broadcast relay started, watching %d patterns", len(patterns))
	return nil
}

// Stop gracefully shuts down all Redis subscriptions.
func (r *BroadcastRelay) Stop() error {
	r.mu.Lock()
	defer r.mu.Unlock()

	for pattern, sub := range r.subs {
		if err := sub.Close(); err != nil {
			logger.L().Warnf("ws: failed to close Redis subscription for %s: %v", pattern, err)
		}
	}
	r.subs = make(map[string]*redis.PubSub)
	logger.L().Info("ws: broadcast relay stopped")
	return nil
}

// Publish broadcasts a WSMessage to a Redis Pub/Sub channel so all
// gochat instances can receive it. This is the "publish" side of the
// cross-instance relay, called when a local event needs to reach
// clients on other server instances.
//
// Reference: Chatwoot ActionCableBroadcastJob.perform — publishes
// to Redis channel "room_{pubsub_token}" or "account_{account_id}".
func (r *BroadcastRelay) Publish(ctx context.Context, room string, msg *WSMessage) error {
	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("failed to marshal WSMessage for Redis publish: %w", err)
	}

	channel := RedisPrefixRoom + room
	result, err := r.rdb.Publish(ctx, channel, data).Result()
	if err != nil {
		return fmt.Errorf("failed to publish to Redis channel %s: %w", channel, err)
	}

	logger.L().Debugf("ws: published to Redis channel %s (%d receivers)", channel, result)
	return nil
}

// PublishAccount broadcasts a WSMessage to an account-level room.
func (r *BroadcastRelay) PublishAccount(ctx context.Context, accountID uint, msg *WSMessage) error {
	data, err := json.Marshal(msg)
	if err != nil {
		return fmt.Errorf("failed to marshal WSMessage for account publish: %w", err)
	}

	channel := fmt.Sprintf(RedisPrefixAccount+"%d", accountID)
	result, err := r.rdb.Publish(ctx, channel, data).Result()
	if err != nil {
		return fmt.Errorf("failed to publish to Redis channel %s: %w", channel, err)
	}

	logger.L().Debugf("ws: published to account Redis channel %s (%d receivers)", channel, result)
	return nil
}

// receiveLoop continuously receives messages from a Redis Pub/Sub subscription
// and forwards them to the MessageHandler for local WebSocket delivery.
func (r *BroadcastRelay) receiveLoop(sub *redis.PubSub, pattern string) {
	ch := sub.Channel()
	for {
		select {
		case <-r.ctx.Done():
			logger.L().Infof("ws: broadcast relay receive loop stopped for pattern %s", pattern)
			return
		case msg, ok := <-ch:
			if !ok {
				logger.L().Warnf("ws: Redis Pub/Sub channel closed for pattern %s", pattern)
				return
			}
			r.handleRedisMessage(msg)
		}
	}
}

// handleRedisMessage parses a Redis Pub/Sub message and forwards to the MessageHandler.
func (r *BroadcastRelay) handleRedisMessage(msg *redis.Message) {
	var wsMsg WSMessage
	if err := json.Unmarshal([]byte(msg.Payload), &wsMsg); err != nil {
		logger.L().Errorf("ws: failed to unmarshal Redis message on channel %s: %v", msg.Channel, err)
		return
	}

	// Determine the target room from the Redis channel name
	// Channel format: gochat:ws:room:{room_name} or gochat:ws:account:{account_id}
	data, err := json.Marshal(wsMsg)
	if err != nil {
		logger.L().Errorf("ws: failed to re-marshal WSMessage: %v", err)
		return
	}

	if accountID := extractAccountIDFromChannel(msg.Channel); accountID > 0 {
		r.hub.SendToAccount(accountID, data)
	} else if room := extractRoomFromChannel(msg.Channel); room != "" {
		r.hub.SendToRoom(room, data)
	} else {
		logger.L().Warnf("ws: unrecognized Redis channel format: %s", msg.Channel)
	}
}

// extractAccountIDFromChannel parses "gochat:ws:account:{id}" → accountID.
func extractAccountIDFromChannel(channel string) uint {
	prefix := RedisPrefixAccount
	if len(channel) <= len(prefix) {
		return 0
	}
	remainder := channel[len(prefix):]
	var id uint
	if _, err := fmt.Sscanf(remainder, "%d", &id); err == nil {
		return id
	}
	return 0
}

// extractRoomFromChannel parses "gochat:ws:room:{room_name}" → room_name.
func extractRoomFromChannel(channel string) string {
	prefix := RedisPrefixRoom
	if len(channel) <= len(prefix) {
		return ""
	}
	return channel[len(prefix):]
}