#include "mathtext.hpp"

#include "latinmodern-fonts.hpp"

#include <jkqtmathtext/jkqtmathtext.h>

#include <QBuffer>
#include <QCoreApplication>
#include <QDir>
#include <QFile>
#include <QFontDatabase>
#include <QHash>
#include <QPointer>
#include <QThreadPool>

namespace ZShell::llm {

namespace {
constexpr int kRenderMargin = 2;
constexpr unsigned int kResolutionDpi = 96;
constexpr int kCacheLimit = 512;

struct LatinModern {
	bool roman = false;
	bool math = false;
};

const LatinModern& loadLatinModern() {
	static const LatinModern fonts = [] {
		LatinModern result;
		const auto add = [](const unsigned char* data,
							size_t size,
							const QString& fileName,
							const QString& family) {
			const QString path = QDir::tempPath() + QLatin1Char('/') + fileName;
			{
				QFile f(path);
				if (!f.open(QIODevice::WriteOnly) ||
					f.write(
						reinterpret_cast<const char*>(data),
						static_cast<qint64>(size)) != static_cast<qint64>(size))
					return false;
			}
			const int key = QFontDatabase::addApplicationFont(path);
			if (key < 0) return false;
			return QFontDatabase::applicationFontFamilies(key).contains(family);
		};
		result.roman = add(lmfont::lmroman10_regular,
						   sizeof(lmfont::lmroman10_regular),
						   QStringLiteral("lmroman10-regular.otf"),
						   QStringLiteral("LMRoman10")) &&
					   add(lmfont::lmroman10_italic,
						   sizeof(lmfont::lmroman10_italic),
						   QStringLiteral("lmroman10-italic.otf"),
						   QStringLiteral("LMRoman10")) &&
					   add(lmfont::lmroman10_bold,
						   sizeof(lmfont::lmroman10_bold),
						   QStringLiteral("lmroman10-bold.otf"),
						   QStringLiteral("LMRoman10")) &&
					   add(lmfont::lmroman10_bolditalic,
						   sizeof(lmfont::lmroman10_bolditalic),
						   QStringLiteral("lmroman10-bolditalic.otf"),
						   QStringLiteral("LMRoman10"));
		result.math =
			add(lmfont::latinmodern_math,
				sizeof(lmfont::latinmodern_math),
				QStringLiteral("latinmodern-math.otf"),
				QStringLiteral("Latin Modern Math"));
		return result;
	}();
	return fonts;
}

} // namespace

LlmMathText::LlmMathText(QObject* parent)
	: QObject(parent)
	, m_renderer(
		  std::make_shared<JKQTMathText>(nullptr, /* useFontsForGUI */ true)) {
	const LatinModern& fonts = loadLatinModern();
	if (fonts.roman)
		m_renderer->setFontRomanAndMath(
			QStringLiteral("LMRoman10"), JKQTMathTextFontEncoding::MTFEUnicode);
	if (fonts.math) {
		m_renderer->setFontMathRoman(
			QStringLiteral("Latin Modern Math"),
			JKQTMathTextFontEncoding::MTFEUnicode);
		m_renderer->setFallbackFontSymbols(
			QStringLiteral("Latin Modern Math"),
			JKQTMathTextFontEncoding::MTFEUnicode);
	}
}

void LlmMathText::setLatex(const QString& value) {
	if (m_latex == value) return;
	m_latex = value;
	reRender();
}

void LlmMathText::setColor(const QColor& value) {
	if (m_color == value) return;
	m_color = value;
	reRender();
}

void LlmMathText::setFontPointSize(double value) {
	if (qFuzzyCompare(m_fontPointSize, value)) return;
	m_fontPointSize = value;
	reRender();
}

void LlmMathText::setDevicePixelRatio(qreal value) {
	if (qFuzzyCompare(m_devicePixelRatio, value)) return;
	m_devicePixelRatio = value;
	reRender();
}

namespace {

struct MathRender {
	bool ok = false;
	QImage image;
	QUrl url;
	qreal width = 0;
	qreal height = 0;
};

QHash<QString, MathRender>& mathCache() {
	static QHash<QString, MathRender> cache;
	return cache;
}

} // namespace

void LlmMathText::reRender() {
	++m_requestId;
	if (m_latex.trimmed().isEmpty()) {
		m_image = QImage();
		m_imageUrl = QUrl();
		m_width = 0;
		m_height = 0;
		m_ok = false;
		Q_EMIT changed();
		return;
	}

	const QString key = m_latex + QLatin1Char(0x1f) + m_color.name() +
						QLatin1Char(0x1f) + QString::number(m_fontPointSize) +
						QLatin1Char(0x1f) + QString::number(m_devicePixelRatio);
	if (auto it = mathCache().find(key); it != mathCache().end()) {
		m_image = it->image;
		m_imageUrl = it->url;
		m_width = it->width;
		m_height = it->height;
		m_ok = it->ok;
		Q_EMIT changed();
		return;
	}

	if (m_inFlight) return;
	m_inFlight = true;

	const QString latex = m_latex;
	const QColor color = m_color;
	const double pointSize = m_fontPointSize;
	const qreal dpr = m_devicePixelRatio;
	auto renderer = m_renderer;
	QThreadPool::globalInstance()->start(
		[this, renderer, id = m_requestId, key, latex, color, pointSize, dpr]() {
			MathRender render;
			renderer->setFontPointSize(pointSize);
			renderer->setFontColor(color);
			if (renderer->parse(
					latex,
					JKQTMathText::LatexParser,
					JKQTMathText::DefaultParseOptions)) {
				const QImage image = renderer->drawIntoImage(
					/* drawBoxes */ false,
					QColor(Qt::transparent),
					kRenderMargin,
					dpr,
					kResolutionDpi);
				if (!image.isNull()) {
					QByteArray png;
					{
						QBuffer buffer(&png);
						buffer.open(QIODevice::WriteOnly);
						image.save(&buffer, "PNG");
					}
					render.image = image;
					render.url = QUrl(
						QStringLiteral("data:image/png;base64,") +
						QString::fromLatin1(png.toBase64()));
					render.width = image.width() / dpr;
					render.height = image.height() / dpr;
					render.ok = true;
				}
			}
			QPointer<LlmMathText> guard(this);
			QMetaObject::invokeMethod(
				QCoreApplication::instance(),
				[guard, id, key, render = std::move(render)]() mutable {
					LlmMathText* self = guard;
					if (!self) return;
					self->m_inFlight = false;
					if (id != self->m_requestId) {
						self->reRender();
						return;
					}
					if (render.ok) {
						auto& cache = mathCache();
						if (cache.size() >= kCacheLimit) cache.clear();
						cache.insert(key, render);
					}
					self->m_image = render.image;
					self->m_imageUrl = render.url;
					self->m_width = render.width;
					self->m_height = render.height;
					self->m_ok = render.ok;
					Q_EMIT self->changed();
				},
				Qt::QueuedConnection);
		});
}

} // namespace ZShell::llm