#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 {
// A few px of breathing room around the equation.
constexpr int kRenderMargin = 2;
constexpr unsigned int kResolutionDpi = 96;
// The render cache is unbounded between clears; cap it so a very long
// session with many distinct equations cannot grow it forever.
constexpr int kCacheLimit = 512;

// Registers the embedded Latin Modern faces with the font database
// (once per process) and reports which families became available.
struct LatinModern {
	bool roman = false;
	bool math = false;
};

const LatinModern& loadLatinModern() {
	static const LatinModern fonts = [] {
		LatinModern result;
		// addApplicationFont(QByteArray) fails in this environment, so
		// the embedded bytes are staged to a per-process temp file and
		// registered through the (stable) file-based API.
		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)
// No parent: the renderer is used from pool threads, and a parented
// QObject would taint children it creates there.
: QObject(parent), m_renderer(std::make_shared<JKQTMathText>(
      nullptr, /* useFontsForGUI */ true)) {
	// Latin Modern is the default font of modern LaTeX; use the embedded
	// faces instead of whatever the system happens to have installed.
	const LatinModern& fonts = loadLatinModern();
	if (fonts.roman)
		m_renderer->setFontRomanAndMath(QStringLiteral("LMRoman10"),
		                                JKQTMathTextFontEncoding::MTFEUnicode);
	if (fonts.math) {
		// Same pattern as JKQTMathText's useXITS(): the OpenType math
		// font supplies the math alphabet and operators from its MATH
		// table.
		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 {

// Process-wide render cache; the key is the full render state, so
// re-opening a chat reuses already-rendered equations. Accessed on the
// GUI thread only.
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;
	}

	// A render is already running; it re-renders the latest state when
	// it completes (id mismatch), so there is nothing to do here.
	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;
	// The worker captures the renderer by value (shared_ptr) so it
	// stays alive even if this object is destroyed mid-render; it is
	// only ever used by the single in-flight worker (m_inFlight),
	// never concurrently.
	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()));
				// drawIntoImage renders at devicePixelRatio; convert
				// back to logical pixels.
				render.width = image.width() / dpr;
				render.height = image.height() / dpr;
				render.ok = true;
			}
		}
		// Deliver through the app instance (never destroyed) and
		// re-check the pointer on the GUI thread: posting to `this`
		// from the pool thread would race with its destruction.
		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) {
					// Superseded while the worker ran; render the
					// latest state.
					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