z-bar-qt/Plugins/ZShell/HyprPlugins/main.cpp

#include <any>
#include <atomic>
#include <chrono>
#include <filesystem>
#include <fstream>
#include <mutex>
#include <optional>
#include <sstream>
#include <string>
#include <thread>
#include <vector>

#define private public
#include <hyprland/src/plugins/PluginAPI.hpp>
#include <hyprland/src/render/OpenGL.hpp>
#include <hyprland/src/render/Renderer.hpp>
#undef private

#include <hyprland/src/managers/EventManager.hpp>

extern "C" {
#include <lua.h>
#include <lauxlib.h>
}

#include "colortemp.hpp"

using ZShell::services::easeInOutCubic;
using ZShell::services::kelvinToGain;

using Render::IFramebuffer;
using Render::GL::CHyprOpenGLImpl;
using Render::GL::g_pHyprOpenGL;

inline HANDLE PHANDLE = nullptr;

static std::filesystem::path g_runtimeDir{"/tmp"};
static std::filesystem::path g_cacheBase{"/tmp/.cache"};

static void initEnvPaths() {
	if (const char* env = getenv("XDG_RUNTIME_DIR"))
		g_runtimeDir = env;

	if (const char* env = getenv("XDG_CACHE_HOME")) {
		if (*env)
			g_cacheBase = env;
	} else if (const char* home = getenv("HOME")) {
		g_cacheBase = std::filesystem::path(home) / ".cache";
	}
}

static std::ofstream& debugLog() {
	static std::ofstream log(
		g_runtimeDir / "zshell-nightlight-debug.log",
		std::ios::app);
	return log;
}

static std::mutex g_logMutex;
static std::atomic<uint64_t> g_logSeq{0};

static void dlog(const std::string& msg) {
	std::lock_guard<std::mutex> lk(g_logMutex);
	auto& log = debugLog();
	log << "[" << g_logSeq++ << "][tid=" << std::this_thread::get_id() << "] "
		<< msg << std::endl;
}

static constexpr int BAKED_STEPS = 24;

static std::string glslFloat(float v) {
	std::ostringstream ss;
	ss << std::fixed << std::setprecision(6) << v;
	return ss.str();
}

static std::filesystem::path cacheDir() {
	return g_cacheBase / "zshell" / "night-light";
}

static std::filesystem::path shaderPath() {
	return cacheDir() / "nightlight.frag";
}

static std::filesystem::path bakedShaderPath(const std::array<float, 3>& gain) {
	return cacheDir() /
		   ("nightlight-" + glslFloat(gain[0]) + "-" + glslFloat(gain[1]) +
			"-" + glslFloat(gain[2]) + ".frag");
}

static std::string defaultShaderSource() {
	return "#version 100\n"
		   "precision highp float;\n"
		   "varying vec2 v_texcoord;\n"
		   "uniform sampler2D tex;\n"
		   "uniform vec3 tint;\n"
		   "void main() {\n"
		   "    vec4 pixColor = texture2D(tex, v_texcoord);\n"
		   "    gl_FragColor = vec4(pixColor.rgb * tint, pixColor.a);\n"
		   "}\n";
}

static std::string bakedShaderSource(const std::array<float, 3>& gain) {
	std::ostringstream src;
	src << "#version 100\n"
		   "precision highp float;\n"
		   "varying vec2 v_texcoord;\n"
		   "uniform sampler2D tex;\n"
		   "const vec3 GAIN = vec3("
		<< glslFloat(gain[0]) << "," << glslFloat(gain[1]) << ","
		<< glslFloat(gain[2])
		<< ");\n"
		   "void main() {\n"
		   "    vec4 pixColor = texture2D(tex, v_texcoord);\n"
		   "    gl_FragColor = vec4(pixColor.rgb * GAIN, pixColor.a);\n"
		   "}\n";
	return src.str();
}

static void ensureShaderFile(
	const std::filesystem::path& path, const std::string& source) {
	std::error_code ec;
	if (std::filesystem::is_regular_file(path, ec)) return;
	std::filesystem::create_directories(path.parent_path(), ec);
	std::ofstream out(path, std::ios::trunc);
	out << source;
}

static bool g_active = false;
static bool g_uniformMode = false;
static bool g_modeDecided = false;
static std::array<float, 3> g_currentGain{1.f, 1.f, 1.f};

static int g_currentKelvin = 6500;
static int g_lastKelvin = 6500;

struct STransition {
	bool active = false;
	std::chrono::steady_clock::time_point start;
	float durationSec = 0.4f;
	std::array<float, 3> fromGain{1.f, 1.f, 1.f};
	std::array<float, 3> toGain{1.f, 1.f, 1.f};
	int lastBakedStep = -1;
};
static STransition g_transition;

static void postStateEvent(bool enabled, bool animating);

static bool gainIsNeutral(const std::array<float, 3>& gain) {
	for (float v : gain)
		if (std::abs(v - 1.f) > 0.001f) return false;
	return true;
}

static CFunctionHook* g_pBeginHook = nullptr;
using origBegin = void (*)(
	CHyprOpenGLImpl*,
	PHLMONITOR,
	const CRegion&,
	SP<IFramebuffer>,
	std::optional<CRegion>);

static CFunctionHook* g_pSaveMirrorHook = nullptr;
using origSaveMirror = bool (*)(CHyprOpenGLImpl*, const CBox&);

static bool hkSaveBufferForMirror(CHyprOpenGLImpl* thisptr, const CBox& box) {
	const bool saved = thisptr->m_applyFinalShader;
	thisptr->m_applyFinalShader = false;
	const bool result = (*reinterpret_cast<origSaveMirror>(
		g_pSaveMirrorHook->m_original))(thisptr, box);
	thisptr->m_applyFinalShader = saved;
	return result;
}

static void installShader() {
	auto* gl = g_pHyprOpenGL.get();
	if (!gl)
		return;

	ensureShaderFile(shaderPath(), defaultShaderSource());
	gl->applyScreenShader(shaderPath().string());

	auto* shader = gl->m_finalScreenShader.get();
	if (shader && shader->program() >= 1 &&
	    shader->getUniformLocation(SHADER_TINT) != -1) {
		g_uniformMode = true;
	} else {
		g_uniformMode = false;
		if (!g_modeDecided)
			dlog(
				"shader has no 'tint' uniform, falling back to "
				"per-step recompiled shaders");
	}
	g_modeDecided = true;
}

static void setUniformGain(const std::array<float, 3>& gain) {
	auto* gl = g_pHyprOpenGL.get();
	auto* shader = gl ? gl->m_finalScreenShader.get() : nullptr;
	if (!gl || !shader)
		return;

	gl->useShader(gl->m_finalScreenShader);
	shader->setUniformFloat3(SHADER_TINT, gain[0], gain[1], gain[2]);
}

static void hkBegin(
	CHyprOpenGLImpl* thisptr,
	PHLMONITOR mon,
	const CRegion& damage,
	SP<IFramebuffer> fb,
	std::optional<CRegion> finalDamage) {
	(*reinterpret_cast<origBegin>(
		g_pBeginHook->m_original))(thisptr, mon, damage, fb, finalDamage);

	if (fb) {
		auto* renderer = g_pHyprRenderer.get();
		if (renderer)
			renderer->m_renderData.blockScreenShader = true;
	}

	if (!g_active) return;

	auto* gl = g_pHyprOpenGL.get();
	if (!gl) return;

	if (gl->m_finalScreenShader.get() &&
	    gl->m_finalScreenShader->program() < 1)
		installShader();
	if (!gl->m_finalScreenShader.get() ||
	    gl->m_finalScreenShader->program() < 1)
		return;

	if (g_transition.active) {
		auto* monPtr = mon.get();
		if (!monPtr)
			return;
		const float elapsed =
			std::chrono::duration<float>(
				std::chrono::steady_clock::now() - g_transition.start)
				.count();
		const float t =
			std::clamp(elapsed / g_transition.durationSec, 0.f, 1.f);
		const float e = easeInOutCubic(t);

		for (size_t i = 0; i < 3; ++i)
			g_currentGain[i] =
				g_transition.fromGain[i] +
				(g_transition.toGain[i] - g_transition.fromGain[i]) * e;

		if (g_uniformMode) {
			setUniformGain(g_currentGain);
		} else {
			const int step = static_cast<int>(e * BAKED_STEPS);
			if (step != g_transition.lastBakedStep) {
				g_transition.lastBakedStep = step;
				ensureShaderFile(
					bakedShaderPath(g_currentGain),
					bakedShaderSource(g_currentGain));
			gl->applyScreenShader(bakedShaderPath(g_currentGain).string());
		}
	}

		if (t >= 1.f) {
			g_transition.active = false;
			if (gainIsNeutral(g_currentGain)) {
				g_active = false;
				gl->applyScreenShader("");
			}
			postStateEvent(g_active, false);
		} else {
			monPtr->m_forceFullFrames = 2;
			monPtr->scheduleFrame();
		}
	} else if (g_uniformMode) {
		setUniformGain(g_currentGain);
	}
}

static constexpr int KELVIN_MIN = 2000;
static constexpr int KELVIN_MAX = 6500;

static int luaNightlightSet(lua_State* L) {
	const int rawKelvin = static_cast<int>(luaL_checknumber(L, 1));
	const float durationSec =
		lua_gettop(L) >= 2 ? static_cast<float>(luaL_checknumber(L, 2)) : 0.4f;

	const int targetKelvin = std::clamp(rawKelvin, KELVIN_MIN, KELVIN_MAX);
	if (targetKelvin != rawKelvin)
		dlog(
			"clamped " + std::to_string(rawKelvin) + "K to " +
			std::to_string(targetKelvin) + "K");

	g_transition.fromGain = g_currentGain;
	g_transition.toGain = kelvinToGain(targetKelvin);
	g_transition.active = true;
	g_transition.start = std::chrono::steady_clock::now();
	g_transition.durationSec = durationSec;
	g_transition.lastBakedStep = -1;
	g_active = true;

	g_lastKelvin = g_currentKelvin;
	g_currentKelvin = targetKelvin;

	postStateEvent(true, true);

	lua_pushinteger(L, targetKelvin);
	return 1;
}

static int luaNightlightDisable(lua_State* L) {
	if (!g_active) return 0;

	const float durationSec =
		lua_gettop(L) >= 1 ? static_cast<float>(luaL_checknumber(L, 1)) : 0.4f;

	g_transition.fromGain = g_currentGain;
	g_transition.toGain = {1.f, 1.f, 1.f};
	g_transition.active = true;
	g_transition.start = std::chrono::steady_clock::now();
	g_transition.durationSec = durationSec;
	g_transition.lastBakedStep = -1;

	g_lastKelvin = g_currentKelvin;
	g_currentKelvin = 6500;

	postStateEvent(false, true);

	return 0;
}

static int luaNightlightState(lua_State* L) {
	lua_pushstring(
		L,
		std::format(
			"{{\"enabled\":{},\"animating\":{},\"last\":{},\"current\":{}}}",
			g_active ? "true" : "false",
			g_transition.active ? "true" : "false",
			g_lastKelvin,
			g_currentKelvin)
			.c_str());
	return 1;
}

static void postStateEvent(bool enabled, bool animating) {
	if (!g_pEventManager) return;
	g_pEventManager->postEvent(
		SHyprIPCEvent{
			.event = "zshell-nightlight",
			.data = std::format(
				"{},{},{},{}",
				enabled ? 1 : 0,
				animating ? 1 : 0,
				g_lastKelvin,
				g_currentKelvin)});
}

static const SFunctionMatch* findExactMethod(
	const std::vector<SFunctionMatch>& methods,
	const std::string& classAndMethod) {
	for (const auto& m : methods) {
		if (m.signature.find(classAndMethod) != std::string::npos) return &m;
	}
	return nullptr;
}

static std::string manglePrefix(
	const std::string& className, const std::string& methodName) {
	return std::to_string(className.size()) + className +
		   std::to_string(methodName.size()) + methodName;
}

APICALL EXPORT PLUGIN_DESCRIPTION_INFO PLUGIN_INIT(HANDLE handle) {
	PHANDLE = handle;
	initEnvPaths();
	dlog("=== PLUGIN_INIT, build timestamp " __DATE__ " " __TIME__ " ===");

	const std::string serverHash = __hyprland_api_get_hash();
	const std::string clientHash = __hyprland_api_get_client_hash();

	if (serverHash != clientHash) {
		HyprlandAPI::addNotification(
			PHANDLE,
			"[zshell-nightlight] Mismatched headers! Can't proceed.",
			CHyprColor{1.0f, 0.2f, 0.2f, 1.0f},
			5000);
		throw std::runtime_error("[zshell-nightlight] version mismatch");
	}

	static const auto SAVEMIRROR_METHODS =
		HyprlandAPI::findFunctionsByName(PHANDLE, "saveBufferForMirror");
	const SFunctionMatch* saveMirrorTarget = findExactMethod(
		SAVEMIRROR_METHODS,
		manglePrefix("CHyprOpenGLImpl", "saveBufferForMirror"));

	if (saveMirrorTarget) {
		g_pSaveMirrorHook = HyprlandAPI::createFunctionHook(
			PHANDLE,
			saveMirrorTarget->address,
			reinterpret_cast<void*>(&hkSaveBufferForMirror));
		if (g_pSaveMirrorHook && g_pSaveMirrorHook->hook())
			dlog(
				"saveBufferForMirror hook installed: " +
				saveMirrorTarget->signature);
	}

	if (!g_pSaveMirrorHook)
		HyprlandAPI::addNotification(
			PHANDLE,
			"[zshell-nightlight] WARNING: couldn't hook saveBufferForMirror — "
			"captures may still show the nightlight",
			CHyprColor{1.f, 0.6f, 0.2f, 1.f},
			8000);

	static const auto METHODS =
		HyprlandAPI::findFunctionsByName(PHANDLE, "begin");
	const SFunctionMatch* target =
		findExactMethod(METHODS, manglePrefix("CHyprOpenGLImpl", "begin"));
	if (!target)
		throw std::runtime_error(
			"[zshell-nightlight] couldn't uniquely resolve "
			"CHyprOpenGLImpl::begin");

	g_pBeginHook = HyprlandAPI::createFunctionHook(
		PHANDLE, target->address, reinterpret_cast<void*>(&hkBegin));
	if (!g_pBeginHook || !g_pBeginHook->hook())
		throw std::runtime_error("[zshell-nightlight] failed to hook begin()");

	if (!HyprlandAPI::addLuaFunction(
			PHANDLE, "zshell", "nlSet", luaNightlightSet))
		throw std::runtime_error(
			"[zshell-nightlight] failed to register Lua function");

	if (!HyprlandAPI::addLuaFunction(
			PHANDLE, "zshell", "nlDisable", luaNightlightDisable))
		throw std::runtime_error(
			"[zshell-nightlight] failed to register Lua function");

	if (!HyprlandAPI::addLuaFunction(
			PHANDLE, "zshell", "nlState", luaNightlightState))
		throw std::runtime_error(
			"[zshell-nightlight] failed to register Lua function");

	return {
		"zshell-nightlight",
		"Animated nightlight, excluded from screencopy",
		"ZShell",
		"1.2"};
}


APICALL EXPORT std::string PLUGIN_API_VERSION() {
	return HYPRLAND_API_VERSION;
}


APICALL EXPORT void PLUGIN_EXIT() {}