MindSift/internal/ws/docroom.go

package ws

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

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

// 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"`
	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
	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),
		}
		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 "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 != "" {
			for i := range r.blocks {
				if r.blocks[i].ID == targetID {
					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
				}
			}
		}

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

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

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

	if store != nil {
		if err := store.FlushDocumentState(docID, title, blocksCopy); err != nil {
			log.Printf("[doc-room] flush error for doc %s: %v", docID, err)
		}
	}
}

// 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 {
		log.Printf("[doc-ws] upgrade error for doc %s: %v", r.docID, err)
		return
	}

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

	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
	title := r.title
	blocksCopy := make([]domain.DocumentBlock, len(r.blocks))
	copy(blocksCopy, r.blocks)
	r.mu.Unlock()

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

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

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

		_ = 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 != "" {
			if !c.canWrite {
				// Read-only clients cannot mutate the document
				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
			}
		}
	}
}