/**
 * 带 Prism 装饰的 Tiptap 代码块。
 * ProseMirror 文档只保存纯文本，语法高亮通过 Decoration 叠加，避免污染编辑内容。
 */
import CodeBlock from "@tiptap/extension-code-block";
import { findChildren, mergeAttributes } from "@tiptap/vue-3";
import type { Node as ProsemirrorNode } from "@tiptap/pm/model";
import type { EditorState, Transaction } from "@tiptap/pm/state";
import { Plugin, PluginKey } from "@tiptap/pm/state";
import { Decoration, DecorationSet } from "@tiptap/pm/view";
import Prism from "prismjs";
import { resolvePrismLanguage } from "../../utils/Prism";

/** Prism 分词后用于生成 Decoration 的扁平节点。 */
type HighlightNode = {
	text: string;
	classes: string[];
};

/** Prism Token 的递归结构。 */
type PrismToken = {
	type: string;
	alias?: string | string[];
	content: string | PrismToken | Array<string | PrismToken>;
};

/** 用于判断事务影响范围的最小结构。 */
type StepRange = {
	from?: number;
	to?: number;
};

/** 将语言名称转换为代码块 class 名。 */
function buildLanguageClassName(languageClassPrefix: string | null | undefined, language?: string | null) {
	if (!languageClassPrefix || !language) {
		return null;
	}
	return `${languageClassPrefix}${language}`;
}

/** 生成 Prism token 对应的 class 列表。 */
function buildTokenClassList(type: string, alias?: string | string[]) {
	const aliasList = Array.isArray(alias) ? alias : alias ? [alias] : [];
	return ["token", type, ...aliasList];
}

/** 将嵌套 Prism token 展平成文本片段，便于生成行内装饰。 */
function parsePrismTokens(tokens: Array<string | PrismToken>, className: string[] = []): HighlightNode[] {
	return tokens.flatMap((token) => {
		if (typeof token === "string") {
			return [{
				text: token,
				classes: className
			}];
		}

		const classes = [...className, ...buildTokenClassList(token.type, token.alias)];
		if (typeof token.content === "string") {
			return [{
				text: token.content,
				classes
			}];
		}
		if (Array.isArray(token.content)) {
			return parsePrismTokens(token.content, classes);
		}
		return parsePrismTokens([token.content], classes);
	});
}

/** 为文档中的全部代码块生成语法高亮装饰。 */
function getDecorations({
	doc,
	name,
	defaultLanguage
}: {
	doc: ProsemirrorNode;
	name: string;
	defaultLanguage: string | null | undefined;
}) {
	const decorations: Decoration[] = [];

	findChildren(doc, node => node.type.name === name).forEach((block) => {
		let from = block.pos + 1;
		const prismLanguage = resolvePrismLanguage(block.node.attrs.language || defaultLanguage);
		const grammar = prismLanguage ? Prism.languages[prismLanguage] : null;
		const tokens = grammar ? Prism.tokenize(block.node.textContent, grammar) : [block.node.textContent];

		parsePrismTokens(tokens).forEach((node) => {
			const to = from + node.text.length;
			if (node.classes.length) {
				decorations.push(Decoration.inline(from, to, {
					class: node.classes.join(" ")
				}));
			}
			from = to;
		});
	});

	return DecorationSet.create(doc, decorations);
}

/** 创建只在代码块内容或结构受影响时重新计算高亮的插件。 */
function createPrismPlugin({
	name,
	defaultLanguage
}: {
	name: string;
	defaultLanguage: string | null | undefined;
}) {
	const prismPlugin: Plugin<any> = new Plugin({
		key: new PluginKey("prism"),
		state: {
			init: (_: unknown, { doc }: EditorState) => getDecorations({
				doc,
				name,
				defaultLanguage
			}),
			apply: (transaction: Transaction, decorationSet: DecorationSet, oldState: EditorState, newState: EditorState) => {
				const oldNodeName = oldState.selection.$head.parent.type.name;
				const newNodeName = newState.selection.$head.parent.type.name;
				const oldNodes = findChildren(oldState.doc, node => node.type.name === name);
				const newNodes = findChildren(newState.doc, node => node.type.name === name);

				if (
					transaction.docChanged
					&& (
						[oldNodeName, newNodeName].includes(name)
						|| newNodes.length !== oldNodes.length
						|| transaction.steps.some((step) => {
							const range = step as StepRange;
							return range.from !== undefined
								&& range.to !== undefined
								&& oldNodes.some(node => node.pos >= range.from! && node.pos + node.node.nodeSize <= range.to!);
						})
					)
				) {
					return getDecorations({
						doc: transaction.doc,
						name,
						defaultLanguage
					});
				}

				return decorationSet.map(transaction.mapping, transaction.doc);
			}
		},
		props: {
			decorations(state: EditorState) {
				return prismPlugin.getState(state);
			}
		}
	});

	return prismPlugin;
}

/** 自定义代码块渲染和 Prism 装饰插件。 */
export default CodeBlock.extend({
	renderHTML({ node, HTMLAttributes }: { node: ProsemirrorNode; HTMLAttributes: Record<string, unknown> }) {
		const languageClassName = buildLanguageClassName(this.options.languageClassPrefix, node.attrs.language);

		return [
			"pre",
			mergeAttributes(this.options.HTMLAttributes, HTMLAttributes, {
				class: languageClassName
			}),
			[
				"code",
				{
					class: languageClassName
				},
				0
			]
		];
	},

	addProseMirrorPlugins() {
		return [
			createPrismPlugin({
				name: this.name,
				defaultLanguage: this.options.defaultLanguage
			})
		];
	}
});