photoplaces/backend/internal/handlers/websocket.go

// Package handlers
package handlers

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

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

type VisitorDot struct {
	UserID string  `json:"user_id,omitempty"`
	Lat    float64 `json:"lat"`
	Lng    float64 `json:"lng"`
}

type WSHub struct {
	mu      sync.RWMutex
	clients map[*websocket.Conn]VisitorDot
}

func NewWSHub() *WSHub {
	return &WSHub{clients: make(map[*websocket.Conn]VisitorDot)}
}

func (h *WSHub) HandleWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("ws upgrade: %v", err)
		return
	}

	dot := VisitorDot{}

	h.mu.Lock()
	h.clients[conn] = dot
	h.mu.Unlock()

	defer func() {
		h.mu.Lock()
		delete(h.clients, conn)
		h.mu.Unlock()
		conn.Close()
	}()

	for {
		var msg VisitorDot
		if err := conn.ReadJSON(&msg); err != nil {
			break
		}

		h.mu.Lock()
		h.clients[conn] = msg
		visitors := make([]VisitorDot, 0, len(h.clients))
		for _, v := range h.clients {
			visitors = append(visitors, v)
		}

		for c := range h.clients {
			if err := c.WriteJSON(map[string]interface{}{"visitors": visitors}); err != nil {
				c.Close()
				delete(h.clients, c)
			}
		}
		h.mu.Unlock()
	}
}