pi-ext

import {
	Alias,
	isMap,
	isScalar,
	isSeq,
	type Node,
	type Pair,
	parseDocument,
} from "yaml";

export type FrontmatterValue =
	| null
	| boolean
	| number
	| string
	| readonly FrontmatterValue[]
	| FrontmatterMap;
export type FrontmatterMap = ReadonlyMap<string, FrontmatterValue>;

export class MarkdownDocumentError extends Error {
	readonly code: string;

	constructor(code: string, message: string) {
		super(message);
		this.code = code;
		this.name = "MarkdownDocumentError";
	}
}

const lineEndingLength = (source: string, at: number): number => {
	if (source[at] === "\r" && source[at + 1] === "\n") return 2;
	return source[at] === "\r" || source[at] === "\n" ? 1 : 0;
};

// Extracts the leading YAML frontmatter block exactly like remark-frontmatter
// with ["yaml"] would: a `---` fence on line 1 (optional trailing spaces or
// tabs) closed by a `---` line, with one trailing line ending removed from the
// value. Lone carriage returns stay literal, matching the mdast value.
const frontmatterBlock = (source: string): string | undefined => {
	let at = source.charCodeAt(0) === 0xfeff ? 1 : 0;
	if (source.slice(at, at + 3) !== "---") return undefined;
	at += 3;
	while (source[at] === " " || source[at] === "\t") at += 1;
	if (lineEndingLength(source, at) === 0) return undefined;
	at += lineEndingLength(source, at);
	const valueStart = at;
	for (;;) {
		if (source.startsWith("---", at)) {
			let cursor = at + 3;
			while (source[cursor] === " " || source[cursor] === "\t") cursor += 1;
			if (cursor === source.length || lineEndingLength(source, cursor) > 0) {
				const value = source.slice(valueStart, at);
				return value.endsWith("\r\n")
					? value.slice(0, -2)
					: /\r|\n/.test(value.at(-1) ?? "")
						? value.slice(0, -1)
						: value;
			}
		}
		const eol = source.slice(at).search(/[\r\n]/);
		if (eol < 0) return undefined;
		at += eol + lineEndingLength(source, at + eol);
	}
};

const unsupported = (message: string): never => {
	throw new MarkdownDocumentError("frontmatter.unsupported", message);
};

const convertNode = (node: Node | null): FrontmatterValue => {
	if (node === null) return null;
	if (node instanceof Alias)
		return unsupported("YAML aliases are not frontmatter values");
	if (isScalar(node)) {
		const value = node.value;
		if (
			value === null ||
			typeof value === "boolean" ||
			typeof value === "string"
		)
			return value;
		if (typeof value === "number" && Number.isFinite(value)) return value;
		return unsupported(`Unsupported YAML scalar: ${String(value)}`);
	}
	if (isSeq(node))
		return node.items.map((item) => convertNode(item as Node | null));
	if (isMap(node)) return convertMap(node.items);
	return unsupported("Unsupported YAML node");
};

const convertMap = (items: readonly Pair[]): FrontmatterMap => {
	const result = new Map<string, FrontmatterValue>();
	for (const pair of items) {
		const key = pair.key as Node | null;
		const keyValue = isScalar(key) ? key.value : undefined;
		if (typeof keyValue === "string") {
			result.set(keyValue, convertNode(pair.value as Node | null));
			continue;
		}
		unsupported("Frontmatter map keys must be strings");
	}
	return result;
};

export function parseFrontmatter(source: string): FrontmatterMap {
	const block = frontmatterBlock(source);
	if (block === undefined)
		throw new MarkdownDocumentError(
			"frontmatter.missing",
			"Document requires one leading YAML frontmatter block",
		);
	const document = parseDocument(block, {
		schema: "core",
		strict: true,
		uniqueKeys: true,
	});
	if (document.errors.length > 0)
		throw new MarkdownDocumentError(
			"frontmatter.invalid",
			document.errors.map((error) => error.message).join("; "),
		);
	if (!isMap(document.contents))
		throw new MarkdownDocumentError(
			"frontmatter.shape",
			"Frontmatter must be a YAML map",
		);
	return convertMap(document.contents.items);
}