import { Extension } from '@tiptap/core';
import type { Node } from '@tiptap/pm/model';
import { Plugin, PluginKey } from '@tiptap/pm/state';
import { Decoration, DecorationSet } from '@tiptap/pm/view';

const pluginKey = new PluginKey<DecorationSet>('nlpHighlighter');

interface NlpHighlighterOptions {
  getTerms: () => string[];
  getEnabled: () => boolean;
}

function escapeRegex(str: string): string {
  return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}

function buildDecorations(doc: Node, terms: string[], enabled: boolean): DecorationSet {
  if (!enabled || terms.length === 0) {
    return DecorationSet.empty;
  }

  // Build a single combined word-boundary regex for all terms (one pass per text node)
  const pattern = terms
    .filter((t) => t.length > 0)
    .map(escapeRegex)
    .join('|');

  if (!pattern) return DecorationSet.empty;

  const regex = new RegExp(`\\b(${pattern})\\b`, 'gi');
  const decorations: Decoration[] = [];

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

    regex.lastIndex = 0;
    let match: RegExpExecArray | null;

    while ((match = regex.exec(node.text)) !== null) {
      const start = pos + match.index;
      const end = start + match[0].length;
      decorations.push(Decoration.inline(start, end, { class: 'nlp-highlight' }));
    }
  });

  return DecorationSet.create(doc, decorations);
}

export const NlpHighlighterExtension = Extension.create<NlpHighlighterOptions>({
  name: 'nlpHighlighter',

  addOptions() {
    return {
      getTerms: () => [],
      getEnabled: () => false,
    };
  },

  addProseMirrorPlugins() {
    const { getTerms, getEnabled } = this.options;

    return [
      new Plugin({
        key: pluginKey,

        state: {
          init(_, state) {
            return buildDecorations(state.doc, getTerms(), getEnabled());
          },
          apply(tr, old) {
            // Only rebuild on explicit highlight update or doc change from external source
            if (tr.getMeta('updateHighlight')) {
              return buildDecorations(tr.doc, getTerms(), getEnabled());
            }
            // Map decoration positions cheaply on all other transactions (typing, cursor)
            return old.map(tr.mapping, tr.doc);
          },
        },

        props: {
          decorations(state) {
            return pluginKey.getState(state) ?? DecorationSet.empty;
          },
        },
      }),
    ];
  },
});