Turbine/pkg/ring/websocket/controller.go

package websocket

import (
	"fmt"

	"github.com/gofiber/contrib/v3/websocket"
	"github.com/google/uuid"
	"github.com/rs/zerolog/log"
	"github.com/samber/lo"

	pb "git.solsynth.dev/goatworks/turbine/pkg/shared/proto/gen"
)

// WebSocketController handles the WebSocket endpoint.
type WebSocketController struct {
	Manager *Manager
}

// NewWebSocketController creates a new WebSocketController.
func NewWebSocketController(manager *Manager) *WebSocketController {
	return &WebSocketController{
		Manager: manager,
	}
}

// HandleWebSocket is the main handler for the /ws endpoint.
func (wc *WebSocketController) HandleWebSocket(c *websocket.Conn) {
	// Mock Authentication for now based on C# example
	// In a real scenario, this would involve JWT verification, session lookup, etc.
	// For demonstration, we'll assume a dummy user and session.
	// The C# code uses HttpContext.Items to get CurrentUser and CurrentSession.
	// In Go Fiber, we can pass this through Locals or use middleware to set it up.
	// For now, let's create a dummy user and session.

	// TODO: Replace with actual authentication logic
	// For now, assume a dummy account and session
	// Based on the C# code, CurrentUser and CurrentSession are pb.Account and pb.AuthSession
	dummyAccount := &pb.Account{
		Id:   uuid.New().String(),
		Name: "dummy_user",
		Nick: "Dummy User",
	}
	dummySession := &pb.AuthSession{
		ClientId: lo.ToPtr(uuid.New().String()), // This is used as deviceId if not present
	}

	// Device ID handling
	deviceAlt := c.Query("deviceAlt")
	if deviceAlt != "" {
		allowedDeviceAlternative := []string{"watch"} // Hardcoded for now
		found := false
		for _, alt := range allowedDeviceAlternative {
			if deviceAlt == alt {
				found = true
				break
			}
		}
		if !found {
			log.Warn().Msgf("Unsupported device alternative: %s", deviceAlt)
			bytes, _ := newErrorPacket("Unsupported device alternative: " + deviceAlt).ToBytes()
			c.WriteMessage(websocket.BinaryMessage, bytes)
			c.Close()
			return
		}
	}

	accountID := uuid.MustParse(dummyAccount.Id)
	deviceIDStr := ""
	if dummySession.ClientId == nil {
		deviceIDStr = uuid.New().String()
	} else {
		deviceIDStr = *dummySession.ClientId
	}
	if deviceAlt != "" {
		deviceIDStr = fmt.Sprintf("%s+%s", deviceIDStr, deviceAlt)
	}

	// Setup connection context for cancellation
	cancel := func() {} // Placeholder
	connectionKey := ConnectionKey{AccountID: accountID, DeviceID: deviceIDStr}

	// Add connection to manager
	if !wc.Manager.TryAdd(connectionKey, c, cancel) {
		bytes, _ := newErrorPacket("Too many connections from the same device and account.").ToBytes()
		c.WriteMessage(websocket.BinaryMessage, bytes)
		c.Close()
		return
	}

	log.Debug().Msgf("Connection established with user @%s#%s and device #%s", dummyAccount.Name, dummyAccount.Id, deviceIDStr)

	// Publish WebSocket connected event
	wc.Manager.PublishWebSocketConnectedEvent(accountID, deviceIDStr, false) // isOffline is false on connection

	defer func() {
		wc.Manager.Disconnect(connectionKey, "Client disconnected.")
		// Publish WebSocket disconnected event
		isOffline := !wc.Manager.GetAccountIsConnected(accountID) // Check if account is completely offline
		wc.Manager.PublishWebSocketDisconnectedEvent(accountID, deviceIDStr, isOffline)
		log.Debug().Msgf("Connection disconnected with user @%s#%s and device #%s", dummyAccount.Name, dummyAccount.Id, deviceIDStr)
	}()

	// Main event loop
	for {
		mt, msg, err := c.ReadMessage()
		if err != nil {
			log.Error().Err(err).Msg("WebSocket read error")
			break
		}

		if mt == websocket.CloseMessage {
			log.Info().Msg("Received close message from client")
			break
		}

		if mt != websocket.BinaryMessage {
			log.Warn().Msgf("Received non-binary message type: %d", mt)
			continue
		}

		packet, err := FromBytes(msg)
		if err != nil {
			log.Error().Err(err).Msg("Failed to deserialize WebSocket packet")
			bytes, _ := newErrorPacket("Failed to deserialize packet").ToBytes()
			c.WriteMessage(websocket.BinaryMessage, bytes)
			continue
		}

		if err := wc.Manager.HandlePacket(dummyAccount, deviceIDStr, packet, c); err != nil {
			log.Error().Err(err).Msg("Failed to handle incoming WebSocket packet")
			bytes, _ := newErrorPacket("Failed to process packet").ToBytes()
			c.WriteMessage(websocket.BinaryMessage, bytes)
			// Depending on error, might want to close connection
		}
	}
}

// newErrorPacket creates a new WebSocketPacket with an error type.
func newErrorPacket(message string) *WebSocketPacket {
	return &WebSocketPacket{
		Type:         WebSocketPacketTypeError,
		ErrorMessage: message,
	}
}