core/server/events/ws.go

package events

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

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

	"armature/metrics"
)

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
	upgrader websocket.Upgrader
}

// 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.
// allowedOrigins controls the WebSocket upgrader's CheckOrigin behavior.
// Empty or "*" allows all origins (dev/test). Comma-separated list
// restricts to those origins (production).
func NewHub(bus *Bus, allowedOrigins string) *Hub {
	h := &Hub{
		bus:   bus,
		conns: make(map[string]*Conn),
	}
	h.upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin:     h.makeCheckOrigin(allowedOrigins),
	}
	return h
}

// makeCheckOrigin returns a CheckOrigin function that validates the
// request Origin header against the configured allowed origins.
func (h *Hub) makeCheckOrigin(allowedOrigins string) func(r *http.Request) bool {
	origins := strings.TrimSpace(allowedOrigins)

	// No restriction: dev/test or explicit "*"
	if origins == "" || origins == "*" {
		return func(r *http.Request) bool { return true }
	}

	// Parse comma-separated origin list
	var allowed []string
	for _, o := range strings.Split(origins, ",") {
		o = strings.TrimSpace(o)
		if o != "" {
			allowed = append(allowed, o)
		}
	}

	return func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		if origin == "" {
			// No Origin header — same-origin request (non-browser client).
			return true
		}
		for _, a := range allowed {
			if a == origin {
				return true
			}
		}
		log.Printf("[ws] rejected WebSocket from origin %q (allowed: %s)", origin, origins)
		return false
	}
}

// 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 := h.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())
	metrics.WebSocketConnections.Inc()

	// 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
}

// Broadcast sends an event to all connected clients via the bus.
// No TargetUserID — every WS subscriber receives it.
func (h *Hub) Broadcast(event Event) {
	h.bus.Publish(event)
}

// PublishToUser sends an event to a specific user via the bus.
// All replicas receive the event; only the one with the user's connection delivers it.
func (h *Hub) PublishToUser(userID string, event Event) {
	event.TargetUserID = userID
	h.bus.Publish(event)
}

// sendToUserLocal sends an event directly to local WebSocket connections for a user.
// Unexported — callers should use PublishToUser for cross-pod delivery.
func (h *Hub) sendToUserLocal(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())
	metrics.WebSocketConnections.Dec()

	// 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
		}

		// Targeted events skip room filtering (user-scoped, not room-scoped).
		if e.TargetUserID != "" && e.TargetUserID != c.userID {
			return
		}

		// must not be forwarded to WebSocket clients — they sent it.
		if strings.HasPrefix(e.Label, "tool.result.") {
			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
		}

		if event.Label == "room.subscribe" {
			var req struct {
				Room string `json:"room"`
			}
			if json.Unmarshal(event.Payload, &req) == nil && req.Room != "" && len(c.rooms) < 100 {
				c.JoinRoom(req.Room)
			}
			continue
		}
		if event.Label == "room.unsubscribe" {
			var req struct {
				Room string `json:"room"`
			}
			if json.Unmarshal(event.Payload, &req) == nil && req.Room != "" {
				c.LeaveRoom(req.Room)
			}
			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
}