/*
 * wlopm - Wayland output power manager
 *
 * Copyright (C) 2021 Leon Henrik Plickat
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as published
 * by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <errno.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <wayland-client.h>

#ifdef __linux__
#include <features.h>
#ifdef __GLIBC__
#include<execinfo.h>
#endif
#endif

#include "wlr-output-power-management-unstable-v1.h"

#define VERSION "0.1.0"

const char usage[] =
	"Usage: wlopm [options...]\n"
	"  -j, --json                  Use JSON format.\n"
	"  -h, --help                  Print this help text and exit.\n"
	"  -v, --version               Print version and exit.\n"
	"      --on     <output-name>  Set the power mode of the specified output to on.\n"
	"      --off    <output-name>  Set the power mode of the specified output to off.\n"
	"      --toggle <output-name>  Toggle the power mode of the specified output.\n"
	"\n";

struct Output
{
	struct wl_list link;
	struct wl_output *wl_output;
	struct zwlr_output_power_v1 *wlr_output_power;
	enum zwlr_output_power_v1_mode mode;
	char *name;
	bool operation_failed;
	uint32_t global_name;
};

enum Action
{
	LIST,
	OPERATIONS,
};

enum Power_mode
{
	ON,
	OFF,
	TOGGLE,
};

struct Operation
{
	struct wl_list link;
	char *name;
	enum Power_mode power_mode;
};

bool json = false;
bool json_prev = false;

struct wl_display *wl_display = NULL;
struct wl_registry *wl_registry = NULL;
struct wl_callback *sync_callback = NULL;

struct wl_list outputs;
struct wl_list operations;

struct zwlr_output_power_manager_v1 *wlr_output_power_manager = NULL;

int ret = EXIT_SUCCESS;
bool loop = true;

static void wlr_output_power_handle_mode (void *data, struct zwlr_output_power_v1 *wlr_output_power,
		enum zwlr_output_power_v1_mode mode)
{
	struct Output *output = (struct Output *)data;
	output->mode = mode;
}

static void wlr_output_power_handle_failed (void *data, struct zwlr_output_power_v1 *wlr_output_power)
{
	struct Output *output = (struct Output *)data;
	output->operation_failed = true;
}

static const struct zwlr_output_power_v1_listener wlr_output_power_listener = {
	.mode   = wlr_output_power_handle_mode,
	.failed = wlr_output_power_handle_failed,
};

static void wl_output_handle_name (void *data, struct wl_output *wl_output,
		const char *name)
{
	struct Output *output = (struct Output *)data;
	if ( output->name != NULL )
		free(output->name);
	output->name = strdup(name);
}

static void wl_output_handle_geometry (void *data, struct wl_output *wl_output,
		int32_t x, int32_t y, int32_t physical_width, int32_t physical_height,
		int32_t subpixel, const char* make, const char* model, int32_t transform)
{
	/* This function is deliberately left empty. */
}

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)
{
	/* This function is deliberately left empty. */
}

static void wl_output_handle_scale (void *data, struct wl_output *wl_output,
		int32_t scale)
{
	/* This function is deliberately left empty. */
}

static void wl_output_handle_description (void *data, struct wl_output *wl_output,
		const char *description)
{
	/* This function is deliberately left empty. */
}

static void wl_output_handle_done (void *data, struct wl_output *wl_output)
{
	/* This function is deliberately left empty. */
}

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

static void registry_handle_global (void *data, struct wl_registry *registry,
		uint32_t name, const char *interface, uint32_t version)
{
	if ( strcmp(interface, wl_output_interface.name) == 0 )
	{
		if ( version < 4 )
		{
			fputs("ERROR: The compositor uses an outdated wl_output version.\n"
				"       Please inform the compositor developers so they can update to the latest version.\n", stderr);
			loop = false;
			return;
		}

		struct Output *output = calloc(1, sizeof(struct Output));
		if ( output == NULL )
		{
			fputs("ERROR: Failed to allocate.\n", stderr);
			return;
		}

		output->wl_output = wl_registry_bind(registry, name,
				&wl_output_interface, 4);
		wl_output_add_listener(output->wl_output, &wl_output_listener, output);
		output->wlr_output_power = NULL;
		output->name = NULL;
		output->global_name = name;

		wl_list_insert(&outputs, &output->link);
	}
	else if ( strcmp(interface, zwlr_output_power_manager_v1_interface.name) == 0 )
		wlr_output_power_manager = wl_registry_bind(registry, name,
				&zwlr_output_power_manager_v1_interface, version);
}

static void registry_handle_global_remove (void *data, struct wl_registry *registry,
    uint32_t name)
{
	/* We don't run long enough to care. */
}

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

static void sync_handle_done (void *data, struct wl_callback *wl_callback, uint32_t other_data);
static const struct wl_callback_listener sync_callback_listener = {
	.done = sync_handle_done,
};

static struct Output *output_from_name (const char *str)
{
	struct Output *output;
	wl_list_for_each(output, &outputs, link)
		if ( strcmp(output->name, str) == 0 )
			return output;
	return NULL;
}

static void output_set_power_mode (struct Output *output, enum Power_mode mode)
{
	enum zwlr_output_power_v1_mode new_mode = ZWLR_OUTPUT_POWER_V1_MODE_ON;
	switch (mode)
	{
		case ON:
			new_mode = ZWLR_OUTPUT_POWER_V1_MODE_ON;
			break;

		case OFF:
			new_mode = ZWLR_OUTPUT_POWER_V1_MODE_OFF;
			break;

		case TOGGLE:
			if ( output->mode == ZWLR_OUTPUT_POWER_V1_MODE_ON )
				new_mode = ZWLR_OUTPUT_POWER_V1_MODE_OFF;
			else
				new_mode = ZWLR_OUTPUT_POWER_V1_MODE_ON;
			break;
	}
	zwlr_output_power_v1_set_mode(output->wlr_output_power, new_mode);
}

static char *power_mode_to_string (enum zwlr_output_power_v1_mode mode)
{
	return mode == ZWLR_OUTPUT_POWER_V1_MODE_ON ? "on" : "off";
}

static void print_json_error (const char *output_name, const char *msg)
{
	fprintf(stdout,
			"%s\n    {\n"
			"      \"output\": \"%s\",\n"
			"      \"error\": \"%s\"\n"
			"    }",
			json_prev ? "," : "", output_name, msg);
	json_prev = true;
}

static void do_operation (struct Operation *operation)
{
	if ( *operation->name == '*' )
	{
		struct Output *output;
		wl_list_for_each(output, &outputs, link)
			output_set_power_mode(output, operation->power_mode);
	}
	else
	{
		struct Output *output = output_from_name(operation->name);
		if ( output == NULL )
		{
			if (json)
				print_json_error(operation->name, "output does not exist");
			else
				fprintf(stderr, "ERROR: Output '%s' does not exist.\n",
						operation->name);
			return;
		}
		output_set_power_mode(output, operation->power_mode);
	}
}

static void sync_handle_done (void *data, struct wl_callback *wl_callback, uint32_t other_data)
{
	wl_callback_destroy(wl_callback);
	sync_callback = NULL;

	static int sync = 0;
	if ( sync == 0 )
	{
		if ( wlr_output_power_manager == NULL )
		{
			fputs("ERROR: Wayland server does not support wlr-output-power-management-v1.\n", stderr);
			loop = false;
			ret = EXIT_FAILURE;
			return;
		}

		struct Output *output;
		wl_list_for_each(output, &outputs, link)
		{
			output->wlr_output_power = zwlr_output_power_manager_v1_get_output_power(
					wlr_output_power_manager, output->wl_output);
			zwlr_output_power_v1_add_listener(output->wlr_output_power,
					&wlr_output_power_listener, output);
		}

		sync_callback = wl_display_sync(wl_display);
		wl_callback_add_listener(sync_callback, &sync_callback_listener, NULL);
	}
	else if ( sync == 1 )
	{
		if (wl_list_empty(&operations))
		{
			/* The operations list is empty, so let's just list all
			 * outputs and their current power mode.
			 */
			struct Output *output;
			if (json)
			{
				fputs("[", stdout);
				wl_list_for_each(output, &outputs, link)
				{
					fprintf(stdout,
							"%s\n  {\n"
							"    \"output\": \"%s\",\n"
							"    \"power-mode\": \"%s\"\n"
							"  }",
							json_prev ? "," : "",
							output->name,
							power_mode_to_string(output->mode));
					json_prev = true;
				}
				fputs("\n]\n", stdout);
			}
			else
				wl_list_for_each(output, &outputs, link)
					fprintf(stdout, "%s %s\n", output->name,
							power_mode_to_string(output->mode));
			loop = false;
		}
		else
		{
			/* There are operations in the operations list. We have
			 * things to do!
			 */

			if (json)
				fputs(
						"{\n"
						"  \"errors\": [",
						stdout);

			struct Operation *operation;
			wl_list_for_each(operation, &operations, link)
				do_operation(operation);

			/* We need to sync yet another time because setting the
			 * power mode might fail and we want to display those
			 * error messages.
			 */
			sync_callback = wl_display_sync(wl_display);
			wl_callback_add_listener(sync_callback, &sync_callback_listener, NULL);
		}
	}
	else
	{
		struct Output *output;
		wl_list_for_each(output, &outputs, link)
			if (output->operation_failed)
			{
				if (json)
					print_json_error(output->name, "setting power mode failed");
				else
					fprintf(stderr, "ERROR: Setting power mode for output '%s' failed.\n",
							output->name);
			}

		if (json)
			fputs(
					"\n  ]\n"
					"}\n",
					stdout);
		loop = false;
	}
	sync++;
}

static void destroy_all_outputs (void)
{
	struct Output *output, *tmp;
	wl_list_for_each_safe(output, tmp, &outputs, link)
	{
		if ( output->wlr_output_power != NULL )
			zwlr_output_power_v1_destroy(output->wlr_output_power);
		wl_output_destroy(output->wl_output);
		wl_list_remove(&output->link);
		free(output->name);
		free(output);
	}
}

static void destroy_all_operations (void)
{
	struct Operation *operation, *tmp;
	wl_list_for_each_safe(operation, tmp, &operations, link)
	{
		wl_list_remove(&operation->link);
		free(operation->name);
		free(operation);
	}
}

static bool create_operation (const char *name, enum Power_mode power_mode)
{
	struct Operation *operation = calloc(1, sizeof(struct Operation));
	if ( operation == NULL )
	{
		fprintf(stderr, "ERROR: calloc: %s\n", strerror(errno));
		return false;
	}

	operation->name = strdup(name);
	if ( operation->name == NULL )
	{
		fprintf(stderr, "ERROR: calloc: %s\n", strerror(errno));
		free(operation);
		return false;
	}

	operation->power_mode = power_mode;

	wl_list_insert(&operations, &operation->link);
	return true;
}

/**
 * Intercept error signals (like SIGSEGV and SIGFPE) so that we can try to
 * print a fancy error message and a backtracke before letting the system kill us.
 */
static void handle_error (int signum)
{
	const char *msg =
		"\n"
		"┌──────────────────────────────────────────┐\n"
		"│                                          │\n"
		"│          wlopm has crashed.              │\n"
		"│                                          │\n"
		"│    This is likely a bug, so please       │\n"
		"│    report this to the mailing list.      │\n"
		"│                                          │\n"
		"│  ~leon_plickat/public-inbox@lists.sr.ht  │\n"
		"│                                          │\n"
		"└──────────────────────────────────────────┘\n"
		"\n";
	fputs(msg, stderr);

#ifdef __linux__
#ifdef __GLIBC__
	fputs("Attempting to get backtrace:\n", stderr);

	/* In some rare cases, getting a backtrace can also cause a segfault.
	 * There is nothing we can or should do about that. All hope is lost at
	 * that point.
	 */
	void *buffer[255];
	const int calls = backtrace(buffer, sizeof(buffer) / sizeof(void *));
	backtrace_symbols_fd(buffer, calls, fileno(stderr));
	fputs("\n", stderr);
#endif
#endif

	/* Let the default handlers deal with the rest. */
	signal(signum, SIG_DFL);
	kill(getpid(), signum);
}

/**
 * Set up signal handlers.
 */
static void init_signals (void)
{
	signal(SIGSEGV, handle_error);
	signal(SIGFPE, handle_error);

}

int main(int argc, char *argv[])
{
	init_signals();

	wl_list_init(&operations);
	for (int i = 1; i < argc; i++)
	{
		if ( strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0 )
		{
			destroy_all_operations();
			fputs(usage, stderr);
			return EXIT_SUCCESS;
		}
		if ( strcmp(argv[i], "-v") == 0 || strcmp(argv[i], "--version") == 0 )
		{
		{
			destroy_all_operations();
			fprintf(stderr, "wlopm version %s\n", VERSION);
			return EXIT_SUCCESS;
		}
		}
		else if ( strcmp(argv[i], "-j") == 0 || strcmp(argv[i], "--json") == 0 )
			json = true;
		else if ( strcmp(argv[i], "--on") == 0 )
		{
			if ( i == argc - 1 )
			{
				fputs("ERROR: '--on' needs an output name.\n", stderr);
				destroy_all_operations();
				return EXIT_FAILURE;
			}
			if (! create_operation(argv[i+1], ON))
			{
				destroy_all_operations();
				return EXIT_FAILURE;
			}
			i++;
		}
		else if ( strcmp(argv[i], "--off") == 0 )
		{
			if ( i == argc - 1 )
			{
				fputs("ERROR: '--off' needs an output name.\n", stderr);
				destroy_all_operations();
				return EXIT_FAILURE;
			}
			if (! create_operation(argv[i+1], OFF))
			{
				destroy_all_operations();
				return EXIT_FAILURE;
			}
			i++;
		}
		else if ( strcmp(argv[i], "--toggle") == 0 )
		{
			if ( i == argc - 1 )
			{
				fputs("ERROR: '--toggle' needs an output name.\n", stderr);
				destroy_all_operations();
				return EXIT_FAILURE;
			}
			if (! create_operation(argv[i+1], TOGGLE))
			{
				destroy_all_operations();
				return EXIT_FAILURE;
			}
			i++;
		}
		else
		{
			fprintf(stderr, "ERROR: Unknown option '%s'\n", argv[i]);
			destroy_all_operations();
			return EXIT_FAILURE;
		}
	}

	/* We query the display name here instead of letting wl_display_connect()
	 * figure it out itself, because libwayland (for legacy reasons) falls
	 * back to using "wayland-0" when $WAYLAND_DISPLAY is not set, which is
	 * generally not desirable.
	 */
	const char *display_name = getenv("WAYLAND_DISPLAY");
	if ( display_name == NULL )
	{
		fputs("ERROR: WAYLAND_DISPLAY is not set.\n", stderr);
		return EXIT_FAILURE;
	}

	wl_display = wl_display_connect(display_name);
	if ( wl_display == NULL )
	{
		fputs("ERROR: Can not connect to wayland display.\n", stderr);
		return EXIT_FAILURE;
	}

	wl_list_init(&outputs);

	wl_registry = wl_display_get_registry(wl_display);
	wl_registry_add_listener(wl_registry, &registry_listener, NULL);

	sync_callback = wl_display_sync(wl_display);
	wl_callback_add_listener(sync_callback, &sync_callback_listener, NULL);

	while ( loop && wl_display_dispatch(wl_display) > 0 );

	destroy_all_operations();
	destroy_all_outputs();

	if ( sync_callback != NULL )
		wl_callback_destroy(sync_callback);
	if ( wlr_output_power_manager != NULL )
		zwlr_output_power_manager_v1_destroy(wlr_output_power_manager);
	if ( wl_registry != NULL )
		wl_registry_destroy(wl_registry);
	wl_display_disconnect(wl_display);

	return ret;
}