core/server/events/ws.go

package events

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

	"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 },
}

const (
	writeWait  = 10 * time.Second
	pongWait   = 60 * time.Second
	pingPeriod = (pongWait * 9) / 10
	maxMsgSize = 4096
)

// Hub manages WebSocket connections and bridges them to the Bus.
type Hub struct {
	bus   *Bus
	mu    sync.RWMutex
	conns map[string]*Conn // connID → Conn
}

// Conn represents a single WebSocket connection.
type Conn struct {
	id     string
	userID string
	rooms  map[string]bool // rooms this connection is subscribed to
	ws     *websocket.Conn
	send   chan []byte
	hub    *Hub
	unsubs []func() // bus unsubscribe functions
}

// NewHub creates a Hub bound to an event bus.
func NewHub(bus *Bus) *Hub {
	return &Hub{
		bus:   bus,
		conns: make(map[string]*Conn),
	}
}

// HandleWebSocket is a Gin handler for WebSocket upgrade.
// Expects userID set in context (by auth middleware or query param).
func (h *Hub) HandleWebSocket(c *gin.Context) {
	// Auth: check token from query param
	userID := c.GetString("user_id")
	if userID == "" {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication required"})
		return
	}

	ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Printf("[ws] upgrade failed: %v", err)
		return
	}

	connID := userID + "-" + time.Now().Format("150405.000")
	conn := &Conn{
		id:     connID,
		userID: userID,
		rooms:  make(map[string]bool),
		ws:     ws,
		send:   make(chan []byte, 64),
		hub:    h,
	}

	h.mu.Lock()
	h.conns[connID] = conn
	h.mu.Unlock()

	log.Printf("[ws] connected: %s (user=%s, total=%d)", connID, userID, h.ConnCount())

	// Subscribe this connection to bus events destined for clients
	conn.subscribeToBus()

	// Publish presence
	h.bus.Publish(Event{
		Label:    "user.presence",
		Payload:  mustJSON(map[string]any{"user_id": userID, "status": "online"}),
		Ts:       time.Now().UnixMilli(),
		SenderID: userID,
	})

	// Start read/write pumps
	go conn.writePump()
	go conn.readPump()
}

// ConnCount returns the number of active connections.
func (h *Hub) ConnCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.conns)
}

// ConnsByUser returns all connections for a given user.
func (h *Hub) ConnsByUser(userID string) []*Conn {
	h.mu.RLock()
	defer h.mu.RUnlock()
	var result []*Conn
	for _, c := range h.conns {
		if c.userID == userID {
			result = append(result, c)
		}
	}
	return result
}

// IsConnected returns true if the user has at least one active WebSocket.
func (h *Hub) IsConnected(userID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, c := range h.conns {
		if c.userID == userID {
			return true
		}
	}
	return false
}

// SendToUser sends an event directly to all WebSocket connections for a user,
// bypassing room filtering. Used for targeted tool.call delivery.
func (h *Hub) SendToUser(userID string, event Event) {
	data, err := json.Marshal(event)
	if err != nil {
		return
	}
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, c := range h.conns {
		if c.userID == userID {
			select {
			case c.send <- data:
			default:
			}
		}
	}
}

// GetBus returns the underlying event bus.
func (h *Hub) GetBus() *Bus {
	return h.bus
}

// removeConn cleans up a connection.
func (h *Hub) removeConn(conn *Conn) {
	h.mu.Lock()
	delete(h.conns, conn.id)
	h.mu.Unlock()

	// Unsubscribe from bus
	for _, unsub := range conn.unsubs {
		unsub()
	}

	log.Printf("[ws] disconnected: %s (total=%d)", conn.id, h.ConnCount())

	// Publish presence offline (only if no other conns for this user)
	if len(h.ConnsByUser(conn.userID)) == 0 {
		h.bus.Publish(Event{
			Label:    "user.presence",
			Payload:  mustJSON(map[string]any{"user_id": conn.userID, "status": "offline"}),
			Ts:       time.Now().UnixMilli(),
			SenderID: conn.userID,
		})
	}
}

// subscribeToBus registers this conn as a bus subscriber for client-facing events.
func (c *Conn) subscribeToBus() {
	unsub := c.hub.bus.Subscribe("*", func(e Event) {
		// Only forward events destined for clients
		if !ShouldSendToClient(e.Label) {
			return
		}

		// Don't echo typing events back to the sender
		if e.Label == "chat.typing" || e.ConnID == c.id {
			if e.SenderID == c.userID && e.ConnID == c.id {
				return
			}
		}

		// Room filtering: if event has a room, only send if conn is in that room
		if e.Room != "" && !c.rooms[e.Room] {
			return
		}

		data, err := json.Marshal(e)
		if err != nil {
			return
		}

		select {
		case c.send <- data:
		default:
			// send buffer full, drop event
		}
	})

	c.unsubs = append(c.unsubs, unsub)
}

// JoinRoom adds this connection to a named room.
func (c *Conn) JoinRoom(room string) {
	c.rooms[room] = true
}

// LeaveRoom removes this connection from a named room.
func (c *Conn) LeaveRoom(room string) {
	delete(c.rooms, room)
}

// ── Read/Write Pumps ─────────────────────────

func (c *Conn) readPump() {
	defer func() {
		c.hub.removeConn(c)
		c.ws.Close()
	}()

	c.ws.SetReadLimit(maxMsgSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error {
		c.ws.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	for {
		_, message, err := c.ws.ReadMessage()
		if err != nil {
			break
		}

		var event Event
		if err := json.Unmarshal(message, &event); err != nil {
			continue
		}

		// Handle ping
		if event.Label == "ping" {
			pong, _ := json.Marshal(Event{Label: "pong", Ts: time.Now().UnixMilli()})
			select {
			case c.send <- pong:
			default:
			}
			continue
		}

		// Validate: only accept events allowed from clients
		if !ShouldAcceptFromClient(event.Label) {
			continue
		}

		// Tag with sender info
		event.SenderID = c.userID
		event.ConnID = c.id
		if event.Ts == 0 {
			event.Ts = time.Now().UnixMilli()
		}

		// Publish into bus
		c.hub.bus.Publish(event)
	}
}

func (c *Conn) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()

	for {
		select {
		case msg, ok := <-c.send:
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				c.ws.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.ws.WriteMessage(websocket.TextMessage, msg); err != nil {
				return
			}

		case <-ticker.C:
			c.ws.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.ws.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// ── Helpers ──────────────────────────────────

func mustJSON(v any) json.RawMessage {
	data, err := json.Marshal(v)
	if err != nil {
		return json.RawMessage(`{}`)
	}
	return data
}