z-bar-qt/Plugins/ZShell/Llm/markdownparser.cpp

#include "markdownparser.hpp"

#include "markdownblock.hpp"

#include <cmark-gfm-extension_api.h>
#include <cmark-gfm.h>

#include <QRegularExpression>
#include <QStringList>
#include <QVariantMap>

namespace ZShell::llm {

QVariantList MarkdownParser::parse(const QString& source) {
	QVariantList blocks;
	if (source.trimmed().isEmpty()) return blocks;

	const QStringList lines = source.split('\n');

	auto sliceSource = [&](int startLine, int endLine) -> QString {
		if (startLine < 1 || endLine < startLine) return QString();
		const int from = startLine;
		const int to = qMin(endLine, static_cast<int>(lines.size()));
		return lines.mid(from - 1, to - from + 1).join('\n').trimmed();
	};

	const QByteArray utf8 = source.toUtf8();
	cmark_node* doc = cmark_parse_document(
		utf8.constData(),
		static_cast<size_t>(utf8.size()),
		CMARK_OPT_DEFAULT | CMARK_OPT_SOURCEPOS);
	if (!doc) return blocks;

	const QRegularExpression mathRe(
		QStringLiteral("\\$\\$(.+?)\\$\\$"),
		QRegularExpression::DotMatchesEverythingOption);

	auto makeBlock = [&](LlmMarkdown::Type type) {
		QVariantMap block;
		block.insert("type", static_cast<int>(type));
		block.insert(
			"id",
			QString::number(static_cast<int>(blocks.size())) +
				QLatin1Char(':') + QString::number(static_cast<int>(type)));
		return block;
	};

	auto appendText = [&](const QString& text) {
		if (text.trimmed().isEmpty()) return;
		QVariantMap block = makeBlock(LlmMarkdown::Type::Text);
		block.insert("text", text);
		blocks.append(block);
	};

	auto appendMath = [&](const QString& latex) {
		if (latex.trimmed().isEmpty()) return;
		QVariantMap block = makeBlock(LlmMarkdown::Type::Math);
		block.insert("latex", latex);
		blocks.append(block);
	};

	auto appendCode = [&](const QString& language, const QString& code) {
		QVariantMap block = makeBlock(LlmMarkdown::Type::Code);
		block.insert("language", language);
		block.insert("code", code);
		blocks.append(block);
	};

	auto appendHeading = [&](int level, const QString& text) {
		if (text.trimmed().isEmpty()) return;
		QVariantMap block = makeBlock(LlmMarkdown::Type::Heading);
		block.insert("level", level);
		block.insert("text", text);
		blocks.append(block);
	};

	for (cmark_node* node = cmark_node_first_child(doc); node;
		 node = cmark_node_next(node)) {
		const cmark_node_type type = cmark_node_get_type(node);
		const int startLine = cmark_node_get_start_line(node);
		const int endLine = cmark_node_get_end_line(node);

		if (type == CMARK_NODE_CODE_BLOCK) {
			const char* literal = cmark_node_get_literal(node);
			QString code = literal ? QString::fromUtf8(literal) : QString();
			if (code.endsWith('\n')) code.chop(1);

			QString language;
			if (const char* info = cmark_node_get_fence_info(node); info)
				language = QString::fromUtf8(info)
							   .section(' ', 0, 0)
							   .trimmed()
							   .toLower();

			appendCode(language, code);
			continue;
		}

		if (type == CMARK_NODE_HEADING) {
			const char* content = cmark_node_get_string_content(node);
			appendHeading(
				cmark_node_get_heading_level(node),
				content ? QString::fromUtf8(content).trimmed() : QString());
			continue;
		}

		if (type == CMARK_NODE_PARAGRAPH) {
			const QString text = sliceSource(startLine, endLine);
			int cursor = 0;
			bool anyMath = false;
			for (auto it = mathRe.globalMatch(text, cursor); it.hasNext();
				 it = mathRe.globalMatch(text, cursor)) {
				const QRegularExpressionMatch m = it.next();
				anyMath = true;
				appendText(text.mid(
					cursor, static_cast<int>(m.capturedStart() - cursor)));
				appendMath(m.captured(1).trimmed());
				cursor = static_cast<int>(m.capturedEnd());
			}
			if (!anyMath)
				appendText(text);
			else
				appendText(text.mid(cursor));
			continue;
		}

		appendText(sliceSource(startLine, endLine));
	}

	cmark_node_free(doc);
	return blocks;
}

} // namespace ZShell::llm