GopherGate/internal/server/websocket.go

package server

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

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

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // In production, refine this
	},
}

type Hub struct {
	clients     map[*websocket.Conn]bool
	broadcast   chan interface{}
	register    chan *websocket.Conn
	unregister  chan *websocket.Conn
	mu          sync.Mutex
	clientCount int32
}

func NewHub() *Hub {
	return &Hub{
		clients:    make(map[*websocket.Conn]bool),
		broadcast:  make(chan interface{}),
		register:   make(chan *websocket.Conn),
		unregister: make(chan *websocket.Conn),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client] = true
			atomic.AddInt32(&h.clientCount, 1)
			h.mu.Unlock()
			log.Println("WebSocket client registered")
		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				client.Close()
				atomic.AddInt32(&h.clientCount, -1)
			}
			h.mu.Unlock()
			log.Println("WebSocket client unregistered")
		case message := <-h.broadcast:
			h.mu.Lock()
			for client := range h.clients {
				err := client.WriteJSON(message)
				if err != nil {
					log.Printf("WebSocket error: %v", err)
					client.Close()
					delete(h.clients, client)
					atomic.AddInt32(&h.clientCount, -1)
				}
			}
			h.mu.Unlock()
		}
	}
}

func (h *Hub) GetClientCount() int {
	return int(atomic.LoadInt32(&h.clientCount))
}

func (s *Server) handleWebSocket(c *gin.Context) {
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("Failed to set websocket upgrade: %v", err)
		return
	}

	s.hub.register <- conn

	defer func() {
		s.hub.unregister <- conn
	}()

	// Initial message
	conn.WriteJSON(gin.H{
		"type":    "connected",
		"message": "Connected to LLM Proxy Dashboard",
	})

	for {
		var msg map[string]interface{}
		err := conn.ReadJSON(&msg)
		if err != nil {
			break
		}
		
		if msg["type"] == "ping" {
			conn.WriteJSON(gin.H{"type": "pong", "payload": gin.H{}})
		}
	}
}