rspace-online/modules/rnotes/components/suggestion-plugin.ts

/**
 * ProseMirror plugin that intercepts transactions in "suggesting" mode
 * and converts them into track-changes marks instead of direct edits.
 *
 * In suggesting mode:
 * - Insertions → wraps inserted text with `suggestionInsert` mark
 * - Deletions → converts to `suggestionDelete` mark instead of deleting
 *
 * Accept/reject operations remove marks (and optionally the text).
 */

import { Plugin, PluginKey, type Transaction } from '@tiptap/pm/state';
import type { Editor } from '@tiptap/core';

const pluginKey = new PluginKey('suggestion-plugin');

interface SuggestionPluginState {
	suggesting: boolean;
	authorId: string;
	authorName: string;
}

/**
 * Create the suggestion mode ProseMirror plugin.
 * @param getSuggesting - callback that returns current suggesting mode state
 * @param getAuthor - callback that returns { authorId, authorName }
 */
export function createSuggestionPlugin(
	getSuggesting: () => boolean,
	getAuthor: () => { authorId: string; authorName: string },
): Plugin {
	return new Plugin({
		key: pluginKey,

		filterTransaction(tr: Transaction, state) {
			if (!getSuggesting()) return true;
			if (!tr.docChanged) return true;
			if (tr.getMeta('suggestion-accept') || tr.getMeta('suggestion-reject')) return true;
			if (tr.getMeta('suggestion-applied')) return true;

			// Intercept the transaction and convert it to suggestion marks
			const { authorId, authorName } = getAuthor();
			const suggestionId = `s_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;

			// We need to rebuild the transaction with suggestion marks
			const newTr = state.tr;
			let blocked = false;

			tr.steps.forEach((step, i) => {
				const stepMap = step.getMap();
				let hasInsert = false;
				let hasDelete = false;

				stepMap.forEach((oldStart: number, oldEnd: number, newStart: number, newEnd: number) => {
					if (newEnd > newStart) hasInsert = true;
					if (oldEnd > oldStart) hasDelete = true;
				});

				if (hasInsert && !hasDelete) {
					// Pure insertion — let it through but add the suggestionInsert mark
					blocked = true;
					const doc = tr.docs[i];
					stepMap.forEach((_oldStart: number, _oldEnd: number, newStart: number, newEnd: number) => {
						if (newEnd > newStart) {
							const insertedText = tr.docs[i + 1]?.textBetween(newStart, newEnd, '', '') || '';
							if (insertedText) {
								const insertMark = state.schema.marks.suggestionInsert.create({
									suggestionId,
									authorId,
									authorName,
									createdAt: Date.now(),
								});
								newTr.insertText(insertedText, newStart);
								newTr.addMark(newStart, newStart + insertedText.length, insertMark);
								newTr.setMeta('suggestion-applied', true);
							}
						}
					});
				} else if (hasDelete && !hasInsert) {
					// Pure deletion — convert to suggestionDelete mark instead
					blocked = true;
					stepMap.forEach((oldStart: number, oldEnd: number) => {
						if (oldEnd > oldStart) {
							const deleteMark = state.schema.marks.suggestionDelete.create({
								suggestionId,
								authorId,
								authorName,
								createdAt: Date.now(),
							});
							newTr.addMark(oldStart, oldEnd, deleteMark);
							newTr.setMeta('suggestion-applied', true);
						}
					});
				} else if (hasInsert && hasDelete) {
					// Replacement (delete + insert) — mark old text as deleted, new text as inserted
					blocked = true;
					stepMap.forEach((oldStart: number, oldEnd: number, newStart: number, newEnd: number) => {
						if (oldEnd > oldStart) {
							const deleteMark = state.schema.marks.suggestionDelete.create({
								suggestionId,
								authorId,
								authorName,
								createdAt: Date.now(),
							});
							newTr.addMark(oldStart, oldEnd, deleteMark);
						}
						if (newEnd > newStart) {
							const insertedText = tr.docs[i + 1]?.textBetween(newStart, newEnd, '', '') || '';
							if (insertedText) {
								const insertMark = state.schema.marks.suggestionInsert.create({
									suggestionId,
									authorId,
									authorName,
									createdAt: Date.now(),
								});
								// Insert after the "deleted" text
								const insertPos = oldEnd;
								newTr.insertText(insertedText, insertPos);
								newTr.addMark(insertPos, insertPos + insertedText.length, insertMark);
							}
						}
					});
					newTr.setMeta('suggestion-applied', true);
				}
			});

			if (blocked && newTr.docChanged) {
				// Dispatch our modified transaction instead
				// We need to use view.dispatch in the next tick
				setTimeout(() => {
					const view = (state as any).view;
					if (view) view.dispatch(newTr);
				}, 0);
				return false; // Block the original transaction
			}

			return !blocked;
		},
	});
}

/**
 * Accept a suggestion by removing the mark.
 * - For insertions: remove the mark (text stays)
 * - For deletions: remove the text and the mark
 */
export function acceptSuggestion(editor: Editor, suggestionId: string) {
	const { state } = editor;
	const { doc, tr } = state;

	// Find all marks with this suggestionId
	doc.descendants((node, pos) => {
		if (!node.isText) return;

		// Check for suggestionDelete marks — accept = remove the text
		const deleteMark = node.marks.find(
			m => m.type.name === 'suggestionDelete' && m.attrs.suggestionId === suggestionId
		);
		if (deleteMark) {
			tr.delete(pos, pos + node.nodeSize);
			tr.setMeta('suggestion-accept', true);
			return false;
		}

		// Check for suggestionInsert marks — accept = remove the mark, keep text
		const insertMark = node.marks.find(
			m => m.type.name === 'suggestionInsert' && m.attrs.suggestionId === suggestionId
		);
		if (insertMark) {
			tr.removeMark(pos, pos + node.nodeSize, insertMark);
			tr.setMeta('suggestion-accept', true);
		}
	});

	if (tr.docChanged) {
		editor.view.dispatch(tr);
	}
}

/**
 * Reject a suggestion by reverting it.
 * - For insertions: remove the text and the mark
 * - For deletions: remove the mark (text stays)
 */
export function rejectSuggestion(editor: Editor, suggestionId: string) {
	const { state } = editor;
	const { doc, tr } = state;

	doc.descendants((node, pos) => {
		if (!node.isText) return;

		// Check for suggestionInsert marks — reject = remove text + mark
		const insertMark = node.marks.find(
			m => m.type.name === 'suggestionInsert' && m.attrs.suggestionId === suggestionId
		);
		if (insertMark) {
			tr.delete(pos, pos + node.nodeSize);
			tr.setMeta('suggestion-reject', true);
			return false;
		}

		// Check for suggestionDelete marks — reject = remove the mark, keep text
		const deleteMark = node.marks.find(
			m => m.type.name === 'suggestionDelete' && m.attrs.suggestionId === suggestionId
		);
		if (deleteMark) {
			tr.removeMark(pos, pos + node.nodeSize, deleteMark);
			tr.setMeta('suggestion-reject', true);
		}
	});

	if (tr.docChanged) {
		editor.view.dispatch(tr);
	}
}

/**
 * Accept all suggestions in the document.
 */
export function acceptAllSuggestions(editor: Editor) {
	const ids = new Set<string>();
	editor.state.doc.descendants((node) => {
		if (!node.isText) return;
		for (const mark of node.marks) {
			if (mark.type.name === 'suggestionInsert' || mark.type.name === 'suggestionDelete') {
				ids.add(mark.attrs.suggestionId);
			}
		}
	});
	for (const id of ids) {
		acceptSuggestion(editor, id);
	}
}

/**
 * Reject all suggestions in the document.
 */
export function rejectAllSuggestions(editor: Editor) {
	const ids = new Set<string>();
	editor.state.doc.descendants((node) => {
		if (!node.isText) return;
		for (const mark of node.marks) {
			if (mark.type.name === 'suggestionInsert' || mark.type.name === 'suggestionDelete') {
				ids.add(mark.attrs.suggestionId);
			}
		}
	});
	for (const id of ids) {
		rejectSuggestion(editor, id);
	}
}