public-tiptap-comments

tiptap-comment-extension.ts

import { Mark, mergeAttributes, type Range } from '@tiptap/core';
import { Mark as PMMark } from '@tiptap/pm/model';

declare module '@tiptap/core' {
	interface Commands<ReturnType> {
		comment: {
			addComment: (commentId: string) => ReturnType;
			removeComment: (commentId: string) => ReturnType;
			highlightComment: (commentId: string) => ReturnType;
			unhighlightComment: (commentId: string) => ReturnType;
		};
	}
}

export interface MarkWithRange {
	mark: PMMark;
	range: Range;
}

export interface CommentOptions {
	HTMLAttributes: Record<string, unknown>;
	onCommentActivated: (commentId: string) => void;
	onCommentRemoved: (commentId: string) => void;
}

export interface CommentStorage {
	activeCommentId: string | null;
}

export const CommentExtension = Mark.create<CommentOptions, CommentStorage>({
	name: 'comment',

	addOptions() {
		return {
			HTMLAttributes: {},
			onCommentActivated: () => {},
			onCommentRemoved: () => {}
		};
	},

	addAttributes() {
		return {
			commentId: {
				default: null,
				parseHTML: (el) => (el as HTMLSpanElement).getAttribute('data-comment-id'),
				renderHTML: (attrs) => ({ 'data-comment-id': attrs.commentId })
			},
			highlighted: {
				default: false,
				parseHTML: (el) => (el as HTMLSpanElement).classList.contains('highlighted-comment'),
				renderHTML: (attrs) => (attrs.highlighted ? { class: 'highlighted-comment' } : {})
			}
		};
	},

	parseHTML() {
		return [
			{
				tag: 'span[data-comment-id]',
				getAttrs: (el) => {
					const commentId = (el as HTMLSpanElement).getAttribute('data-comment-id')?.trim();
					return commentId ? { commentId } : false;
				}
			}
		];
	},

	renderHTML({ HTMLAttributes }) {
		const attrs = mergeAttributes(this.options.HTMLAttributes, HTMLAttributes);
		const classes = new Set(attrs.class ? (attrs.class as string).split(' ') : ['comment']);

		if (attrs.highlighted) {
			classes.add('highlighted-comment');
		}

		return ['span', { ...attrs, class: [...classes].join(' ') }, 0];
	},

	onSelectionUpdate() {
		const { $from } = this.editor.state.selection;

		const marks = $from.marks();

		if (!marks.length) {
			this.storage.activeCommentId = null;
			this.options.onCommentActivated(this.storage.activeCommentId!);
			return;
		}

		const commentMark = this.editor.schema.marks.comment;

		const activeCommentMark = marks.find((mark) => mark.type === commentMark);

		this.storage.activeCommentId = activeCommentMark?.attrs.commentId || null;

		this.options.onCommentActivated(this.storage.activeCommentId!);
	},

	addStorage() {
		return {
			activeCommentId: null
		};
	},

	addCommands() {
		return {
			addComment:
				(commentId) =>
				({ commands }) => {
					if (!commentId) return false;
					commands.setMark('comment', { commentId });
					return true;
				},
			removeComment:
				(commentId) =>
				({ tr: transaction, dispatch }) => {
					if (!commentId) return false;

					const commentMarksWithRange: MarkWithRange[] = [];

					transaction.doc.descendants((node, pos) => {
						const commentMark = node.marks.find(
							(mark) => mark.type.name === 'comment' && mark.attrs.commentId === commentId
						);

						if (!commentMark) return;

						commentMarksWithRange.push({
							mark: commentMark,
							range: {
								from: pos,
								to: pos + node.nodeSize
							}
						});
					});

					commentMarksWithRange.forEach(({ mark, range }) => {
						transaction.removeMark(range.from, range.to, mark);
					});

					dispatch?.(transaction);
					this.options.onCommentRemoved(commentId);
					return true;
				},
			highlightComment:
				(commentId) =>
				({ tr: transaction, dispatch }) => {
					if (!commentId) return false;

					transaction.doc.descendants((node, pos) => {
						const commentMark = node.marks.find(
							(mark) => mark.type.name === 'comment' && mark.attrs.commentId === commentId
						);

						if (commentMark) {
							const mark = this.editor.schema.marks.comment.create({
								...commentMark.attrs,
								highlighted: true
							});
							transaction.addMark(pos, pos + node.nodeSize, mark);
						}
					});

					dispatch?.(transaction);
					return true;
				},
			unhighlightComment:
				(commentId) =>
				({ tr: transaction, dispatch }) => {
					if (!commentId) return false;

					transaction.doc.descendants((node, pos) => {
						const commentMark = node.marks.find(
							(mark) => mark.type.name === 'comment' && mark.attrs.commentId === commentId
						);

						if (commentMark) {
							const mark = this.editor.schema.marks.comment.create({
								...commentMark.attrs,
								highlighted: false
							});
							transaction.addMark(pos, pos + node.nodeSize, mark);
						}
					});

					dispatch?.(transaction);
					return true;
				}
		};
	}
});