#define EGL_EGLEXT_PROTOTYPES
#define GL_GLEXT_PROTOTYPES

#include "wallpaperservice.hpp"

#include <EGL/eglext.h>
#include <GL/gl.h>
#include <GL/glext.h>

#include <fcntl.h>
#include <poll.h>
#include <unistd.h>

#include <QDebug>
#include <QGuiApplication>
#include <QImageReader>
#include <QScreen>

#include <qguiapplication_platform.h>

#include <private/qwaylandscreen_p.h>

#include <cstring>

namespace ZShell::internal {

namespace {

const char* kVertexShader = R"(
#version 300 es
void main() {
	// The corners are (-1,-1), (3,-1) and (-1,3), which cover the whole buffer.
	vec2 corner = vec2(float((gl_VertexID << 1) & 2), float(gl_VertexID & 2));
	gl_Position = vec4(corner * 2.0 - 1.0, 0.0, 1.0);
}
)";

const char* kFragmentShader = R"(
#version 300 es
precision highp float;
uniform sampler2D source;
uniform vec2 size;
uniform vec4 crop;
out vec4 outColor;
void main() {
	// The pixels are stored top-down, so the fragment coordinate is flipped.
	vec2 p = vec2(gl_FragCoord.x, size.y - gl_FragCoord.y) / size;
	outColor = texture(source, crop.xy + p * crop.zw);
}
)";

constexpr int kWaitTimeoutMs = 20;

void logShader(uint32_t object, bool program) {
	int ok = 0;
	if (program)
		glGetProgramiv(object, GL_LINK_STATUS, &ok);
	else
		glGetShaderiv(object, GL_COMPILE_STATUS, &ok);

	if (ok) return;

	char info[512] = {};
	if (program)
		glGetProgramInfoLog(object, int(sizeof(info)), nullptr, info);
	else
		glGetShaderInfoLog(object, int(sizeof(info)), nullptr, info);

	qWarning() << "WallpaperService:" << (program ? "linking" : "compiling")
			   << "failed:" << info;
}

QVariantMap cropFor(const QVariantMap& crops, const QString& screen) {
	QVariantMap crop = crops.value(screen).toMap();

	if (crop.isEmpty()) {
		crop.insert("x", 0.0);
		crop.insert("y", 0.0);
		crop.insert("width", 1.0);
		crop.insert("height", 1.0);
		crop.insert("zoom", 1.0);
	}

	return crop;
}

} // namespace

WallpaperService::WallpaperService(QObject* parent) : QObject(parent) {}

WallpaperService::~WallpaperService() {
	m_stop = true;
	wake();

	if (m_thread.joinable()) m_thread.join();

	for (auto* surface : m_surfaces)
		delete surface;
	m_surfaces.clear();

	if (m_eglDisplay != EGL_NO_DISPLAY && m_context != EGL_NO_CONTEXT) {
		if (eglMakeCurrent(
				m_eglDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, m_context)) {
			if (m_texture) glDeleteTextures(1, &m_texture);
			if (m_program) glDeleteProgram(m_program);
		}

		eglDestroyContext(m_eglDisplay, m_context);
	}

	if (m_queue) wl_event_queue_destroy(m_queue);

	if (m_wake[0] >= 0) {
		close(m_wake[0]);
		close(m_wake[1]);
	}

	if (m_layerShell) m_layerShell->destroy();
	if (m_viewporter) m_viewporter->destroy();
	if (m_scaleManager) m_scaleManager->destroy();
	if (m_registry) wl_registry_destroy(m_registry);
}

QString WallpaperService::path() const {
	std::lock_guard<std::recursive_mutex> lock(m_mutex);
	return m_path;
}

void WallpaperService::setPath(const QString& path) {
	if (m_path == path) return;

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);
		m_path = path;

		m_forceDecode = true;
		m_imageChanged = true;
	}

	emit pathChanged();

	start();
	queueAll();
}

void WallpaperService::addScreen(const QString& screen) {
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		if (m_screens.contains(screen)) return;

		m_screens.append(screen);
	}

	start();

	queueAll();
}

void WallpaperService::removeScreen(const QString& screen) {
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		if (!m_screens.removeOne(screen)) return;
	}

	queueAll();
}

const QVariantMap& WallpaperService::crops() const {
	std::lock_guard<std::recursive_mutex> lock(m_mutex);
	return m_crops;
}

void WallpaperService::setCrops(const QVariantMap& crops) {
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		if (m_crops == crops) return;

		m_crops = crops;
	}

	emit cropsChanged();
	queueAll();
}

void WallpaperService::setCrop(const QString& screen, const QVariantMap& crop) {
	WallpaperSurface* surface = nullptr;

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		if (m_crops.value(screen) == crop) return;

		m_crops.insert(screen, crop);

		for (auto* candidate : m_surfaces) {
			if (candidate->screen() == screen) {
				surface = candidate;
				break;
			}
		}
	}

	emit cropsChanged();

	if (surface) queue(surface);
}

void WallpaperService::redraw() {
	queueAll();
}

void WallpaperService::start() {
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		if (started || m_path.isEmpty() || m_screens.isEmpty()) return;

		started = true;
	}

	auto* interface =
		qApp->nativeInterface<QNativeInterface::QWaylandApplication>();
	if (!interface) {
		qWarning() << "WallpaperService: this is not a Wayland session";
		return;
	}

	m_display = interface->display();
	m_compositor = interface->compositor();
	if (!m_display || !m_compositor) return;

	m_queue = wl_display_create_queue(m_display);
	if (!m_queue) return;

	if (pipe2(m_wake, O_CLOEXEC) < 0) {
		qWarning() << "WallpaperService: could not make the pipe to wake on";
		return;
	}

	bindGlobals();

	if (!setupGl()) return;

	m_drawing = true;
	m_thread = std::thread([this] { loop(); });

	wake();
}

void WallpaperService::wake() {
	if (m_wake[1] < 0) return;

	char byte = 0;

	if (::write(m_wake[1], &byte, 1) < 0)
		qWarning() << "WallpaperService: could not wake the drawing thread";
}

void WallpaperService::queue(WallpaperSurface* surface) {
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);
		m_pending.push_back(surface);
	}

	wake();
}

void WallpaperService::queueAll() {
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		for (auto* surface : m_surfaces)
			m_pending.push_back(surface);
	}

	wake();
}

void WallpaperService::loop() {
	while (!m_stop) {
		present();

		if (wl_display_prepare_read(m_display) == 0) {
			wl_display_flush(m_display);

			struct pollfd fds[2] = {
				{wl_display_get_fd(m_display), POLLIN, 0},
				{m_wake[0], POLLIN, 0},
			};

			int ready = poll(fds, 2, kWaitTimeoutMs);

			if (ready > 0 && fds[0].revents) {
				wl_display_read_events(m_display);
				continue;
			}

			if (ready > 0 && fds[1].revents) {
				char byte = 0;
				::read(m_wake[0], &byte, 1);
			}

			wl_display_cancel_read(m_display);
			continue;
		}

		struct pollfd fds[1] = {{m_wake[0], POLLIN, 0}};

		if (poll(fds, 1, kWaitTimeoutMs) <= 0) continue;

		char byte = 0;
		::read(m_wake[0], &byte, 1);
	}
}

void WallpaperService::present() {
	wl_display_dispatch_queue_pending(m_display, m_queue);

	addSurfaces();

	ensureDecodeSize();

	std::vector<WallpaperSurface*> pending;
	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);
		pending.swap(m_pending);
	}

	for (auto* surface : pending)
		drawSurface(surface);
}

void WallpaperService::addSurfaces() {
	if (!started || !m_layerShell || !m_viewporter || !m_scaleManager) return;

	QList<WallpaperSurface*> gone;
	QList<QString> needed;

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		QList<WallpaperSurface*> kept;
		for (auto* surface : m_surfaces) {
			if (m_screens.contains(surface->screen()))
				kept.append(surface);
			else
				gone.append(surface);
		}

		m_surfaces = kept;

		for (const auto& screen : m_screens) {
			bool exists = false;
			for (auto* surface : m_surfaces) {
				if (surface->screen() == screen) {
					exists = true;
					break;
				}
			}

			if (!exists) needed.append(screen);
		}
	}

	for (auto* surface : gone)
		delete surface;

	for (const auto& screen : needed)
		addSurface(screen);
}

void WallpaperService::registryGlobal(
	void* data,
	struct wl_registry* registry,
	uint32_t name,
	const char* interface,
	uint32_t version) {
	auto* service = static_cast<WallpaperService*>(data);

	void* proxy = nullptr;

	if (strcmp(interface, "zwlr_layer_shell_v1") == 0) {
		proxy = new QtWayland::zwlr_layer_shell_v1(
			registry, int(name), int(version));
	} else if (strcmp(interface, "wp_viewporter") == 0) {
		proxy = new QtWayland::wp_viewporter(registry, int(name), int(version));
	} else if (strcmp(interface, "wp_fractional_scale_manager_v1") == 0) {
		proxy = new QtWayland::wp_fractional_scale_manager_v1(
			registry, int(name), int(version));
	} else {
		return;
	}

	service->storeGlobal(interface, proxy);

	service->queueAll();
}

void WallpaperService::storeGlobal(const char* interface, void* proxy) {
	std::lock_guard<std::recursive_mutex> lock(m_mutex);

	if (strcmp(interface, "zwlr_layer_shell_v1") == 0)
		m_layerShell = static_cast<QtWayland::zwlr_layer_shell_v1*>(proxy);
	else if (strcmp(interface, "wp_viewporter") == 0)
		m_viewporter = static_cast<QtWayland::wp_viewporter*>(proxy);
	else if (strcmp(interface, "wp_fractional_scale_manager_v1") == 0)
		m_scaleManager =
			static_cast<QtWayland::wp_fractional_scale_manager_v1*>(proxy);
}

void WallpaperService::bindGlobals() {
	m_registry = wl_display_get_registry(m_display);

	static const struct wl_registry_listener listener = {
		registryGlobal,
		nullptr,
	};

	wl_registry_add_listener(m_registry, &listener, this);
}

bool WallpaperService::setupGl() {
	const char* extensions = eglQueryString(EGL_NO_DISPLAY, EGL_EXTENSIONS);
	if (extensions && strstr(extensions, "EGL_EXT_platform_wayland")) {
		m_eglDisplay =
			eglGetPlatformDisplay(EGL_PLATFORM_WAYLAND_EXT, m_display, nullptr);
	} else {
		m_eglDisplay =
			eglGetDisplay(reinterpret_cast<EGLNativeDisplayType>(m_display));
	}

	if (m_eglDisplay == EGL_NO_DISPLAY) {
		qWarning() << "WallpaperService: no EGL display";
		return false;
	}

	EGLint major = 0, minor = 0;
	if (!eglInitialize(m_eglDisplay, &major, &minor)) {
		qWarning() << "WallpaperService: eglInitialize failed";
		return false;
	}

	const EGLint configAttributes[] = {
		EGL_SURFACE_TYPE,
		EGL_WINDOW_BIT,
		EGL_RENDERABLE_TYPE,
		EGL_OPENGL_ES2_BIT,
		EGL_RED_SIZE,
		8,
		EGL_GREEN_SIZE,
		8,
		EGL_BLUE_SIZE,
		8,
		EGL_ALPHA_SIZE,
		8,
		EGL_NONE};

	EGLint count = 0;
	if (!eglChooseConfig(m_eglDisplay, configAttributes, &m_config, 1, &count) ||
		count == 0) {
		qWarning() << "WallpaperService: no EGL config for window surfaces";
		return false;
	}

	const EGLint contextAttributes[] = {EGL_CONTEXT_MAJOR_VERSION, 3, EGL_NONE};
	m_context = eglCreateContext(
		m_eglDisplay, m_config, EGL_NO_CONTEXT, contextAttributes);
	if (m_context == EGL_NO_CONTEXT) {
		qWarning() << "WallpaperService: eglCreateContext failed";
		return false;
	}

	return true;
}

uint32_t WallpaperService::createProgram() {
	uint32_t vertex = glCreateShader(GL_VERTEX_SHADER);
	glShaderSource(vertex, 1, &kVertexShader, nullptr);
	glCompileShader(vertex);
	logShader(vertex, false);

	uint32_t fragment = glCreateShader(GL_FRAGMENT_SHADER);
	glShaderSource(fragment, 1, &kFragmentShader, nullptr);
	glCompileShader(fragment);
	logShader(fragment, false);

	uint32_t program = glCreateProgram();
	glAttachShader(program, vertex);
	glAttachShader(program, fragment);
	glLinkProgram(program);
	logShader(program, true);

	if (glGetUniformLocation(program, "size") < 0 ||
		glGetUniformLocation(program, "crop") < 0) {
		qWarning() << "WallpaperService: the program does not have the "
					  "uniforms it needs";
		glDeleteProgram(program);
		return 0;
	}

	m_sizeUniform = glGetUniformLocation(program, "size");
	m_cropUniform = glGetUniformLocation(program, "crop");

	glDeleteShader(vertex);
	glDeleteShader(fragment);
	return program;
}

void WallpaperService::uploadImage() {
	if (m_image.isNull()) return;

	if (m_texture == 0) {
		glGenTextures(1, &m_texture);
		glBindTexture(GL_TEXTURE_2D, m_texture);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
		glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
	} else {
		glBindTexture(GL_TEXTURE_2D, m_texture);
	}

	glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
	glTexImage2D(
		GL_TEXTURE_2D,
		0,
		GL_RGBA8,
		m_image.width(),
		m_image.height(),
		0,
		GL_RGBA,
		GL_UNSIGNED_BYTE,
		m_image.constBits());
}

bool WallpaperService::decode(const QString& path) {
	if (path.isEmpty()) return false;

	QImageReader reader(path);
	reader.setAutoTransform(true);

	QSize size = reader.size();
	if (size.isValid()) m_sourceSize = size;

	if (!m_decodeSize.isEmpty() && !m_sourceSize.isEmpty())
		reader.setScaledSize(
			m_sourceSize.scaled(m_decodeSize, Qt::KeepAspectRatioByExpanding));

	QImage image = reader.read();
	if (image.isNull()) {
		qWarning() << "WallpaperService: could not read the image" << path;
		return false;
	}

	m_image = image.convertToFormat(QImage::Format_RGBA8888);
	return !m_image.isNull();
}

bool WallpaperService::ensureDecodeSize() {
	QSize largest;
	QString path;

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		largest = largestBufferSize();
		if (largest.isEmpty()) return false;

		if (!m_forceDecode && !m_decodeSize.isEmpty() &&
			qint64(largest.width()) * largest.height() <=
				qint64(m_decodeSize.width()) * m_decodeSize.height())
			return false;

		QSize target =
			m_sourceSize.isEmpty()
				? largest
				: m_sourceSize.scaled(largest, Qt::KeepAspectRatioByExpanding);

		if (!m_forceDecode && !m_image.isNull() && target == m_image.size())
			return false;

		m_forceDecode = false;
		m_decodeSize = largest;
		path = m_path;
	}

	if (!decode(path)) return false;

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);
		m_imageChanged = true;
	}

	return true;
}

QSize WallpaperService::largestBufferSize() {
	QSize largest;

	for (auto* surface : m_surfaces) {
		QSize buffer = surface->bufferSize();
		if (buffer.isEmpty()) continue;

		if (largest.isEmpty() || qint64(buffer.width()) * buffer.height() >
									 qint64(largest.width()) * largest.height())
			largest = buffer;
	}

	return largest;
}

QRectF WallpaperService::sourceRect(
	const QVariantMap& crops, const QString& screen, const QSize& size) const {
	QVariantMap crop = cropFor(crops, screen);

	QRectF rect(
		crop.value("x", 0.0).toDouble(),
		crop.value("y", 0.0).toDouble(),
		crop.value("width", 1.0).toDouble(),
		crop.value("height", 1.0).toDouble());

	if (size.isEmpty() || rect.isEmpty() || m_image.isNull()) return rect;

	qreal imageWidth = qreal(m_image.width());
	qreal imageHeight = qreal(m_image.height());

	QRectF pixels(
		rect.x() * imageWidth,
		rect.y() * imageHeight,
		rect.width() * imageWidth,
		rect.height() * imageHeight);

	qreal targetRatio = qreal(size.width()) / qreal(size.height());
	qreal rectRatio = pixels.width() / pixels.height();

	if (rectRatio > targetRatio) {
		qreal newWidth = pixels.height() * targetRatio;
		pixels.setX(pixels.x() + (pixels.width() - newWidth) / 2.0);
		pixels.setWidth(newWidth);
	} else if (rectRatio < targetRatio) {
		qreal newHeight = pixels.width() / targetRatio;
		pixels.setY(pixels.y() + (pixels.height() - newHeight) / 2.0);
		pixels.setHeight(newHeight);
	}

	return QRectF(
		pixels.x() / imageWidth,
		pixels.y() / imageHeight,
		pixels.width() / imageWidth,
		pixels.height() / imageHeight);
}

void WallpaperService::drawSurface(WallpaperSurface* surface) {
	EGLSurface target = surface->eglSurface();
	if (m_context == EGL_NO_CONTEXT || target == EGL_NO_SURFACE ||
		surface->bufferSize().isEmpty())
		return;

	if (!eglMakeCurrent(m_eglDisplay, target, target, m_context)) {
		qWarning() << "WallpaperService: the context cannot be made current";
		return;
	}

	// Release the context when this scope ends (see FrameService::drawSurface).
	// Holding it across the draw loop's sleep serializes the two raw-EGL clients on
	// NVIDIA, which makes eglSwapBuffers fail and starves other processes' contexts.
	struct ContextRelease {
		EGLDisplay display;
		~ContextRelease() {
			eglMakeCurrent(
				display, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
		}
	} contextRelease{m_eglDisplay};

	QVariantMap crops;
	QSize size;
	bool upload = false;

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);

		crops = m_crops;
		size = surface->bufferSize();

		upload = m_imageChanged;
		m_imageChanged = false;
	}

	if (size.isEmpty()) return;

	if (m_program == 0) m_program = createProgram();
	if (m_texture == 0 || upload) uploadImage();

	if (m_program == 0 || m_texture == 0 || m_image.isNull()) return;

	glUseProgram(m_program);
	glBindTexture(GL_TEXTURE_2D, m_texture);

	glUniform2f(
		GLint(m_sizeUniform), GLfloat(size.width()), GLfloat(size.height()));

	QRectF rect = sourceRect(crops, surface->screen(), size);
	glUniform4f(
		GLint(m_cropUniform),
		GLfloat(rect.x()),
		GLfloat(rect.y()),
		GLfloat(rect.width()),
		GLfloat(rect.height()));

	glViewport(0, 0, size.width(), size.height());
	glDrawArrays(GL_TRIANGLES, 0, 3);

	if (!eglSwapBuffers(m_eglDisplay, target))
		qWarning()
			<< "WallpaperService: eglSwapBuffers failed, eglGetError = 0x"
			<< QString::number(unsigned(eglGetError()), 16);
}

void WallpaperService::addSurface(const QString& screen) {
	if (!m_layerShell || !m_viewporter || !m_scaleManager) return;

	QScreen* qscreen = nullptr;
	for (auto* candidate : QGuiApplication::screens()) {
		if (candidate->name() == screen) {
			qscreen = candidate;
			break;
		}
	}

	if (!qscreen) {
		qWarning() << "WallpaperService:" << screen << "is not a screen";
		return;
	}

	auto* waylandScreen =
		dynamic_cast<QtWaylandClient::QWaylandScreen*>(qscreen->handle());
	if (!waylandScreen || !waylandScreen->output()) {
		qWarning() << "WallpaperService:" << screen
				   << "is not a Wayland screen";
		return;
	}

	wl_surface* wlSurface = wl_compositor_create_surface(m_compositor);
	if (!wlSurface) return;

	auto* surface =
		new WallpaperSurface(this, wlSurface, waylandScreen->output(), screen);

	{
		std::lock_guard<std::recursive_mutex> lock(m_mutex);
		m_surfaces.append(surface);
	}
}

} // namespace ZShell::internal