Can't restore selection after deleting/modifying lexical node

Question

I wrote a custom plugin that highlights incorrect words in the text by adding a specific class (case 1) or removes the class and merges the node with adjacent ones when the label is no longer needed (case 2).

const SpellcheckPlugin = (props) => {
  const [editor] = useLexicalComposerContext()

  createEffect(() => {
    const rootElement = editor.getRootElement()
    if (!rootElement) return

    const spellcheckHandler = (event) => {
      const { misspelled } = event.detail

      editor.update(() => {
        const selection = $getSelection()

        const misspelledWordsRegexp = new RegExp(
          `${misspelled.map((key) => `(${key})`).join('|')}`,
          'gi'
        )

        rootElement.firstElementChild?.childNodes.forEach((node) => {
          const lexicalNode = $getNearestNodeFromDOMNode(node)
          // Case 1: Add misspelled
          if (
            misspelled.some((word) => node.textContent.includes(word)) &&
            !node.classList.contains('misspelled')
          ) {
            const textTokens = node.textContent.split(misspelledWordsRegexp).filter(Boolean)
            const newNodes = textTokens.map((text) => {
              const spanNode = $createExtendedTextNode(text)
              if (misspelled.includes(text)) {
                spanNode.addClass('misspelled')
              }
              return spanNode
            })
            newNodes.reverse().forEach((node) => {
              lexicalNode.insertAfter(node)
            })
            lexicalNode.remove()
          }

          // Case 2: Remove misspelled и merge
          if (!misspelled.includes(node.textContent) && node.classList.contains('misspelled')) {
            node.classList.remove('misspelled')
            if (node.getAttribute('class') === '') node.removeAttribute('class')
            const previousNode = lexicalNode.getPreviousSibling()
            const nextNode = lexicalNode.getNextSibling()
            let textContent = lexicalNode.getTextContent()
            lexicalNode.select()
            if ($isExtendedTextNode(previousNode)) {
              textContent = previousNode.getTextContent() + textContent
              offset += textContent.length
              previousNode.remove()
            }
            if ($isExtendedTextNode(nextNode)) {
              textContent += nextNode.getTextContent()
              nextNode.remove()
            }
            lexicalNode.setTextContent(textContent)
            // Trying to restore selection
            const updatedSelection = $createRangeSelection()
            updatedSelection.anchor.set(lexicalNode, offset)
            updatedSelection.focus.set(lexicalNode, offset)
            $setSelection(updatedSelection)
          }
        })
      })
    }

    rootElement.addEventListener('spellcheck', spellcheckHandler)

    onCleanup(() => rootElement.removeEventListener('spellcheck', spellcheckHandler))
  })

  return null
}

export default SpellcheckPlugin

When trying to restore the cursor (selection) to its original position or at least close to it, I encounter the following error: Lexical error: Error: updateEditor: selection has been lost because the previously selected nodes have been removed and selection wasn't moved to another node. Ensure selection changes after removing/replacing a selected node.

I also tried a different approach by calculating the absolute offset of the cursor and then restoring it after the changes, but I ended up with the same error.

const selection = $getSelection()
let cursorIndex = null;
if (selection && selection.isCollapsed()) {
  const anchorNode = selection.anchor.getNode();
          
  const anchorOffset = anchorNode.getPreviousSiblings().reduce((acc, node) => acc + node.getTextContentSize(), 0) + selection.anchor.offset;
  cursorIndex = anchorNode.getTextContentSize() + anchorOffset;
}

//............
if (cursorIndex !== null) {
  const newSelection = $getSelection();
  let currentIndex = 0;
    
  rootElement.firstElementChild?.childNodes.forEach((node) => {
    const lexicalNode = $getNearestNodeFromDOMNode(node)
    const nodeTextLength = lexicalNode.getTextContentSize();
    if (cursorIndex <= currentIndex + nodeTextLength) {
      const offset = cursorIndex - currentIndex;
      newSelection.anchor.set(lexicalNode, offset);
      newSelection.focus.set(lexicalNode, offset);
      return; 
    }
    currentIndex += nodeTextLength;
  });
}

What could be causing this issue, and how can I fix it to restore the selection properly?