MindSift/internal/ws/docroom.go

package ws

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

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

// CursorPresence represents an active user's ephemeral cursor position in a document room.
type CursorPresence struct {
	ClientID       string `json:"clientId"`
	UserID         string `json:"userId"`
	UserName       string `json:"userName"`
	UserColor      string `json:"userColor"`
	BlockID        string `json:"blockId,omitempty"`
	SelectionStart *int   `json:"selectionStart,omitempty"`
	SelectionEnd   *int   `json:"selectionEnd,omitempty"`
	UpdatedAt      int64  `json:"updatedAt"`
}

// DocWSClient represents an active WebSocket connection to a specific document room.
type DocWSClient struct {
	hub      *Hub
	room     *DocRoom
	conn     *websocket.Conn
	userID   string
	clientID string
	canWrite bool
	send     chan []byte
}

// DocSyncEvent is the lightweight JSON payload exchanged over /api/v1/ws/docs/{id}.
type DocSyncEvent struct {
	Type           string                 `json:"type"`
	DocID          string                 `json:"docId,omitempty"`
	ClientID       string                 `json:"clientId,omitempty"`
	UserID         string                 `json:"userId,omitempty"`
	UserName       string                 `json:"userName,omitempty"`
	UserColor      string                 `json:"userColor,omitempty"`
	SelectionStart *int                   `json:"selectionStart,omitempty"`
	SelectionEnd   *int                   `json:"selectionEnd,omitempty"`
	ActiveCursors  []CursorPresence       `json:"activeCursors,omitempty"`
	Title          string                 `json:"title,omitempty"`
	Block          *domain.DocumentBlock  `json:"block,omitempty"`
	BlockID        string                 `json:"blockId,omitempty"`
	AfterBlockID   string                 `json:"afterBlockId,omitempty"`
	NewIndex       *int                   `json:"newIndex,omitempty"`
	Content        *string                `json:"content,omitempty"`
	Level          *int                   `json:"level,omitempty"`
	Checked        *bool                  `json:"checked,omitempty"`
	Collapsed      *bool                  `json:"collapsed,omitempty"`
	BlockType      *string                `json:"blockType,omitempty"`
	Blocks         []domain.DocumentBlock `json:"blocks,omitempty"`
}

// DocRoom coordinates in-memory real-time block state and 1.5s debounced SQLite persistence for one document.
type DocRoom struct {
	mu            sync.Mutex
	docID         string
	hub           *Hub
	title         string
	blocks        []domain.DocumentBlock
	loaded        bool
	wsClients     map[*DocWSClient]bool
	activeCursors map[string]*CursorPresence
	debounceTimer *time.Timer
	dirty         bool
}

// GetOrCreateDocRoom returns the active in-memory room for docID or initializes one from SQLite.
func (h *Hub) GetOrCreateDocRoom(docID string) *DocRoom {
	h.mu.Lock()
	if h.docRooms == nil {
		h.docRooms = make(map[string]*DocRoom)
	}
	room, exists := h.docRooms[docID]
	if !exists {
		room = &DocRoom{
			docID:         docID,
			hub:           h,
			wsClients:     make(map[*DocWSClient]bool),
			activeCursors: make(map[string]*CursorPresence),
		}
		h.docRooms[docID] = room
	}
	store := h.docStore
	h.mu.Unlock()

	room.mu.Lock()
	if !room.loaded && store != nil {
		title, blocks, err := store.LoadDocumentState(docID)
		if err == nil {
			room.title = title
			room.blocks = normalizeRoomBlocks(docID, blocks)
			room.loaded = true
		}
	}
	room.mu.Unlock()

	return room
}

// GetActiveDocState returns the live in-memory title and blocks if an active DocRoom exists in RAM.
func (h *Hub) GetActiveDocState(docID string) (string, []domain.DocumentBlock, bool) {
	h.mu.RLock()
	room, exists := h.docRooms[docID]
	h.mu.RUnlock()
	if !exists || room == nil {
		return "", nil, false
	}

	room.mu.Lock()
	defer room.mu.Unlock()
	if !room.loaded {
		return "", nil, false
	}
	blocksCopy := make([]domain.DocumentBlock, len(room.blocks))
	copy(blocksCopy, room.blocks)
	return room.title, blocksCopy, true
}

// SyncActiveDocRoom updates an in-memory room if one is currently loaded (e.g. after REST PUT/PATCH).
func (h *Hub) SyncActiveDocRoom(docID string, title string, blocks []domain.DocumentBlock) {
	h.mu.RLock()
	room, exists := h.docRooms[docID]
	h.mu.RUnlock()
	if !exists || room == nil {
		return
	}

	room.mu.Lock()
	if title != "" {
		room.title = title
	}
	if blocks != nil {
		room.blocks = normalizeRoomBlocks(docID, blocks)
	}
	room.loaded = true
	room.mu.Unlock()
}

// CloseDocRoom removes a document room from memory when a document is deleted.
func (h *Hub) CloseDocRoom(docID string) {
	h.mu.Lock()
	room, exists := h.docRooms[docID]
	if exists {
		delete(h.docRooms, docID)
	}
	h.mu.Unlock()

	if exists && room != nil {
		room.mu.Lock()
		if room.debounceTimer != nil {
			room.debounceTimer.Stop()
		}
		room.dirty = false
		for c := range room.wsClients {
			_ = c.conn.Close()
		}
		room.wsClients = make(map[*DocWSClient]bool)
		room.mu.Unlock()
	}
}

func normalizeRoomBlocks(docID string, blocks []domain.DocumentBlock) []domain.DocumentBlock {
	if len(blocks) == 0 {
		return []domain.DocumentBlock{
			{
				ID:         docID + "-b0",
				DocumentID: docID,
				Type:       "text",
				Content:    "",
				Level:      0,
				Position:   1000,
			},
		}
	}
	out := make([]domain.DocumentBlock, len(blocks))
	for i, b := range blocks {
		if b.ID == "" {
			b.ID = domain.NewID("blk")
		}
		b.DocumentID = docID
		if b.Type == "" {
			b.Type = "text"
		}
		if b.Level < 0 {
			b.Level = 0
		}
		b.Position = float64((i + 1) * 1000)
		out[i] = b
	}
	return out
}

// GetStateSnapshot returns a copy of the room's current live in-memory title and blocks.
func (r *DocRoom) GetStateSnapshot() (string, []domain.DocumentBlock) {
	r.mu.Lock()
	defer r.mu.Unlock()
	blocksCopy := make([]domain.DocumentBlock, len(r.blocks))
	copy(blocksCopy, r.blocks)
	return r.title, blocksCopy
}

// HandleEvent applies a granular block or title mutation to the in-memory room,
// broadcasts it to all OTHER clients in the room (excluding senderClientID for echo protection),
// and resets the 1.5-second inactivity debounce timer for SQLite persistence.
func (r *DocRoom) HandleEvent(evt DocSyncEvent, sender *DocWSClient) {
	r.mu.Lock()

	evt.DocID = r.docID
	if sender != nil {
		if evt.ClientID == "" {
			evt.ClientID = sender.clientID
		}
		evt.UserID = sender.userID
	}

	mutated := false

	switch evt.Type {
	case "cursor_moved":
		if evt.ClientID != "" {
			cp := &CursorPresence{
				ClientID:       evt.ClientID,
				UserID:         evt.UserID,
				UserName:       evt.UserName,
				UserColor:      evt.UserColor,
				BlockID:        evt.BlockID,
				SelectionStart: evt.SelectionStart,
				SelectionEnd:   evt.SelectionEnd,
				UpdatedAt:      time.Now().UnixMilli(),
			}
			if r.activeCursors == nil {
				r.activeCursors = make(map[string]*CursorPresence)
			}
			r.activeCursors[evt.ClientID] = cp

			msgBytes, err := json.Marshal(evt)
			if err == nil {
				for client := range r.wsClients {
					if sender != nil && (client == sender || (sender.clientID != "" && client.clientID == sender.clientID)) {
						continue
					}
					if evt.ClientID != "" && client.clientID == evt.ClientID {
						continue
					}
					select {
					case client.send <- msgBytes:
					default:
					}
				}
			}
		}
		r.mu.Unlock()
		return

	case "title_updated":
		if evt.Title != "" || r.title != evt.Title {
			r.title = evt.Title
			mutated = true
		}

	case "block_created":
		if evt.Block != nil {
			newBlock := *evt.Block
			if newBlock.ID == "" {
				newBlock.ID = domain.NewID("blk")
			}
			newBlock.DocumentID = r.docID
			if newBlock.Type == "" {
				newBlock.Type = "text"
			}
			if newBlock.Level < 0 {
				newBlock.Level = 0
			}

			// Avoid duplicate ID if already created
			existsIdx := -1
			for i, b := range r.blocks {
				if b.ID == newBlock.ID {
					existsIdx = i
					break
				}
			}

			if existsIdx >= 0 {
				r.blocks[existsIdx] = newBlock
			} else {
				insertIdx := len(r.blocks)
				if evt.AfterBlockID != "" {
					for i, b := range r.blocks {
						if b.ID == evt.AfterBlockID {
							insertIdx = i + 1
							break
						}
					}
				} else if evt.NewIndex != nil && *evt.NewIndex >= 0 && *evt.NewIndex <= len(r.blocks) {
					insertIdx = *evt.NewIndex
				}

				updated := make([]domain.DocumentBlock, 0, len(r.blocks)+1)
				updated = append(updated, r.blocks[:insertIdx]...)
				updated = append(updated, newBlock)
				updated = append(updated, r.blocks[insertIdx:]...)
				r.blocks = normalizeRoomBlocks(r.docID, updated)
			}
			evt.Block = &newBlock
			mutated = true
		}

	case "block_updated":
		targetID := evt.BlockID
		if targetID == "" && evt.Block != nil {
			targetID = evt.Block.ID
		}
		if targetID != "" {
			found := false
			for i := range r.blocks {
				if r.blocks[i].ID == targetID {
					found = true
					if evt.Block != nil {
						r.blocks[i].Content = evt.Block.Content
						r.blocks[i].Level = evt.Block.Level
						if evt.Block.Type != "" {
							r.blocks[i].Type = evt.Block.Type
						}
						r.blocks[i].Checked = evt.Block.Checked
						r.blocks[i].Collapsed = evt.Block.Collapsed
						r.blocks[i].Lang = evt.Block.Lang
					} else {
						if evt.Content != nil {
							r.blocks[i].Content = *evt.Content
						}
						if evt.Level != nil {
							lvl := *evt.Level
							if lvl < 0 {
								lvl = 0
							}
							r.blocks[i].Level = lvl
						}
						if evt.Checked != nil {
							r.blocks[i].Checked = *evt.Checked
						}
						if evt.Collapsed != nil {
							r.blocks[i].Collapsed = *evt.Collapsed
						}
						if evt.BlockType != nil && *evt.BlockType != "" {
							r.blocks[i].Type = *evt.BlockType
						}
					}
					bCopy := r.blocks[i]
					evt.Block = &bCopy
					evt.BlockID = targetID
					mutated = true
					break
				}
			}
			if !found && evt.Block != nil {
				newBlock := *evt.Block
				newBlock.DocumentID = r.docID
				if newBlock.Type == "" {
					newBlock.Type = "text"
				}
				if newBlock.Level < 0 {
					newBlock.Level = 0
				}
				r.blocks = append(r.blocks, newBlock)
				r.blocks = normalizeRoomBlocks(r.docID, r.blocks)
				bCopy := newBlock
				evt.Block = &bCopy
				evt.BlockID = targetID
				mutated = true
			}
		}

	case "block_deleted":
		if evt.BlockID != "" {
			filtered := make([]domain.DocumentBlock, 0, len(r.blocks))
			for _, b := range r.blocks {
				if b.ID != evt.BlockID {
					filtered = append(filtered, b)
				}
			}
			r.blocks = normalizeRoomBlocks(r.docID, filtered)
			mutated = true
		}

	case "block_moved":
		if evt.BlockID != "" && evt.NewIndex != nil {
			fromIdx := -1
			for i, b := range r.blocks {
				if b.ID == evt.BlockID {
					fromIdx = i
					break
				}
			}
			if fromIdx >= 0 {
				moving := r.blocks[fromIdx]
				if evt.Level != nil {
					lvl := *evt.Level
					if lvl < 0 {
						lvl = 0
					}
					moving.Level = lvl
				}
				without := make([]domain.DocumentBlock, 0, len(r.blocks)-1)
				without = append(without, r.blocks[:fromIdx]...)
				without = append(without, r.blocks[fromIdx+1:]...)

				toIdx := *evt.NewIndex
				if toIdx < 0 {
					toIdx = 0
				}
				if toIdx > len(without) {
					toIdx = len(without)
				}

				reordered := make([]domain.DocumentBlock, 0, len(r.blocks))
				reordered = append(reordered, without[:toIdx]...)
				reordered = append(reordered, moving)
				reordered = append(reordered, without[toIdx:]...)
				r.blocks = normalizeRoomBlocks(r.docID, reordered)
				mutated = true
			}
		}

	case "blocks_replaced":
		if evt.Blocks != nil {
			r.blocks = normalizeRoomBlocks(r.docID, evt.Blocks)
			evt.Blocks = r.blocks
			mutated = true
		}
	}

	if !mutated {
		r.mu.Unlock()
		return
	}

	r.dirty = true
	r.loaded = true

	msgBytes, err := json.Marshal(evt)
	if err == nil {
		for client := range r.wsClients {
			// Strict Echo Protection: never send back to the originating connection or clientId
			if sender != nil && (client == sender || (sender.clientID != "" && client.clientID == sender.clientID)) {
				continue
			}
			if evt.ClientID != "" && client.clientID == evt.ClientID {
				continue
			}
			select {
			case client.send <- msgBytes:
			default:
				slog.Warn("Doc room WS client send buffer full, dropping message",
					"doc_id", r.docID,
					"client_id", client.clientID,
					"user_id", client.userID,
				)
			}
		}
	}

	r.scheduleDebouncedFlushLocked()
	r.mu.Unlock()
}

func (r *DocRoom) scheduleDebouncedFlushLocked() {
	if r.debounceTimer != nil {
		r.debounceTimer.Stop()
	}
	r.debounceTimer = time.AfterFunc(1500*time.Millisecond, func() {
		r.Flush()
	})
}

// Flush persists the current in-memory document state (blocks, #tags in block_tags, and markdown_cache) to SQLite.
func (r *DocRoom) Flush() {
	r.mu.Lock()
	if !r.dirty {
		r.mu.Unlock()
		return
	}
	if r.debounceTimer != nil {
		r.debounceTimer.Stop()
		r.debounceTimer = nil
	}

	docID := r.docID
	title := r.title
	blocksCopy := make([]domain.DocumentBlock, len(r.blocks))
	copy(blocksCopy, r.blocks)
	r.dirty = false
	r.mu.Unlock()

	tagsCount := 0
	for _, b := range blocksCopy {
		tagsCount += len(domain.ExtractTags(b.Content))
	}

	start := time.Now()
	r.hub.mu.RLock()
	store := r.hub.docStore
	r.hub.mu.RUnlock()

	if store != nil {
		if err := store.FlushDocumentState(docID, title, blocksCopy); err != nil {
			slog.Error("SQLite document flush failed",
				"doc_id", docID,
				"error", err,
			)
		} else {
			durationMs := float64(time.Since(start).Microseconds()) / 1000.0
			slog.Debug("SQLite document flushed",
				"doc_id", docID,
				"blocks_count", len(blocksCopy),
				"tags_count", tagsCount,
				"duration_ms", durationMs,
			)
		}
	}
}

// ServeDocWS handles WebSocket connections for GET /api/v1/ws/docs/{id}
func (r *DocRoom) ServeDocWS(w http.ResponseWriter, req *http.Request, userID string, clientID string, canWrite bool) {
	conn, err := upgrader.Upgrade(w, req, nil)
	if err != nil {
		slog.Error("WebSocket upgrade error", "doc_id", r.docID, "error", err)
		return
	}

	if clientID == "" {
		clientID = req.URL.Query().Get("clientId")
	}
	if clientID == "" {
		clientID = domain.NewID("cli")
	}

	userName := req.URL.Query().Get("userName")
	userColor := req.URL.Query().Get("userColor")

	client := &DocWSClient{
		hub:      r.hub,
		room:     r,
		conn:     conn,
		userID:   userID,
		clientID: clientID,
		canWrite: canWrite,
		send:     make(chan []byte, 256),
	}

	r.mu.Lock()
	r.wsClients[client] = true
	activeClients := len(r.wsClients)
	if r.activeCursors == nil {
		r.activeCursors = make(map[string]*CursorPresence)
	}
	if userName != "" {
		r.activeCursors[clientID] = &CursorPresence{
			ClientID:  clientID,
			UserID:    userID,
			UserName:  userName,
			UserColor: userColor,
			UpdatedAt: time.Now().UnixMilli(),
		}
	}
	title := r.title
	blocksCopy := make([]domain.DocumentBlock, len(r.blocks))
	copy(blocksCopy, r.blocks)

	activeList := make([]CursorPresence, 0, len(r.activeCursors))
	for _, cp := range r.activeCursors {
		if cp != nil {
			activeList = append(activeList, *cp)
		}
	}
	r.mu.Unlock()

	slog.Debug("WS client joined doc room",
		"doc_id", r.docID,
		"client_id", clientID,
		"user_id", userID,
		"active_clients", activeClients,
	)

	// Immediately send live in-memory state and current cursors to the newly connected client
	initEvt := DocSyncEvent{
		Type:          "doc_state",
		DocID:         r.docID,
		ClientID:      clientID,
		Title:         title,
		Blocks:        blocksCopy,
		ActiveCursors: activeList,
	}
	if initBytes, err := json.Marshal(initEvt); err == nil {
		client.send <- initBytes
	}

	// If userName was provided, broadcast initial presence to existing clients in room
	if userName != "" {
		joinEvt := DocSyncEvent{
			Type:      "cursor_moved",
			DocID:     r.docID,
			ClientID:  clientID,
			UserID:    userID,
			UserName:  userName,
			UserColor: userColor,
		}
		if joinBytes, err := json.Marshal(joinEvt); err == nil {
			r.mu.Lock()
			for c := range r.wsClients {
				if c != client && c.clientID != clientID {
					select {
					case c.send <- joinBytes:
					default:
					}
				}
			}
			r.mu.Unlock()
		}
	}

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

func (c *DocWSClient) readPump() {
	defer func() {
		c.room.mu.Lock()
		delete(c.room.wsClients, c)
		delete(c.room.activeCursors, c.clientID)
		remaining := len(c.room.wsClients)

		// Broadcast client_left to all remaining clients in room
		leftEvt := DocSyncEvent{
			Type:     "client_left",
			DocID:    c.room.docID,
			ClientID: c.clientID,
			UserID:   c.userID,
		}
		if leftBytes, err := json.Marshal(leftEvt); err == nil {
			for client := range c.room.wsClients {
				if client != c && client.clientID != c.clientID {
					select {
					case client.send <- leftBytes:
					default:
					}
				}
			}
		}
		c.room.mu.Unlock()

		slog.Debug("WS client left doc room",
			"doc_id", c.room.docID,
			"client_id", c.clientID,
			"user_id", c.userID,
			"active_clients", remaining,
		)

		_ = c.conn.Close()

		// Flush immediately when the last client leaves the document room
		if remaining == 0 {
			c.room.Flush()
		}
	}()

	c.conn.SetReadLimit(4 * 1024 * 1024)
	_ = 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
		}

		var evt DocSyncEvent
		if err := json.Unmarshal(message, &evt); err == nil && evt.Type != "" {
			targetBlockID := evt.BlockID
			if targetBlockID == "" && evt.Block != nil {
				targetBlockID = evt.Block.ID
			}

			slog.Debug("WS doc event received",
				"type", evt.Type,
				"doc_id", c.room.docID,
				"block_id", targetBlockID,
				"sender_client_id", c.clientID,
			)

			// Read-only clients can still move their cursor and send presence updates
			if !c.canWrite && evt.Type != "cursor_moved" {
				continue
			}
			c.room.HandleEvent(evt, c)
		}
	}
}

func (c *DocWSClient) writePump() {
	ticker := time.NewTicker(25 * 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
			}
		}
	}
}