gochat/backend/internal/handler/ws/handler.go

package ws

import (
	"context"
	"encoding/json"
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	wspkg "github.com/gochat/gochat/internal/ws"
	"github.com/gochat/gochat/pkg/logger"
)

// uintToStr converts a uint to its string representation.
// Kept for potential use in identifier formatting.
func uintToStr(u uint) string {
	return strconv.FormatUint(uint64(u), 10)
}

// Handler manages WebSocket connections, upgrade, JWT auth, and message dispatch.
// Reference: Chatwoot ActionCable — WebSocket upgrade at /cable, JWT-based auth,
// room-based subscription model (AccountChannel, ConversationChannel).
//
// Supports two authentication paths:
//  1. Agent/User auth: JWT token (from query param or Authorization header)
//  2. Contact auth: pubsub_token (Chatwoot RoomChannel pattern)
type Handler struct {
	hub           *Hub
	authenticator *wspkg.WSAuthenticator
	upgrader      websocket.Upgrader
}

// NewHandler creates a WebSocket handler with the given hub and authenticator.
// The authenticator provides both JWT (agent) and pubsub_token (contact) auth paths.
func NewHandler(hub *Hub, authenticator *wspkg.WSAuthenticator) *Handler {
	return &Handler{
		hub:           hub,
		authenticator: authenticator,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			// Allow all origins — CORS is handled at the Gin middleware layer
			CheckOrigin: func(r *http.Request) bool { return true },
			// ActionCable clients send Sec-WebSocket-Protocol: actioncable-v1-json.
			// If the server doesn't echo back a supported subprotocol, the JS
			// client immediately closes the connection ("Protocol is unsupported")
			// and enters a reconnect loop.  gorilla/websocket picks the first
			// requested protocol listed here that the client also offered.
			Subprotocols: []string{"actioncable-v1-json"},
		},
	}
}

// ServeWS handles the WebSocket upgrade request at /ws.
// URL: /ws?token=<JWT_ACCESS_TOKEN>  OR  /ws?pubsub_token=<TOKEN>
// On successful upgrade, the handler:
//  1. Authenticates the request via WSAuthenticator (JWT or pubsub_token)
//  2. Authorizes the user/contact for the requested account
//  3. Upgrades HTTP connection to WebSocket
//  4. Creates a Client and registers it with the Hub
//  5. Starts readPump (incoming commands) and writePump (outgoing events) goroutines
func (h *Handler) ServeWS(c *gin.Context) {
	// Step 1: Authenticate (JWT or pubsub_token)
	claims, err := h.authenticator.Authenticate(c)
	if err != nil {
		logger.L().Errorf("ws: authentication failed: %v", err)
		c.JSON(http.StatusUnauthorized, gin.H{"error": err.Error()})
		return
	}

	// Step 2: Authorize (verify account access)
	if err := h.authenticator.Authorize(claims, c); err != nil {
		logger.L().Errorf("ws: authorization failed: %v", err)
		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
		return
	}

	// Step 3: Upgrade HTTP connection to WebSocket
	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.L().Errorf("ws: upgrade failed: %v", err)
		return
	}

	// Step 4: Create client and register with hub
	client := NewClient(claims.UserID, claims.AccountID, conn, h.hub)
	// Populate extended claims fields on the client
	client.Role = claims.Role
	client.IsContact = claims.IsContact
	client.PubsubToken = claims.PubsubToken
	client.ContactID = claims.ContactID
	client.InboxID = claims.InboxID
	client.ClientID = claims.ClientID

	h.hub.Register(client)

	logger.L().Infof("ws: connection established (user=%d, account=%d, is_contact=%v)", claims.UserID, claims.AccountID, claims.IsContact)

	// Send ActionCable welcome frame — the JS client expects this immediately
	// after upgrade.  Without it the client's ConnectionMonitor considers the
	// connection stale and enters a reconnect loop.
	welcomeData, _ := json.Marshal(WelcomeFrame{Type: ServerWelcome})
	client.Send <- welcomeData

	// Start pumps in separate goroutines
	go h.writePump(client)
	go h.readPump(client)
}

// ServeCable handles the WebSocket upgrade request at /cable.
// This is the ActionCable-compatible endpoint (Chatwoot uses /cable for WS).
// It is functionally identical to ServeWS but uses the ActionCable naming convention.
// Some Chatwoot clients specifically connect to /cable.
func (h *Handler) ServeCable(c *gin.Context) {
	h.ServeWS(c)
}

// readPump reads messages from the WebSocket connection and dispatches commands.
// Reference: Chatwoot ActionCable consumer — processes subscribe/unsubscribe/ping commands.
// One readPump runs per client connection.
func (h *Handler) readPump(client *Client) {
	defer func() {
		h.hub.Unregister(client)
		client.Conn.Close()
	}()

	client.Conn.SetReadLimit(MaxMessageSize)
	if err := client.Conn.SetReadDeadline(time.Now().Add(PongWait)); err != nil {
		logger.L().Errorf("ws: set initial read deadline for user=%d: %v", client.UserID, err)
		return
	}
	client.Conn.SetPongHandler(func(string) error {
		return client.Conn.SetReadDeadline(time.Now().Add(PongWait))
	})

	for {
		_, message, err := client.Conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				logger.L().Errorf("ws: unexpected close for user=%d: %v", client.UserID, err)
			}
			break
		}
		if err := h.validateClient(client); err != nil {
			logger.L().Infof("ws: closing invalid agent connection for user=%d: %v", client.UserID, err)
			break
		}

		// Decode the command frame
		var cmd CommandFrame
		if err := json.Unmarshal(message, &cmd); err != nil {
			logger.L().Warnf("ws: invalid command from user=%d: %v", client.UserID, err)
			continue
		}

		// Dispatch command
		switch cmd.Command {
		case CommandSubscribe:
			h.handleSubscribe(client, cmd)
		case CommandUnsubscribe:
			h.handleUnsubscribe(client, cmd)
		case CommandPing:
			h.handlePing(client)
		case CommandMessage:
			// ActionCable "message" command — client performs a channel action
			// (e.g. update_presence).  We acknowledge but don't require a
			// specific handler for presence yet.
			logger.L().Debugf("ws: message command from user=%d, data=%s", client.UserID, cmd.Data)
		default:
			logger.L().Warnf("ws: unknown command '%s' from user=%d", cmd.Command, client.UserID)
		}
	}
}

// writePump sends messages from the client's Send channel to the WebSocket connection.
// It also sends periodic ping frames for heartbeat detection.
// One writePump runs per client connection.
func (h *Handler) writePump(client *Client) {
	ticker := time.NewTicker(time.Duration(PingInterval) * time.Second)
	defer func() {
		ticker.Stop()
		client.Conn.Close()
	}()

	for {
		select {
		case message, ok := <-client.Send:
			if err := client.Conn.SetWriteDeadline(time.Now().Add(WriteWait)); err != nil {
				logger.L().Errorf("ws: set write deadline for user=%d: %v", client.UserID, err)
				return
			}
			if !ok {
				// Hub closed the channel — send close frame
				if err := client.Conn.WriteMessage(websocket.CloseMessage, []byte{}); err != nil {
					logger.L().Debugf("ws: close frame for user=%d: %v", client.UserID, err)
				}
				return
			}

			// Write text message (all our frames are JSON text)
			if err := client.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
				logger.L().Errorf("ws: write error for user=%d: %v", client.UserID, err)
				return
			}

		case <-ticker.C:
			if err := h.validateClient(client); err != nil {
				logger.L().Infof("ws: closing invalid agent connection for user=%d: %v", client.UserID, err)
				return
			}
			// Send WebSocket protocol-level ping control frame.
			// The browser automatically responds with a Pong, which triggers
			// the PongHandler in readPump and resets the ReadDeadline.
			// Without this, the ReadDeadline (PongWait=60s) expires and the
			// connection is forcibly closed, causing the client to show
			// "offline" / "reconnecting" notifications every ~60 seconds.
			if err := client.Conn.SetWriteDeadline(time.Now().Add(WriteWait)); err != nil {
				logger.L().Errorf("ws: set ping deadline for user=%d: %v", client.UserID, err)
				return
			}
			if err := client.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				logger.L().Errorf("ws: ws-ping control frame failed for user=%d: %v", client.UserID, err)
				return
			}

			// Also send ActionCable-level ping message (JSON text frame).
			// The JS ConnectionMonitor expects periodic ping messages to
			// keep the connection alive (staleThreshold = 6s by default).
			pingMsg, err := json.Marshal(PingFrame{
				Type:    ServerPing,
				Message: time.Now().UTC().Format(time.RFC3339),
			})
			if err != nil {
				logger.L().Errorf("ws: marshal ActionCable ping for user=%d: %v", client.UserID, err)
				return
			}
			if err := client.Conn.SetWriteDeadline(time.Now().Add(WriteWait)); err != nil {
				logger.L().Errorf("ws: set ActionCable ping deadline for user=%d: %v", client.UserID, err)
				return
			}
			if err := client.Conn.WriteMessage(websocket.TextMessage, pingMsg); err != nil {
				logger.L().Errorf("ws: actioncable ping write failed for user=%d: %v", client.UserID, err)
				return
			}
		}
	}
}

func (h *Handler) validateClient(client *Client) error {
	if client.IsContact || h.authenticator == nil {
		return nil
	}
	return h.authenticator.ValidateAgentAccess(context.Background(), client.UserID, client.ClientID)
}

// handleSubscribe processes a subscribe command.
// Validates the ChannelIdentifier and adds the client to the appropriate room.
func (h *Handler) handleSubscribe(client *Client, cmd CommandFrame) {
	var identifier ChannelIdentifier
	if err := json.Unmarshal([]byte(cmd.Identifier), &identifier); err != nil {
		logger.L().Warnf("ws: invalid identifier from user=%d: %v", client.UserID, err)
		// Send reject frame
		rejectData, _ := json.Marshal(RejectFrame{
			Type:       ServerRejectSubscribe,
			Identifier: cmd.Identifier,
			Reason:     "invalid identifier format",
		})
		client.Send <- rejectData
		return
	}

	room := ""
	if client.IsContact {
		if identifier.Channel != ChannelRoom || identifier.PubsubToken == "" || identifier.PubsubToken != client.PubsubToken {
			rejectData, _ := json.Marshal(RejectFrame{
				Type:       ServerRejectSubscribe,
				Identifier: cmd.Identifier,
				Reason:     "invalid contact RoomChannel subscription",
			})
			client.Send <- rejectData
			return
		}
		room = pubsubTokenRoomName(identifier.PubsubToken)
	} else if identifier.AccountID != client.AccountID {
		rejectData, _ := json.Marshal(RejectFrame{
			Type:       ServerRejectSubscribe,
			Identifier: cmd.Identifier,
			Reason:     "account_id mismatch",
		})
		client.Send <- rejectData
		return
	}

	// Determine room name based on channel type (uses Hub's canonical naming).
	// Contact RoomChannel was resolved above from its authenticated token.
	switch {
	case room != "":
	case identifier.Channel == ChannelAccount || identifier.Channel == ChannelRoom:
		// RoomChannel is Chatwoot's single-subscription model — it maps
		// to the account room (all account-level events are delivered).
		room = accountRoomName(identifier.AccountID)
	case identifier.Channel == ChannelConversation:
		if identifier.ConversationID == 0 {
			rejectData, _ := json.Marshal(RejectFrame{
				Type:       ServerRejectSubscribe,
				Identifier: cmd.Identifier,
				Reason:     "conversation_id required for ConversationChannel",
			})
			client.Send <- rejectData
			return
		}
		room = conversationRoomName(identifier.AccountID, identifier.ConversationID)
	default:
		rejectData, _ := json.Marshal(RejectFrame{
			Type:       ServerRejectSubscribe,
			Identifier: cmd.Identifier,
			Reason:     "unknown channel type: " + identifier.Channel,
		})
		client.Send <- rejectData
		return
	}

	// Subscribe the client to the room
	client.Subscribe(room)

	// Store the ActionCable subscription identifier on the client so that
	// event frames can be wrapped with it for correct client-side routing.
	client.Identifier = cmd.Identifier

	// Send confirmation frame
	confirmData, _ := json.Marshal(ConfirmFrame{
		Type:       ServerConfirmSubscribe,
		Identifier: cmd.Identifier,
	})
	client.Send <- confirmData
}

// handleUnsubscribe processes an unsubscribe command.
func (h *Handler) handleUnsubscribe(client *Client, cmd CommandFrame) {
	var identifier ChannelIdentifier
	if err := json.Unmarshal([]byte(cmd.Identifier), &identifier); err != nil {
		logger.L().Warnf("ws: invalid identifier in unsubscribe from user=%d: %v", client.UserID, err)
		return
	}

	// Determine room name (uses Hub's canonical naming)
	room := ""
	switch identifier.Channel {
	case ChannelAccount, ChannelRoom:
		room = accountRoomName(identifier.AccountID)
	case ChannelConversation:
		room = conversationRoomName(identifier.AccountID, identifier.ConversationID)
	default:
		return
	}

	client.Unsubscribe(room)

	confirmData, _ := json.Marshal(ConfirmFrame{
		Type:       ServerConfirmUnsubscribe,
		Identifier: cmd.Identifier,
	})
	client.Send <- confirmData
}

// handlePing responds to a client-initiated ping command with a pong frame.
func (h *Handler) handlePing(client *Client) {
	pingData, _ := json.Marshal(PingFrame{
		Type:    ServerPing,
		Message: time.Now().UTC().Format(time.RFC3339),
	})
	client.Send <- pingData
}