#pragma once

#include <QMutex>
#include <QObject>
#include <QString>
#include <QVariantList>
#include <QtQml>

#include <cstdint>
#include <memory>
#include <vector>

class QQmlEngine;
class QJSEngine;

namespace ZShell::llm {

// Syntax highlighting for LLM code blocks via tree-sitter.
//
// The tree-sitter runtime is linked. Grammar libraries are dlopen()'d
// lazily, so a missing grammar degrades that language to plain text
// instead of breaking the build or the app. At configure time CMake
// discovers installed grammars (system packages plus the parsers
// Neovim's nvim-treesitter installs) and pairs each with highlight
// queries (vendored, Neovim's, or fetched from
// tree-sitter/highlighting); both are embedded in the generated
// highlight-queries.hpp.
//
// For each grammar the first loadable (ABI-compatible) library wins
// and the first query that compiles against it wins, so a version
// skew between a library and its query degrades gracefully.
//
// The fence language the LLM wrote (```cpp, ```python, ...) is mapped
// to a grammar through an alias table; tags not in the table are used
// as grammar ids as-is.
//
// highlight() parses the code off the GUI thread and delivers a list of
// span maps by calling target's "onHighlightSpans(token, spans)" method
// (on the GUI thread):
//   { "start": int, "length": int, "kind": QString }
// where kind is a semantic role (keyword, string, comment, number,
// function, type, ...) that QML maps to theme colors. An empty list
// means "no highlighting" (unknown language or grammar not installed).
// token is passed back unchanged so the caller can drop results for
// superseded code; a destroyed target is simply skipped.
//
// Successful results are cached by (language, code). A request for
// unchanged code delivers the cached spans directly, without re-parsing
// — while a segment streams, only the grown tail is ever re-parsed.
class CodeHighlighter : public QObject {
	Q_OBJECT
	QML_ELEMENT
	QML_SINGLETON

	public:
	Q_INVOKABLE void highlight(
		const QString& code, const QString& language, QObject* target, int token);

	static CodeHighlighter* create(QQmlEngine*, QJSEngine*);

	struct Grammar {
		// Library candidates in priority order (system package, then
		// Neovim copies); libs[i] pairs with symbols[i].
		std::vector<std::string> libs;
		std::vector<std::string> symbols;
		// Candidate query sources in priority order; first that
		// compiles against the loaded grammar wins.
		std::vector<const char*> queries;
	};

	private:
	struct State {
		bool bad = false; // permanent failure, do not retry
		void* lib = nullptr;
		const void* lang = nullptr; // const TSLanguage*
		void* query = nullptr; // TSQuery*
	};

	[[nodiscard]] static const QHash<QString, QString>& aliases();
	// Maps a tree-sitter capture name to a role index (0 = unstyled).
	[[nodiscard]] static uint8_t roleFor(const char* name, uint32_t length);
	[[nodiscard]] static const char* roleName(uint8_t role);
	// Maps a fence language tag to the grammar id (see aliases()).
	[[nodiscard]] static QString resolveId(const QString& language);
	[[nodiscard]] static QString cacheKey(const QString& id, const QString& code);
	// The parsing work; runs on worker threads, so the per-language
	// state must be initialized under m_stateMutex and is shared as an
	// immutable object afterwards.
	[[nodiscard]] QVariantList doHighlight(const QString& code, const QString& language) const;
	// Exact-match span cache. lookupSpans() runs on the GUI thread,
	// storeSpans() on worker threads; both take m_cacheMutex.
	[[nodiscard]] QVariantList lookupSpans(const QString& code, const QString& language) const;
	void storeSpans(const QString& code, const QString& language,
	               const QVariantList& spans) const;

	struct SpanCacheEntry {
		QString code; // re-compared on lookup; a hash collision can
		             // never deliver the wrong spans
		QVariantList spans;
	};

	mutable QHash<QString, std::shared_ptr<const State>> m_states;
	mutable QMutex m_stateMutex;
	mutable QHash<QString, SpanCacheEntry> m_spanCache;
	mutable QStringList m_spanCacheOrder; // LRU order, oldest first
	mutable int m_spanCacheBytes = 0;
	mutable QMutex m_cacheMutex;
	static CodeHighlighter* s_instance;
};

} // namespace ZShell::llm