MindSift/internal/ws/hub.go

package ws

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

	"github.com/gorilla/websocket"
	"github.com/sandrodach/mindsift/internal/domain"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  2048,
	WriteBufferSize: 2048,
	CheckOrigin: func(r *http.Request) bool {
		return true // Self-hosted reverse proxy compatible
	},
}

type Client struct {
	hub     *Hub
	conn    *websocket.Conn
	spaceID string
	userID  string
	send    chan []byte
}

type docLockItem struct {
	docID     string
	userID    string
	userName  string
	lockedAt  time.Time
	expiresAt time.Time
}

type Hub struct {
	mu         sync.RWMutex
	spaceRooms map[string]map[*Client]bool
	userRooms  map[string]map[*Client]bool
	allClients map[*Client]bool
	locks      map[string]*docLockItem
}

func NewHub() *Hub {
	h := &Hub{
		spaceRooms: make(map[string]map[*Client]bool),
		userRooms:  make(map[string]map[*Client]bool),
		allClients: make(map[*Client]bool),
		locks:      make(map[string]*docLockItem),
	}
	go h.cleanupLocksLoop()
	return h
}

func (h *Hub) Register(c *Client) {
	h.mu.Lock()
	h.allClients[c] = true
	if c.spaceID != "" {
		if h.spaceRooms[c.spaceID] == nil {
			h.spaceRooms[c.spaceID] = make(map[*Client]bool)
		}
		h.spaceRooms[c.spaceID][c] = true
	}
	if c.userID != "" && c.userID != "anonymous" {
		if h.userRooms[c.userID] == nil {
			h.userRooms[c.userID] = make(map[*Client]bool)
		}
		h.userRooms[c.userID][c] = true
	}
	onlineCount := len(h.allClients)
	h.mu.Unlock()

	if c.spaceID != "" {
		h.Broadcast(c.spaceID, "", "USER_PRESENCE", map[string]interface{}{
			"onlineCount": onlineCount,
		})
	}
}

func (h *Hub) Unregister(c *Client) {
	h.mu.Lock()
	delete(h.allClients, c)
	if c.spaceID != "" {
		if room, ok := h.spaceRooms[c.spaceID]; ok {
			delete(room, c)
			if len(room) == 0 {
				delete(h.spaceRooms, c.spaceID)
			}
		}
	}
	userRemaining := 0
	if c.userID != "" && c.userID != "anonymous" {
		if uRoom, ok := h.userRooms[c.userID]; ok {
			delete(uRoom, c)
			userRemaining = len(uRoom)
			if userRemaining == 0 {
				delete(h.userRooms, c.userID)
			}
		}
	}
	close(c.send)

	// If this user has no other active client connections, release their document locks
	var releasedDocIDs []string
	if userRemaining == 0 && c.userID != "" {
		now := time.Now()
		for docID, l := range h.locks {
			if l.userID == c.userID || l.expiresAt.Before(now) {
				releasedDocIDs = append(releasedDocIDs, docID)
				delete(h.locks, docID)
			}
		}
	}
	h.mu.Unlock()

	// Broadcast unlock events for released locks
	for _, docID := range releasedDocIDs {
		h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID})
	}
}

// Document Locking Methods

func (h *Hub) AcquireDocLock(docID, userID, userName string, ttl time.Duration) (*domain.DocLockInfo, bool, *domain.DocLockInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()

	now := time.Now()
	existing, exists := h.locks[docID]
	if exists && existing.expiresAt.After(now) {
		if existing.userID != userID {
			// Locked by another user
			info := &domain.DocLockInfo{
				DocID:     existing.docID,
				UserID:    existing.userID,
				UserName:  existing.userName,
				ExpiresAt: existing.expiresAt.Format(time.RFC3339),
			}
			return nil, false, info
		}
		// Same user extending lock
		existing.expiresAt = now.Add(ttl)
		if userName != "" {
			existing.userName = userName
		}
		info := &domain.DocLockInfo{
			DocID:     existing.docID,
			UserID:    existing.userID,
			UserName:  existing.userName,
			ExpiresAt: existing.expiresAt.Format(time.RFC3339),
		}
		return info, true, nil
	}

	// Create fresh lock
	item := &docLockItem{
		docID:     docID,
		userID:    userID,
		userName:  userName,
		lockedAt:  now,
		expiresAt: now.Add(ttl),
	}
	h.locks[docID] = item
	info := &domain.DocLockInfo{
		DocID:     item.docID,
		UserID:    item.userID,
		UserName:  item.userName,
		ExpiresAt: item.expiresAt.Format(time.RFC3339),
	}
	return info, true, nil
}

func (h *Hub) HeartbeatDocLock(docID, userID string, ttl time.Duration) bool {
	h.mu.Lock()
	defer h.mu.Unlock()

	now := time.Now()
	item, ok := h.locks[docID]
	if !ok || item.userID != userID || item.expiresAt.Before(now) {
		return false
	}
	item.expiresAt = now.Add(ttl)
	return true
}

func (h *Hub) ReleaseDocLock(docID, userID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()

	item, ok := h.locks[docID]
	if !ok {
		return false
	}
	if item.userID == userID || userID == "admin" {
		delete(h.locks, docID)
		return true
	}
	return false
}

func (h *Hub) GetDocLock(docID string) *domain.DocLockInfo {
	h.mu.RLock()
	defer h.mu.RUnlock()

	item, ok := h.locks[docID]
	if !ok || item.expiresAt.Before(time.Now()) {
		return nil
	}
	return &domain.DocLockInfo{
		DocID:     item.docID,
		UserID:    item.userID,
		UserName:  item.userName,
		ExpiresAt: item.expiresAt.Format(time.RFC3339),
	}
}

func (h *Hub) cleanupLocksLoop() {
	ticker := time.NewTicker(4 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		h.mu.Lock()
		now := time.Now()
		var expiredDocIDs []string
		for docID, l := range h.locks {
			if l.expiresAt.Before(now) {
				expiredDocIDs = append(expiredDocIDs, docID)
				delete(h.locks, docID)
			}
		}
		h.mu.Unlock()

		for _, docID := range expiredDocIDs {
			h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID})
		}
	}
}

// Broadcasting Methods

func (h *Hub) BroadcastToAll(event string, payload interface{}) {
	msg := domain.WSMessage{
		Event:     event,
		Payload:   payload,
		Timestamp: domain.NowISO(),
	}
	data, err := json.Marshal(msg)
	if err != nil {
		return
	}

	h.mu.RLock()
	defer h.mu.RUnlock()

	for client := range h.allClients {
		select {
		case client.send <- data:
		default:
		}
	}
}

func (h *Hub) BroadcastToUser(userID, event string, payload interface{}) {
	msg := domain.WSMessage{
		Event:     event,
		Payload:   payload,
		Timestamp: domain.NowISO(),
	}
	data, err := json.Marshal(msg)
	if err != nil {
		return
	}

	h.mu.RLock()
	defer h.mu.RUnlock()

	if room, ok := h.userRooms[userID]; ok {
		for client := range room {
			select {
			case client.send <- data:
			default:
			}
		}
	}
}

func (h *Hub) BroadcastTaskEvent(userID, event string, payload interface{}) {
	h.BroadcastToUser(userID, event, payload)
}

func (h *Hub) BroadcastDocEvent(docID, event string, payload interface{}) {
	h.BroadcastToAll(event, payload)
}

func (h *Hub) Broadcast(spaceID, listID, event string, payload interface{}) {
	msg := domain.WSMessage{
		Event:     event,
		SpaceID:   spaceID,
		ListID:    listID,
		Payload:   payload,
		Timestamp: domain.NowISO(),
	}
	data, err := json.Marshal(msg)
	if err != nil {
		return
	}

	h.mu.RLock()
	defer h.mu.RUnlock()

	room, ok := h.spaceRooms[spaceID]
	if !ok {
		return
	}
	for client := range room {
		select {
		case client.send <- data:
		default:
			// Slow client buffer full
		}
	}
}

func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request, spaceID, userID string) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("[ws] upgrade error: %v", err)
		return
	}

	client := &Client{
		hub:     h,
		conn:    conn,
		spaceID: spaceID,
		userID:  userID,
		send:    make(chan []byte, 128),
	}

	h.Register(client)
	go client.writePump()
	go client.readPump()
}

func (c *Client) readPump() {
	defer func() {
		c.hub.Unregister(c)
		_ = c.conn.Close()
	}()

	c.conn.SetReadLimit(16384)
	_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

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

		// Handle client-initiated lock and unlock messages
		var incoming struct {
			Event   string          `json:"event"`
			DocID   string          `json:"docId"`
			Payload json.RawMessage `json:"payload"`
		}
		if json.Unmarshal(message, &incoming) == nil && c.userID != "" && c.userID != "anonymous" {
			docID := incoming.DocID
			if docID == "" && len(incoming.Payload) > 0 {
				var p struct {
					DocID string `json:"docId"`
				}
				_ = json.Unmarshal(incoming.Payload, &p)
				docID = p.DocID
			}

			if docID != "" {
				switch incoming.Event {
				case "LOCK_DOC":
					if info, ok, _ := c.hub.AcquireDocLock(docID, c.userID, "", 30*time.Second); ok && info != nil {
						c.hub.BroadcastDocEvent(docID, "DOC_LOCKED", info)
					}
				case "HEARTBEAT_DOC":
					c.hub.HeartbeatDocLock(docID, c.userID, 30*time.Second)
				case "UNLOCK_DOC":
					if c.hub.ReleaseDocLock(docID, c.userID) {
						c.hub.BroadcastDocEvent(docID, "DOC_UNLOCKED", map[string]string{"docId": docID})
					}
				}
			}
		}
	}
}

func (c *Client) writePump() {
	ticker := time.NewTicker(30 * time.Second)
	defer func() {
		ticker.Stop()
		_ = c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			_ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				_ = c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}