frontend/src/editor/extensions/citation.ts

import { Node, mergeAttributes } from "@tiptap/core";
import { ReactNodeViewRenderer } from "@tiptap/react";
import { Plugin, PluginKey } from "@tiptap/pm/state";
import { CitationView } from "../CitationView";

declare module "@tiptap/core" {
  interface Commands<ReturnType> {
    citation: {
      insertCitation: (key: string) => ReturnType;
    };
  }
}

export const citationLabelKey = new PluginKey("citation-labels");

type CitationStyle = "apa" | "ieee" | "vancouver" | "chicago-author-date" | "harvard1";
const NUMERIC_STYLES = new Set<CitationStyle>(["ieee", "vancouver"]);

function formatAuthorDate(entry: any): string {
  const authors = entry.author;
  if (!authors?.length) return `(${entry.id || "?"})`;
  const year = entry.issued?.["date-parts"]?.[0]?.[0] || "n.d.";
  const first = authors[0].family || authors[0].literal || "?";
  if (authors.length === 1) return `(${first}, ${year})`;
  if (authors.length === 2) {
    const second = authors[1].family || authors[1].literal || "?";
    return `(${first} & ${second}, ${year})`;
  }
  return `(${first} et al., ${year})`;
}

export const Citation = Node.create({
  name: "citation",
  group: "inline",
  inline: true,
  atom: true,

  addStorage() {
    return {
      citationsMap: null as any,
      settingsMap: null as any,
    };
  },

  addAttributes() {
    return {
      key: { default: "" },
      label: { default: "" },
    };
  },

  parseHTML() {
    return [{ tag: 'span[data-type="citation"]' }];
  },

  renderHTML({ HTMLAttributes }) {
    return [
      "span",
      mergeAttributes(HTMLAttributes, { "data-type": "citation" }),
    ];
  },

  addCommands() {
    return {
      insertCitation:
        (key: string) =>
        ({ commands }) => {
          return commands.insertContent({
            type: this.name,
            attrs: { key },
          });
        },
    };
  },

  addProseMirrorPlugins() {
    const ext = this;

    return [
      new Plugin({
        key: citationLabelKey,

        appendTransaction(transactions, _oldState, newState) {
          const docChanged = transactions.some((tr) => tr.docChanged);
          const signaled = transactions.some((tr) => tr.getMeta(citationLabelKey));
          if (!docChanged && !signaled) return null;

          const storage = ext.editor.storage.citation;
          const citationsMap = storage?.citationsMap;
          const settingsMap = storage?.settingsMap;
          const style = (settingsMap?.get("citationStyle") || "apa") as CitationStyle;

          // Single pass: collect all citations with their positions
          const order: string[] = [];
          const nodes: { pos: number; key: string; currentLabel: string }[] = [];

          newState.doc.descendants((node, pos) => {
            if (node.type.name === "citation" && node.attrs.key) {
              const k = node.attrs.key as string;
              if (!order.includes(k)) order.push(k);
              nodes.push({ pos, key: k, currentLabel: node.attrs.label as string });
            }
          });

          if (nodes.length === 0) return null;

          const updates: { pos: number; key: string; label: string }[] = [];

          for (const { pos, key, currentLabel } of nodes) {
            const entry = citationsMap?.get(key);
            let newLabel: string;

            if (!entry) {
              newLabel = `[${key}]`;
            } else if (NUMERIC_STYLES.has(style)) {
              const idx = order.indexOf(key);
              const num = idx >= 0 ? idx + 1 : "?";
              newLabel = style === "ieee" ? `[${num}]` : `(${num})`;
            } else {
              newLabel = formatAuthorDate(entry);
            }

            if (newLabel !== currentLabel) {
              updates.push({ pos, key, label: newLabel });
            }
          }

          if (updates.length === 0) return null;

          const tr = newState.tr;
          for (let i = updates.length - 1; i >= 0; i--) {
            const { pos, key, label } = updates[i];
            tr.setNodeMarkup(pos, undefined, { key, label });
          }
          return tr;
        },
      }),
    ];
  },

  addNodeView() {
    return ReactNodeViewRenderer(CitationView);
  },
});