#define _DEFAULT_SOURCE
#define _XOPEN_SOURCE 700
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <unistd.h>
#include <wayland-client-protocol.h>
#include <wayland-client.h>

#include "color.h"
#include "wlr-gamma-control-unstable-v1-client-protocol.h"

#ifndef WLSUNSET_VERSION
#define WLSUNSET_VERSION "dev"
#endif

#define WSCT_GAMMA 1.0

struct app {
	struct wl_list outputs;
	struct zwlr_gamma_control_manager_v1 *gamma_control_manager;
	int temperature;
};

struct output {
	struct wl_list link;

	struct app *app;
	struct wl_output *wl_output;
	struct zwlr_gamma_control_v1 *gamma_control;

	uint32_t id;
	uint32_t ramp_size;
	int table_fd;
	uint16_t *table;
	bool applied;
};

static size_t gamma_table_size(uint32_t ramp_size) {
	return (size_t)ramp_size * 3 * sizeof(uint16_t);
}

static void reset_output_table(struct output *output) {
	if (output->table != NULL) {
		munmap(output->table, gamma_table_size(output->ramp_size));
		output->table = NULL;
	}
	if (output->table_fd != -1) {
		close(output->table_fd);
		output->table_fd = -1;
	}
	output->ramp_size = 0;
}

static void destroy_output(struct output *output) {
	if (output->gamma_control != NULL) {
		zwlr_gamma_control_v1_destroy(output->gamma_control);
		output->gamma_control = NULL;
	}
	reset_output_table(output);
	if (output->wl_output != NULL) {
		wl_output_destroy(output->wl_output);
		output->wl_output = NULL;
	}
	wl_list_remove(&output->link);
	free(output);
}

static int create_anonymous_file(off_t size) {
	char template[] = "/tmp/wsct-shared-XXXXXX";
	int fd = mkstemp(template);
	if (fd < 0) {
		return -1;
	}

	if (unlink(template) == -1) {
		close(fd);
		return -1;
	}

	int ret;
	do {
		errno = 0;
		ret = ftruncate(fd, size);
	} while (ret == -1 && errno == EINTR);

	if (ret == -1) {
		close(fd);
		return -1;
	}

	return fd;
}

static int create_gamma_table(uint32_t ramp_size, uint16_t **table) {
	size_t table_size = gamma_table_size(ramp_size);
	int fd = create_anonymous_file(table_size);
	if (fd < 0) {
		perror("failed to create anonymous file");
		return -1;
	}

	void *data = mmap(NULL, table_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
	if (data == MAP_FAILED) {
		perror("failed to mmap gamma table");
		close(fd);
		return -1;
	}

	*table = data;
	return fd;
}

static void fill_gamma_table(uint16_t *table, uint32_t ramp_size, double rw,
		double gw, double bw, double gamma) {
	uint16_t *r = table;
	uint16_t *g = table + ramp_size;
	uint16_t *b = table + 2 * ramp_size;

	for (uint32_t i = 0; i < ramp_size; ++i) {
		double val = (double)i / (double)(ramp_size - 1);
		r[i] = (uint16_t)(UINT16_MAX * pow(val * rw, 1.0 / gamma));
		g[i] = (uint16_t)(UINT16_MAX * pow(val * gw, 1.0 / gamma));
		b[i] = (uint16_t)(UINT16_MAX * pow(val * bw, 1.0 / gamma));
	}
}

static int output_apply_temperature(struct output *output) {
	if (output->gamma_control == NULL || output->table_fd == -1 || output->table == NULL) {
		return -1;
	}

	struct rgb wp = calc_whitepoint(output->app->temperature);
	fill_gamma_table(output->table, output->ramp_size, wp.r, wp.g, wp.b, WSCT_GAMMA);

	if (lseek(output->table_fd, 0, SEEK_SET) == -1) {
		perror("failed to rewind gamma table");
		return -1;
	}

	zwlr_gamma_control_v1_set_gamma(output->gamma_control, output->table_fd);
	output->applied = true;
	fprintf(stderr, "applied %d K to output %u\n", output->app->temperature, output->id);
	return 0;
}

static void gamma_control_handle_gamma_size(void *data,
		struct zwlr_gamma_control_v1 *gamma_control, uint32_t ramp_size) {
	(void)gamma_control;
	struct output *output = data;

	reset_output_table(output);
	if (ramp_size == 0) {
		fprintf(stderr, "output %u reported a zero-length gamma ramp\n", output->id);
		return;
	}

	output->ramp_size = ramp_size;
	output->table_fd = create_gamma_table(ramp_size, &output->table);
	if (output->table_fd < 0) {
		exit(EXIT_FAILURE);
	}

	if (output_apply_temperature(output) == -1) {
		exit(EXIT_FAILURE);
	}
}

static void gamma_control_handle_failed(void *data,
		struct zwlr_gamma_control_v1 *gamma_control) {
	(void)gamma_control;
	struct output *output = data;

	fprintf(stderr, "gamma control failed for output %u\n", output->id);
	zwlr_gamma_control_v1_destroy(output->gamma_control);
	output->gamma_control = NULL;
	reset_output_table(output);
	output->applied = false;
}

static const struct zwlr_gamma_control_v1_listener gamma_control_listener = {
	.gamma_size = gamma_control_handle_gamma_size,
	.failed = gamma_control_handle_failed,
};

static void setup_gamma_control(struct app *app, struct output *output) {
	if (app->gamma_control_manager == NULL || output->gamma_control != NULL) {
		return;
	}

	output->gamma_control = zwlr_gamma_control_manager_v1_get_gamma_control(
		app->gamma_control_manager, output->wl_output);
	zwlr_gamma_control_v1_add_listener(output->gamma_control, &gamma_control_listener, output);
}

static void wl_output_handle_geometry(void *data, struct wl_output *wl_output, int32_t x,
		int32_t y, int32_t phys_width, int32_t phys_height, int32_t subpixel,
		const char *make, const char *model, int32_t transform) {
	(void)data;
	(void)wl_output;
	(void)x;
	(void)y;
	(void)phys_width;
	(void)phys_height;
	(void)subpixel;
	(void)make;
	(void)model;
	(void)transform;
}

static void wl_output_handle_mode(void *data, struct wl_output *wl_output, uint32_t flags,
		int32_t width, int32_t height, int32_t refresh) {
	(void)data;
	(void)wl_output;
	(void)flags;
	(void)width;
	(void)height;
	(void)refresh;
}

static void wl_output_handle_done(void *data, struct wl_output *wl_output) {
	(void)data;
	(void)wl_output;
}

static void wl_output_handle_scale(void *data, struct wl_output *wl_output, int32_t factor) {
	(void)data;
	(void)wl_output;
	(void)factor;
}

static void wl_output_handle_name(void *data, struct wl_output *wl_output, const char *name) {
	(void)data;
	(void)wl_output;
	(void)name;
}

static void wl_output_handle_description(void *data, struct wl_output *wl_output,
		const char *description) {
	(void)data;
	(void)wl_output;
	(void)description;
}

static const struct wl_output_listener output_listener = {
	.geometry = wl_output_handle_geometry,
	.mode = wl_output_handle_mode,
	.done = wl_output_handle_done,
	.scale = wl_output_handle_scale,
	.name = wl_output_handle_name,
	.description = wl_output_handle_description,
};

static void registry_handle_global(void *data, struct wl_registry *registry, uint32_t name,
		const char *interface, uint32_t version) {
	struct app *app = data;
	(void)version;

	if (strcmp(interface, wl_output_interface.name) == 0) {
		struct output *output = calloc(1, sizeof(*output));
		if (output == NULL) {
			perror("failed to allocate output");
			exit(EXIT_FAILURE);
		}

		output->app = app;
		output->id = name;
		output->table_fd = -1;
		output->wl_output = wl_registry_bind(registry, name, &wl_output_interface, 1);
		wl_output_add_listener(output->wl_output, &output_listener, output);
		wl_list_insert(&app->outputs, &output->link);

		setup_gamma_control(app, output);
		return;
	}

	if (strcmp(interface, zwlr_gamma_control_manager_v1_interface.name) == 0) {
		app->gamma_control_manager = wl_registry_bind(registry, name,
			&zwlr_gamma_control_manager_v1_interface, 1);

		struct output *output;
		wl_list_for_each(output, &app->outputs, link) {
			setup_gamma_control(app, output);
		}
	}
}

static void registry_handle_global_remove(void *data, struct wl_registry *registry,
		uint32_t name) {
	(void)registry;
	struct app *app = data;
	struct output *output, *tmp;

	wl_list_for_each_safe(output, tmp, &app->outputs, link) {
		if (output->id == name) {
			destroy_output(output);
			return;
		}
	}
}

static const struct wl_registry_listener registry_listener = {
	.global = registry_handle_global,
	.global_remove = registry_handle_global_remove,
};

static int parse_temperature(const char *arg, int *temperature) {
	char *end = NULL;
	long value;

	errno = 0;
	value = strtol(arg, &end, 10);
	if (arg[0] == '\0' || end == arg || *end != '\0') {
		return -1;
	}
	if ((value == LONG_MIN || value == LONG_MAX) && errno == ERANGE) {
		return -1;
	}
	if (value <= 0 || value > INT_MAX) {
		return -1;
	}

	*temperature = (int)value;
	return 0;
}

static bool any_output_applied(const struct app *app) {
	const struct output *output;
	wl_list_for_each(output, &app->outputs, link) {
		if (output->applied) {
			return true;
		}
	}
	return false;
}

static void destroy_outputs(struct app *app) {
	struct output *output, *tmp;
	wl_list_for_each_safe(output, tmp, &app->outputs, link) {
		destroy_output(output);
	}
}

static const char usage[] =
	"usage: %s <temperature>\n"
	"       %s -h\n"
	"       %s -v\n"
	"\n"
	"Set the color temperature for all Wayland outputs and keep the process\n"
	"running so the gamma tables stay active.\n";

int main(int argc, char *argv[]) {
	if (argc == 2 && strcmp(argv[1], "-h") == 0) {
		printf(usage, argv[0], argv[0], argv[0]);
		return EXIT_SUCCESS;
	}

	if (argc == 2 && strcmp(argv[1], "-v") == 0) {
		printf("wsct version %s\n", WLSUNSET_VERSION);
		return EXIT_SUCCESS;
	}

	if (argc != 2) {
		fprintf(stderr, usage, argv[0], argv[0], argv[0]);
		return EXIT_FAILURE;
	}

	struct app app = { 0 };
	wl_list_init(&app.outputs);

	if (parse_temperature(argv[1], &app.temperature) == -1) {
		fprintf(stderr, "invalid temperature: %s\n", argv[1]);
		return EXIT_FAILURE;
	}

	struct wl_display *display = wl_display_connect(NULL);
	if (display == NULL) {
		fprintf(stderr, "failed to connect to the Wayland display\n");
		return EXIT_FAILURE;
	}

	struct wl_registry *registry = wl_display_get_registry(display);
	wl_registry_add_listener(registry, &registry_listener, &app);

	if (wl_display_roundtrip(display) == -1) {
		fprintf(stderr, "failed to query the Wayland registry\n");
		wl_registry_destroy(registry);
		wl_display_disconnect(display);
		return EXIT_FAILURE;
	}

	if (app.gamma_control_manager == NULL) {
		fprintf(stderr, "compositor doesn't support wlr-gamma-control-unstable-v1\n");
		wl_registry_destroy(registry);
		destroy_outputs(&app);
		wl_display_disconnect(display);
		return EXIT_FAILURE;
	}

	if (wl_list_empty(&app.outputs)) {
		fprintf(stderr, "no Wayland outputs found\n");
		zwlr_gamma_control_manager_v1_destroy(app.gamma_control_manager);
		wl_registry_destroy(registry);
		wl_display_disconnect(display);
		return EXIT_FAILURE;
	}

	struct output *output;
	wl_list_for_each(output, &app.outputs, link) {
		setup_gamma_control(&app, output);
	}

	if (wl_display_roundtrip(display) == -1 || wl_display_roundtrip(display) == -1) {
		fprintf(stderr, "failed to apply the requested temperature\n");
		zwlr_gamma_control_manager_v1_destroy(app.gamma_control_manager);
		wl_registry_destroy(registry);
		destroy_outputs(&app);
		wl_display_disconnect(display);
		return EXIT_FAILURE;
	}

	if (!any_output_applied(&app)) {
		fprintf(stderr, "failed to apply %d K to any output\n", app.temperature);
		zwlr_gamma_control_manager_v1_destroy(app.gamma_control_manager);
		wl_registry_destroy(registry);
		destroy_outputs(&app);
		wl_display_disconnect(display);
		return EXIT_FAILURE;
	}

	fprintf(stderr, "holding gamma controls open; terminate wsct to restore the original ramps\n");
	while (wl_display_dispatch(display) != -1) {
	}

	zwlr_gamma_control_manager_v1_destroy(app.gamma_control_manager);
	wl_registry_destroy(registry);
	destroy_outputs(&app);
	wl_display_disconnect(display);
	return EXIT_SUCCESS;
}