services/socket_service.go

package services

import (
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

type SocketService interface {
	HandleConnection(w http.ResponseWriter, r *http.Request, userID uuid.UUID)
	EmitEvent(userID uuid.UUID, eventType string, payload interface{})
	BroadcastEvent(eventType string, payload interface{})
}

type socketService struct {
	clients    map[uuid.UUID][]*websocket.Conn
	clientsMux sync.RWMutex
	upgrader   websocket.Upgrader
}

func NewSocketService() SocketService {
	return &socketService{
		clients: make(map[uuid.UUID][]*websocket.Conn),
		upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			CheckOrigin: func(r *http.Request) bool {
				return true // Development mode
			},
		},
	}
}

func (s *socketService) HandleConnection(w http.ResponseWriter, r *http.Request, userID uuid.UUID) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Failed to upgrade connection: %v", err)
		return
	}

	s.registerClient(userID, conn)
	defer s.unregisterClient(userID, conn)

	// Keep connection alive and listen for disconnects
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error: %v", err)
			}
			break
		}
	}
}

func (s *socketService) registerClient(userID uuid.UUID, conn *websocket.Conn) {
	s.clientsMux.Lock()
	defer s.clientsMux.Unlock()

	s.clients[userID] = append(s.clients[userID], conn)
	log.Printf("Client connected: %s", userID)
}

func (s *socketService) unregisterClient(userID uuid.UUID, conn *websocket.Conn) {
	s.clientsMux.Lock()
	defer s.clientsMux.Unlock()

	conn.Close()

	if conns, ok := s.clients[userID]; ok {
		for i, c := range conns {
			if c == conn {
				s.clients[userID] = append(conns[:i], conns[i+1:]...)
				break
			}
		}
		if len(s.clients[userID]) == 0 {
			delete(s.clients, userID)
		}
	}
	log.Printf("Client disconnected: %s", userID)
}

func (s *socketService) EmitEvent(userID uuid.UUID, eventType string, payload interface{}) {
	s.clientsMux.RLock()
	conns, ok := s.clients[userID]
	s.clientsMux.RUnlock()

	if !ok {
		return
	}

	message := map[string]interface{}{
		"type":    eventType,
		"payload": payload,
		"time":    time.Now(),
	}

	for _, conn := range conns {
		err := conn.WriteJSON(message)
		if err != nil {
			log.Printf("Error sending message to user %s: %v", userID, err)
			conn.Close()
		}
	}
}

func (s *socketService) BroadcastEvent(eventType string, payload interface{}) {
	s.clientsMux.RLock()
	defer s.clientsMux.RUnlock()

	message := map[string]interface{}{
		"type":    eventType,
		"payload": payload,
		"time":    time.Now(),
	}

	for userID, conns := range s.clients {
		for _, conn := range conns {
			err := conn.WriteJSON(message)
			if err != nil {
				log.Printf("Error broadcasting to user %s: %v", userID, err)
			}
		}
	}
}