MindSift/internal/ws/hub.go

package ws

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

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

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

type Client struct {
	hub             *Hub
	conn            *websocket.Conn
	roomID          string
	userID          string
	send            chan []byte
	subscribedRooms map[string]bool
	sendMu          sync.RWMutex
	closed          bool
}

func (c *Client) closeSend() {
	c.sendMu.Lock()
	defer c.sendMu.Unlock()
	if !c.closed {
		c.closed = true
		close(c.send)
	}
}

func (c *Client) trySend(data []byte) {
	c.sendMu.RLock()
	defer c.sendMu.RUnlock()
	if c.closed {
		return
	}
	select {
	case c.send <- data:
	default:
	}
}

type SSEClient struct {
	id              string
	userID          string
	roomID          string
	send            chan []byte
	subscribedRooms map[string]bool
	sendMu          sync.RWMutex
	closed          bool
}

func (c *SSEClient) closeSend() {
	c.sendMu.Lock()
	defer c.sendMu.Unlock()
	if !c.closed {
		c.closed = true
		close(c.send)
	}
}

func (c *SSEClient) trySend(data []byte) {
	c.sendMu.RLock()
	defer c.sendMu.RUnlock()
	if c.closed {
		return
	}
	select {
	case c.send <- data:
	default:
	}
}

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

type blockLockItem struct {
	docID     string
	blockID   string
	userID    string
	userName  string
	userColor string
	lockedAt  time.Time
	expiresAt time.Time
}

type liveSessionItem struct {
	docID         string
	initiatorID   string
	initiatorName string
	startedAt     time.Time
	expiresAt     time.Time
}

type DocStore interface {
	LoadDocumentState(docID string) (string, []domain.DocumentBlock, error)
	FlushDocumentState(docID string, title string, blocks []domain.DocumentBlock) error
}

type Hub struct {
	mu           sync.RWMutex
	rooms        map[string]map[*Client]bool
	userRooms    map[string]map[*Client]bool
	allClients   map[*Client]bool
	sseClients   map[*SSEClient]bool
	sseUserRooms map[string]map[*SSEClient]bool
	sseRooms     map[string]map[*SSEClient]bool
	locks        map[string]*docLockItem
	blockLocks   map[string]*blockLockItem   // key: docID + ":" + blockID
	liveSessions map[string]*liveSessionItem // key: docID
	docRooms     map[string]*DocRoom
	docStore     DocStore
	stopCleanup  chan struct{}
	closed       bool
}

func NewHub() *Hub {
	h := &Hub{
		rooms:        make(map[string]map[*Client]bool),
		userRooms:    make(map[string]map[*Client]bool),
		allClients:   make(map[*Client]bool),
		sseClients:   make(map[*SSEClient]bool),
		sseUserRooms: make(map[string]map[*SSEClient]bool),
		sseRooms:     make(map[string]map[*SSEClient]bool),
		locks:        make(map[string]*docLockItem),
		blockLocks:   make(map[string]*blockLockItem),
		liveSessions: make(map[string]*liveSessionItem),
		docRooms:     make(map[string]*DocRoom),
		stopCleanup:  make(chan struct{}),
	}
	go h.cleanupLocksLoop()
	return h
}

func (h *Hub) SetDocStore(store DocStore) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.docStore = store
}

func (h *Hub) Register(c *Client) {
	h.mu.Lock()
	h.allClients[c] = true
	if c.roomID != "" {
		if h.rooms[c.roomID] == nil {
			h.rooms[c.roomID] = make(map[*Client]bool)
		}
		h.rooms[c.roomID][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) + len(h.sseClients)
	h.mu.Unlock()

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

func (h *Hub) Unregister(c *Client) {
	h.mu.Lock()
	delete(h.allClients, c)
	if c.roomID != "" {
		if room, ok := h.rooms[c.roomID]; ok {
			delete(room, c)
			if len(room) == 0 {
				delete(h.rooms, c.roomID)
			}
		}
	}

	var docRoomsLeft []string
	if c.subscribedRooms != nil {
		for rID := range c.subscribedRooms {
			if room, ok := h.rooms[rID]; ok {
				delete(room, c)
				if len(room) == 0 {
					delete(h.rooms, rID)
				}
			}
			if strings.HasPrefix(rID, "doc:") {
				docRoomsLeft = append(docRoomsLeft, strings.TrimPrefix(rID, "doc:"))
			}
		}
	}

	userRemaining := 0
	if c.userID != "" && c.userID != "anonymous" {
		if uRoom, ok := h.userRooms[c.userID]; ok {
			delete(uRoom, c)
			userRemaining += len(uRoom)
			if len(uRoom) == 0 {
				delete(h.userRooms, c.userID)
			}
		}
		if sseRoom, ok := h.sseUserRooms[c.userID]; ok {
			userRemaining += len(sseRoom)
		}
	}
	c.closeSend()

	// 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, "doc_id": docID})
	}

	// Broadcast awareness leave for doc rooms
	for _, docID := range docRoomsLeft {
		h.BroadcastRoom("doc:"+docID, "CRDT_AWARENESS_LEAVE", map[string]interface{}{
			"docId":  docID,
			"userId": c.userID,
		}, nil)
	}
}

func (h *Hub) RegisterSSE(c *SSEClient) {
	h.mu.Lock()
	h.sseClients[c] = true
	if c.roomID != "" {
		if h.sseRooms[c.roomID] == nil {
			h.sseRooms[c.roomID] = make(map[*SSEClient]bool)
		}
		h.sseRooms[c.roomID][c] = true
	}
	if c.userID != "" && c.userID != "anonymous" {
		if h.sseUserRooms[c.userID] == nil {
			h.sseUserRooms[c.userID] = make(map[*SSEClient]bool)
		}
		h.sseUserRooms[c.userID][c] = true
	}
	if c.subscribedRooms != nil {
		for rID := range c.subscribedRooms {
			if h.sseRooms[rID] == nil {
				h.sseRooms[rID] = make(map[*SSEClient]bool)
			}
			h.sseRooms[rID][c] = true
		}
	}
	onlineCount := len(h.allClients) + len(h.sseClients)
	h.mu.Unlock()

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

func (h *Hub) UnregisterSSE(c *SSEClient) {
	h.mu.Lock()
	delete(h.sseClients, c)
	if c.roomID != "" {
		if room, ok := h.sseRooms[c.roomID]; ok {
			delete(room, c)
			if len(room) == 0 {
				delete(h.sseRooms, c.roomID)
			}
		}
	}
	var docRoomsLeft []string
	if c.subscribedRooms != nil {
		for rID := range c.subscribedRooms {
			if room, ok := h.sseRooms[rID]; ok {
				delete(room, c)
				if len(room) == 0 {
					delete(h.sseRooms, rID)
				}
			}
			if strings.HasPrefix(rID, "doc:") {
				docRoomsLeft = append(docRoomsLeft, strings.TrimPrefix(rID, "doc:"))
			}
		}
	}
	userRemaining := 0
	if c.userID != "" && c.userID != "anonymous" {
		if uRoom, ok := h.sseUserRooms[c.userID]; ok {
			delete(uRoom, c)
			userRemaining += len(uRoom)
			if len(uRoom) == 0 {
				delete(h.sseUserRooms, c.userID)
			}
		}
		if wsRoom, ok := h.userRooms[c.userID]; ok {
			userRemaining += len(wsRoom)
		}
	}
	c.closeSend()

	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()

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

	for _, docID := range docRoomsLeft {
		h.BroadcastRoom("doc:"+docID, "CRDT_AWARENESS_LEAVE", map[string]interface{}{
			"docId":  docID,
			"userId": c.userID,
		}, nil)
	}
}

func (h *Hub) JoinSSERoomByID(sseID, roomID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	for client := range h.sseClients {
		if client.id == sseID {
			if h.sseRooms[roomID] == nil {
				h.sseRooms[roomID] = make(map[*SSEClient]bool)
			}
			h.sseRooms[roomID][client] = true
			if client.subscribedRooms == nil {
				client.subscribedRooms = make(map[string]bool)
			}
			client.subscribedRooms[roomID] = true
			return true
		}
	}
	return false
}

func (h *Hub) LeaveSSERoomByID(sseID, roomID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	for client := range h.sseClients {
		if client.id == sseID {
			if room, ok := h.sseRooms[roomID]; ok {
				delete(room, client)
				if len(room) == 0 {
					delete(h.sseRooms, roomID)
				}
			}
			if client.subscribedRooms != nil {
				delete(client.subscribedRooms, roomID)
			}
			return true
		}
	}
	return false
}

// 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, isAdmin bool) bool {
	h.mu.Lock()
	defer h.mu.Unlock()

	item, ok := h.locks[docID]
	if !ok {
		return false
	}
	if item.userID == userID || isAdmin || 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),
	}
}

// Block Locking Methods (Option B)

func (h *Hub) AcquireBlockLock(docID, blockID, userID, userName, userColor string, ttl time.Duration) (*domain.BlockLockInfo, bool, *domain.BlockLockInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()

	now := time.Now()
	key := docID + ":" + blockID
	existing, exists := h.blockLocks[key]
	if exists && existing.expiresAt.After(now) {
		if existing.userID != userID {
			info := &domain.BlockLockInfo{
				DocID:     existing.docID,
				BlockID:   existing.blockID,
				UserID:    existing.userID,
				UserName:  existing.userName,
				UserColor: existing.userColor,
				ExpiresAt: existing.expiresAt.Format(time.RFC3339),
			}
			return nil, false, info
		}
		// Refresh lease
		existing.expiresAt = now.Add(ttl)
		if userName != "" {
			existing.userName = userName
		}
		if userColor != "" {
			existing.userColor = userColor
		}
		info := &domain.BlockLockInfo{
			DocID:     existing.docID,
			BlockID:   existing.blockID,
			UserID:    existing.userID,
			UserName:  existing.userName,
			UserColor: existing.userColor,
			ExpiresAt: existing.expiresAt.Format(time.RFC3339),
		}
		return info, true, nil
	}

	item := &blockLockItem{
		docID:     docID,
		blockID:   blockID,
		userID:    userID,
		userName:  userName,
		userColor: userColor,
		lockedAt:  now,
		expiresAt: now.Add(ttl),
	}
	h.blockLocks[key] = item
	info := &domain.BlockLockInfo{
		DocID:     item.docID,
		BlockID:   item.blockID,
		UserID:    item.userID,
		UserName:  item.userName,
		UserColor: item.userColor,
		ExpiresAt: item.expiresAt.Format(time.RFC3339),
	}
	return info, true, nil
}

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

	key := docID + ":" + blockID
	existing, exists := h.blockLocks[key]
	if !exists {
		return false
	}
	if userID != "" && existing.userID != userID && userID != "admin" {
		return false
	}
	delete(h.blockLocks, key)
	return true
}

func (h *Hub) GetDocBlockLocks(docID string) []domain.BlockLockInfo {
	h.mu.RLock()
	defer h.mu.RUnlock()

	now := time.Now()
	var result []domain.BlockLockInfo
	prefix := docID + ":"
	for k, l := range h.blockLocks {
		if strings.HasPrefix(k, prefix) && l.expiresAt.After(now) {
			result = append(result, domain.BlockLockInfo{
				DocID:     l.docID,
				BlockID:   l.blockID,
				UserID:    l.userID,
				UserName:  l.userName,
				UserColor: l.userColor,
				ExpiresAt: l.expiresAt.Format(time.RFC3339),
			})
		}
	}
	return result
}

// Live Session Methods (Opt-In Live Collaboration)

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

	now := time.Now()
	existing, exists := h.liveSessions[docID]
	if exists && existing.expiresAt.After(now) {
		if existing.initiatorID != userID {
			info := &domain.LiveSessionInfo{
				DocID:         existing.docID,
				InitiatorID:   existing.initiatorID,
				InitiatorName: existing.initiatorName,
				StartedAt:     existing.startedAt.Format(time.RFC3339),
				ExpiresAt:     existing.expiresAt.Format(time.RFC3339),
			}
			return nil, false, info
		}
		// Refresh lease
		existing.expiresAt = now.Add(ttl)
		info := &domain.LiveSessionInfo{
			DocID:         existing.docID,
			InitiatorID:   existing.initiatorID,
			InitiatorName: existing.initiatorName,
			StartedAt:     existing.startedAt.Format(time.RFC3339),
			ExpiresAt:     existing.expiresAt.Format(time.RFC3339),
		}
		return info, true, nil
	}

	item := &liveSessionItem{
		docID:         docID,
		initiatorID:   userID,
		initiatorName: userName,
		startedAt:     now,
		expiresAt:     now.Add(ttl),
	}
	h.liveSessions[docID] = item
	info := &domain.LiveSessionInfo{
		DocID:         item.docID,
		InitiatorID:   item.initiatorID,
		InitiatorName: item.initiatorName,
		StartedAt:     item.startedAt.Format(time.RFC3339),
		ExpiresAt:     item.expiresAt.Format(time.RFC3339),
	}
	return info, true, nil
}

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

	existing, exists := h.liveSessions[docID]
	if !exists {
		return false
	}
	if userID != "" && existing.initiatorID != userID && userID != "admin" {
		return false
	}
	delete(h.liveSessions, docID)
	return true
}

func (h *Hub) GetLiveSession(docID string) *domain.LiveSessionInfo {
	h.mu.RLock()
	defer h.mu.RUnlock()

	now := time.Now()
	existing, exists := h.liveSessions[docID]
	if !exists || existing.expiresAt.Before(now) {
		return nil
	}
	return &domain.LiveSessionInfo{
		DocID:         existing.docID,
		InitiatorID:   existing.initiatorID,
		InitiatorName: existing.initiatorName,
		StartedAt:     existing.startedAt.Format(time.RFC3339),
		ExpiresAt:     existing.expiresAt.Format(time.RFC3339),
	}
}

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

	for {
		select {
		case <-h.stopCleanup:
			return
		case <-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)
				}
			}

			type expBlock struct {
				docID   string
				blockID string
			}
			var expiredBlocks []expBlock
			for k, bl := range h.blockLocks {
				if bl.expiresAt.Before(now) {
					expiredBlocks = append(expiredBlocks, expBlock{docID: bl.docID, blockID: bl.blockID})
					delete(h.blockLocks, k)
				}
			}

			var expiredLiveSessions []string
			for docID, ls := range h.liveSessions {
				if ls.expiresAt.Before(now) {
					expiredLiveSessions = append(expiredLiveSessions, docID)
					delete(h.liveSessions, docID)
				}
			}
			h.mu.Unlock()

			for _, docID := range expiredDocIDs {
				h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": docID})
			}
			for _, b := range expiredBlocks {
				h.BroadcastRoom("doc:"+b.docID, "DOC_BLOCK_UNLOCKED", map[string]string{
					"docId":   b.docID,
					"blockId": b.blockID,
				}, nil)
			}
			for _, docID := range expiredLiveSessions {
				h.BroadcastRoom("doc:"+docID, "DOC_LIVE_ENDED", map[string]string{
					"docId": docID,
				}, nil)
			}
		}
	}
}

// Close gracefully shuts down the Hub, stops background tickers, flushes dirty document rooms to DB,
// and immediately disconnects all WebSocket and SSE connections so HTTP shutdown completes in milliseconds.
func (h *Hub) Close() {
	h.mu.Lock()
	if h.closed {
		h.mu.Unlock()
		return
	}
	h.closed = true

	if h.stopCleanup != nil {
		select {
		case <-h.stopCleanup:
		default:
			close(h.stopCleanup)
		}
	}

	docRooms := make([]*DocRoom, 0, len(h.docRooms))
	for _, r := range h.docRooms {
		docRooms = append(docRooms, r)
	}

	allWsClients := make([]*Client, 0, len(h.allClients))
	for c := range h.allClients {
		allWsClients = append(allWsClients, c)
	}

	allSSEClients := make([]*SSEClient, 0, len(h.sseClients))
	for c := range h.sseClients {
		allSSEClients = append(allSSEClients, c)
	}
	h.mu.Unlock()

	// 1. Flush any active dirty document rooms and close doc WS connections
	for _, r := range docRooms {
		r.Flush()
		r.mu.Lock()
		if r.debounceTimer != nil {
			r.debounceTimer.Stop()
			r.debounceTimer = nil
		}
		for c := range r.wsClients {
			_ = c.conn.Close()
		}
		r.wsClients = make(map[*DocWSClient]bool)
		r.mu.Unlock()
	}

	// 2. Close all WebSocket connections
	for _, c := range allWsClients {
		_ = c.conn.Close()
		c.closeSend()
	}

	// 3. Close all SSE client channels so ServeSSE select loop exits immediately
	for _, c := range allSSEClients {
		c.closeSend()
	}
}

// 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()
	if h.closed {
		h.mu.RUnlock()
		return
	}
	defer h.mu.RUnlock()

	for client := range h.allClients {
		client.trySend(data)
	}
	for sseClient := range h.sseClients {
		sseClient.trySend(data)
	}
}

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()
	if h.closed {
		h.mu.RUnlock()
		return
	}
	defer h.mu.RUnlock()

	if room, ok := h.userRooms[userID]; ok {
		for client := range room {
			client.trySend(data)
		}
	}
	if sseRoom, ok := h.sseUserRooms[userID]; ok {
		for sseClient := range sseRoom {
			sseClient.trySend(data)
		}
	}
}

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

func (h *Hub) BroadcastDocEvent(docID, event string, payload interface{}) {
	if docID != "" {
		h.BroadcastRoom("doc:"+docID, event, payload, nil)
	} else {
		h.BroadcastToAll(event, payload)
	}
}

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

	h.mu.RLock()
	if h.closed {
		h.mu.RUnlock()
		return
	}
	defer h.mu.RUnlock()

	targetRoom := listID
	if targetRoom == "" {
		targetRoom = spaceID
	}

	if room, ok := h.rooms[targetRoom]; ok {
		for client := range room {
			client.trySend(data)
		}
	}
	if sseRoom, ok := h.sseRooms[targetRoom]; ok {
		for sseClient := range sseRoom {
			sseClient.trySend(data)
		}
	}
}

func (h *Hub) BroadcastListEvent(listID, event string, payload interface{}) {
	h.Broadcast("", listID, event, payload)
}

func (h *Hub) JoinRoom(c *Client, roomID string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.rooms[roomID] == nil {
		h.rooms[roomID] = make(map[*Client]bool)
	}
	h.rooms[roomID][c] = true
	if c.subscribedRooms == nil {
		c.subscribedRooms = make(map[string]bool)
	}
	c.subscribedRooms[roomID] = true
}

func (h *Hub) LeaveRoom(c *Client, roomID string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if room, ok := h.rooms[roomID]; ok {
		delete(room, c)
		if len(room) == 0 {
			delete(h.rooms, roomID)
		}
	}
	if c.subscribedRooms != nil {
		delete(c.subscribedRooms, roomID)
	}
}

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

	h.mu.RLock()
	if h.closed {
		h.mu.RUnlock()
		return
	}
	defer h.mu.RUnlock()

	if room, ok := h.rooms[roomID]; ok {
		for client := range room {
			if excludeClient != nil && client == excludeClient {
				continue
			}
			client.trySend(data)
		}
	}
	if sseRoom, ok := h.sseRooms[roomID]; ok {
		for sseClient := range sseRoom {
			sseClient.trySend(data)
		}
	}
}

func (c *Client) SendEvent(event string, payload interface{}) {
	msg := domain.WSMessage{
		Event:     event,
		Payload:   payload,
		Timestamp: domain.NowISO(),
	}
	data, err := json.Marshal(msg)
	if err == nil {
		c.trySend(data)
	}
}

func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request, roomID, userID string) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		slog.Error("WebSocket upgrade error", "room_id", roomID, "user_id", userID, "error", err)
		return
	}

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

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

func (h *Hub) ServeSSE(w http.ResponseWriter, r *http.Request, roomID, userID string, initialRooms []string) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache, no-transform")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	client := &SSEClient{
		id:              domain.NewID("sse"),
		userID:          userID,
		roomID:          roomID,
		send:            make(chan []byte, 128),
		subscribedRooms: make(map[string]bool),
	}

	for _, room := range initialRooms {
		if strings.TrimSpace(room) != "" {
			client.subscribedRooms[strings.TrimSpace(room)] = true
		}
	}

	h.RegisterSSE(client)
	defer h.UnregisterSSE(client)

	// Send initial connected comment and event
	fmt.Fprintf(w, ": connected (transport: sse, id: %s)\n\n", client.id)
	connectedMsg, _ := json.Marshal(domain.WSMessage{
		Event: "SSE_CONNECTED",
		Payload: map[string]interface{}{
			"sseId":     client.id,
			"transport": "sse",
			"userId":    userID,
		},
		Timestamp: domain.NowISO(),
	})
	fmt.Fprintf(w, "data: %s\n\n", connectedMsg)
	flusher.Flush()

	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()

	notify := r.Context().Done()

	for {
		select {
		case <-notify:
			return
		case msg, ok := <-client.send:
			if !ok {
				return
			}
			fmt.Fprintf(w, "data: %s\n\n", msg)
			flusher.Flush()
		case <-ticker.C:
			// Periodic ping comment to keep HTTP connection alive and prevent proxy timeouts
			fmt.Fprintf(w, ": ping\n\n")
			flusher.Flush()
		}
	}
}

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

	c.conn.SetReadLimit(4 * 1024 * 1024) // 4MB for large document CRDT updates
	_ = 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, CRDT, room subscription, and presence messages
		var incoming struct {
			Event   string          `json:"event"`
			DocID   string          `json:"docId"`
			Room    string          `json:"room"`
			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
			}

			// Room subscription without requiring docID
			if incoming.Event == "JOIN_ROOM" {
				room := incoming.Room
				if room == "" {
					room = docID
				}
				if room != "" {
					c.hub.JoinRoom(c, room)
				}
			} else if incoming.Event == "LEAVE_ROOM" {
				room := incoming.Room
				if room == "" {
					room = docID
				}
				if room != "" {
					c.hub.LeaveRoom(c, room)
				}
			}

			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, false) {
						c.hub.BroadcastDocEvent(docID, "DOC_UNLOCKED", map[string]string{"docId": docID, "doc_id": 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
			}
		}
	}
}