/*
 * 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 as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * 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 <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>

#include <wayland-client.h>

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

#define VERSION "0.0.1"

const char usage[] =
	"Usage:\n"
	"\twlopm                     List outputs and their power modes.\n"
	"\twlopm on  <output-name>   Set output power mode to on.\n"
	"\twlopm off <output-name>   Set output power mode to off.\n"
	"\n";

enum Action
{
	LIST,
	ON,
	OFF,
};

enum Action action = LIST;
char *name = NULL;

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

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

struct wl_list outputs;

struct zxdg_output_manager_v1 *xdg_output_manager = NULL;
struct zwlr_output_power_manager_v1 *wlr_output_power_manager = NULL;

int ret = EXIT_SUCCESS;
bool loop = true;

static void noop () {}

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;
	fprintf(stderr, "ERROR: Setting mode for output \"%s\" failed.\n", output->name);
	loop = false;
	ret = EXIT_FAILURE;
}

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 xdg_output_handle_name (void *data, struct zxdg_output_v1 *xdg_output,
		const char *name)
{
	struct Output *output = (struct Output *)data;
	if ( output->name != NULL )
		free(output->name);
	output->name = strdup(name);
}

static const struct zxdg_output_v1_listener xdg_output_listener = {
	.logical_size     = noop,
	.name             = xdg_output_handle_name,
	.logical_position = noop,
	.description      = noop,
	.done             = noop, /* Deprecated since version 3. */
};

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 )
	{
		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, version);
		output->xdg_output = NULL;
		output->wlr_output_power = NULL;
		output->name = NULL;
		output->global_name = name;

		wl_list_insert(&outputs, &output->link);
	}
	else if ( strcmp(interface, zxdg_output_manager_v1_interface.name) == 0 )
		xdg_output_manager = wl_registry_bind(registry, name,
				&zxdg_output_manager_v1_interface, version);
	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 const struct wl_registry_listener registry_listener = {
	.global        = registry_handle_global,
	.global_remove = noop, /* We don't run long enough to care. */
};

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 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;
		}

		if ( xdg_output_manager == NULL )
		{
			fputs("ERROR: Wayland server does not support xdg-output.\n", stderr);
			loop = false;
			ret = EXIT_FAILURE;
			return;
		}

		struct Output *output;
		wl_list_for_each(output, &outputs, link)
		{
			output->xdg_output = zxdg_output_manager_v1_get_xdg_output(
					xdg_output_manager, output->wl_output);
			zxdg_output_v1_add_listener(output->xdg_output,
					&xdg_output_listener, output);

			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 ( action == LIST )
		{
			struct Output *output;
			wl_list_for_each(output, &outputs, link)
				fprintf(stdout, "%s %s\n", output->name,
						output->mode == ZWLR_OUTPUT_POWER_V1_MODE_ON ? "on" : "off");
			loop = false;
		}
		else
		{
			struct Output *output;
			wl_list_for_each(output, &outputs, link)
				if ( strcmp(output->name, name) == 0 )
					goto found;
			fprintf(stdout, "ERROR: No output with name \"%s\".\n", name);
			ret = EXIT_FAILURE;
			loop = false;
			return;
found:
			zwlr_output_power_v1_set_mode(output->wlr_output_power,
					action == ON ? ZWLR_OUTPUT_POWER_V1_MODE_ON : ZWLR_OUTPUT_POWER_V1_MODE_OFF);

			/* We need to sync yet another time because setting the
			 * power mode might fail.
			 */
			sync_callback = wl_display_sync(wl_display);
			wl_callback_add_listener(sync_callback, &sync_callback_listener, NULL);
		}
	}
	else
		loop = false;
	sync++;
}

int main(int argc, char *argv[])
{
	if ( argc == 3 )
	{
		if ( strcmp(argv[1], "on") == 0 )
			action = ON;
		else if ( strcmp(argv[1], "off") == 0 )
			action = OFF;
		else
		{
			fputs(usage, stderr);
			return EXIT_FAILURE;
		}

		name = strdup(argv[2]);
	}
	else if ( argc != 1 )
	{
		fputs(usage, stderr);
		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);
		if ( name != NULL )
			free(name);
		return EXIT_FAILURE;
	}

	wl_display = wl_display_connect(display_name);
	if ( wl_display == NULL )
	{
		fputs("ERROR: Can not connect to wayland display.\n", stderr);
		if ( name != NULL )
			free(name);
		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 );

	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);
		if ( output->xdg_output != NULL )
			zxdg_output_v1_destroy(output->xdg_output);
		wl_output_destroy(output->wl_output);
		wl_list_remove(&output->link);
		free(output->name);
		free(output);
	}

	if ( name != NULL )
		free(name);
	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 ( xdg_output_manager != NULL )
		zxdg_output_manager_v1_destroy(xdg_output_manager);
	if ( wl_registry != NULL )
		wl_registry_destroy(wl_registry);
	wl_display_disconnect(wl_display);

	return ret;
}