#include "mathtext.hpp"

#include "latinmodern-fonts.hpp"

#include <jkqtmathtext/jkqtmathtext.h>

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

#include <QtQuick/QQuickImageProvider>

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;
}

void configureLatinModern(JKQTMathText& renderer) {
	const LatinModern& fonts = loadLatinModern();
	if (fonts.roman)
		renderer.setFontRomanAndMath(
			QStringLiteral("LMRoman10"), JKQTMathTextFontEncoding::MTFEUnicode);
	if (fonts.math) {
		renderer.setFontMathRoman(
			QStringLiteral("Latin Modern Math"),
			JKQTMathTextFontEncoding::MTFEUnicode);
		renderer.setFallbackFontSymbols(
			QStringLiteral("Latin Modern Math"),
			JKQTMathTextFontEncoding::MTFEUnicode);
	}
}

} // namespace

LlmMathText::LlmMathText(QObject* parent) : QObject(parent) {
	loadLatinModern();
}

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;
	QString id;
	qreal width = 0;
	qreal height = 0;
};

struct MathStore {
	QMutex mutex;
	QHash<QString, MathRender> entries;
	QStringList order;
	QHash<QString, QString> idIndex;
	QHash<QString, int> idRefs;
	qint64 bytes = 0;
	quint64 nextId = 1;
};

MathStore& mathStore() {
	static MathStore store;
	return store;
}

qint64 imageBytes(const QImage& image) {
	return image.isNull() ? 0 : static_cast<qint64>(image.sizeInBytes());
}

void retainMathId(const QString& id) {
	if (id.isEmpty()) return;
	auto& store = mathStore();
	QMutexLocker locker(&store.mutex);
	++store.idRefs[id];
}

void releaseMathId(const QString& id) {
	if (id.isEmpty()) return;
	auto& store = mathStore();
	QMutexLocker locker(&store.mutex);
	auto it = store.idRefs.find(id);
	if (it == store.idRefs.end()) return;
	if (--it.value() <= 0) store.idRefs.erase(it);
}

} // namespace

namespace {

constexpr qint64 kMaxCacheBytes = 64 * 1024 * 1024;

class MathImageProvider : public QQuickImageProvider {
	public:
	MathImageProvider() : QQuickImageProvider(QQuickImageProvider::Image) {}

	QImage requestImage(
		const QString& id, QSize* size, const QSize& requestedSize) override {
		Q_UNUSED(requestedSize)
		auto& store = mathStore();
		QMutexLocker locker(&store.mutex);
		const auto keyIt = store.idIndex.constFind(id);
		if (keyIt == store.idIndex.constEnd()) return {};
		const QString key = keyIt.value();
		const auto entryIt = store.entries.constFind(key);
		if (entryIt == store.entries.constEnd()) return {};
		if (size) *size = entryIt->image.size();
		const qsizetype pos = store.order.indexOf(key);
		if (pos >= 0) store.order.move(pos, store.order.size() - 1);
		return entryIt->image;
	}
};

void ensureImageProvider(QObject* owner) {
	if (auto* engine = qmlEngine(owner))
		if (!engine->imageProvider(QStringLiteral("llmmath")))
			engine->addImageProvider(
				QStringLiteral("llmmath"), new MathImageProvider());
}

} // namespace

void LlmMathText::reRender() {
	ensureImageProvider(this);
	++m_requestId;
	if (m_latex.trimmed().isEmpty()) {
		releaseMathId(m_imageId);
		m_imageId.clear();
		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);
	bool cacheHit = false;
	QString hitId;
	qreal hitW = 0;
	qreal hitH = 0;
	bool hitOk = false;
	{
		auto& store = mathStore();
		QMutexLocker locker(&store.mutex);
		const auto it = store.entries.constFind(key);
		if (it != store.entries.constEnd()) {
			cacheHit = true;
			hitId = it->id;
			hitW = it->width;
			hitH = it->height;
			hitOk = it->ok;
		}
	}
	if (cacheHit) {
		releaseMathId(m_imageId);
		retainMathId(hitId);
		m_imageId = hitId;
		m_width = hitW;
		m_height = hitH;
		m_ok = hitOk;
		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;
	QThreadPool::globalInstance()->start(
		[this, id = m_requestId, key, latex, color, pointSize, dpr]() {
			JKQTMathText renderer(nullptr, true);
			configureLatinModern(renderer);
			MathRender render;
			renderer.setFontPointSize(pointSize);
			renderer.setFontColor(color);
			if (renderer.parse(
					latex,
					JKQTMathText::LatexParser,
					JKQTMathText::DefaultParseOptions)) {
				{
					QImage scratch(1, 1, QImage::Format_ARGB32);
					QPainter painter(&scratch);
					const QSizeF size = renderer.getSize(painter);
					render.width = size.width() + 2 * kRenderMargin;
					render.height = size.height() + 2 * kRenderMargin;
				}
				QPointer<LlmMathText> sizeGuard(this);
				const qreal sizeW = render.width;
				const qreal sizeH = render.height;
				QMetaObject::invokeMethod(
					QCoreApplication::instance(),
					[sizeGuard, id, sizeW, sizeH]() {
						LlmMathText* self = sizeGuard;
						if (!self || id != self->m_requestId) return;
						self->m_width = sizeW;
						self->m_height = sizeH;
						Q_EMIT self->changed();
					},
					Qt::QueuedConnection);

				const QImage image = renderer.drawIntoImage(
					false,
					QColor(Qt::transparent),
					kRenderMargin,
					dpr,
					kResolutionDpi);
				if (!image.isNull()) {
					render.image = image;
					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& store = mathStore();
						QMutexLocker locker(&store.mutex);
						const bool isNew = !store.entries.contains(key);
						if (isNew)
							render.id = QString::number(store.nextId++);
						else
							render.id = store.entries[key].id;
						store.idIndex.insert(render.id, key);
						if (isNew) {
							const qint64 newBytes = imageBytes(render.image);
							while (store.order.size() >= kCacheLimit ||
								   store.bytes + newBytes > kMaxCacheBytes) {
								bool evicted = false;
								for (int i = 0; i < store.order.size(); ++i) {
									const QString oldKey = store.order.at(i);
									const auto oldIt =
										store.entries.constFind(oldKey);
									if (oldIt == store.entries.constEnd()) {
										store.order.removeAt(i);
										continue;
									}
									if (store.idRefs.value(oldIt->id, 0) > 0)
										continue;
									store.order.removeAt(i);
									store.bytes -= imageBytes(oldIt->image);
									store.idIndex.remove(oldIt->id);
									store.entries.remove(oldKey);
									evicted = true;
									break;
								}
								if (!evicted) break;
							}
							store.entries.insert(key, render);
							store.order.append(key);
							store.bytes += imageBytes(render.image);
						}
					}
					releaseMathId(self->m_imageId);
					retainMathId(render.id);
					self->m_imageId = render.id;
					self->m_width = render.width;
					self->m_height = render.height;
					self->m_ok = render.ok;
					Q_EMIT self->changed();
				},
				Qt::QueuedConnection);
		});
}

LlmMathText::~LlmMathText() {
	releaseMathId(m_imageId);
}

} // namespace ZShell::llm