#include "markdownparser.hpp"

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

#include <QRegularExpression>
#include <QStringList>

namespace ZShell::llm {

QVector<MarkdownBlockData> MarkdownParser::parse(const QString& source) {
	QVector<MarkdownBlockData> 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 appendText = [&](const QString& text) {
		if (text.trimmed().isEmpty()) return;
		MarkdownBlockData block;
		block.type = MarkdownBlock::Type::Text;
		block.text = text;
		blocks.append(block);
	};

	auto appendMath = [&](const QString& latex) {
		if (latex.trimmed().isEmpty()) return;
		MarkdownBlockData block;
		block.type = MarkdownBlock::Type::Math;
		block.latex = latex;
		blocks.append(block);
	};

	auto appendCode = [&](const QString& language, const QString& code) {
		MarkdownBlockData block;
		block.type = MarkdownBlock::Type::Code;
		block.language = language;
		block.code = code;
		blocks.append(block);
	};

	auto appendHeading = [&](int level, const QString& text) {
		if (text.trimmed().isEmpty()) return;
		MarkdownBlockData block;
		block.type = MarkdownBlock::Type::Heading;
		block.level = level;
		block.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