/*
 * swmctl.c
 * Control and inspect swm's state.
 */
#include <errno.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <wayland-client.h>

#include "ext-foreign-toplevel-list-v1-client-protocol.h"
#include "swm-toplevel-v1-client-protocol.h"
#include "swm-workspace-v1-client-protocol.h"

#define MAX_WORKSPACES      32
#define MAX_WORKSPACE_TITLE 256

typedef struct {
    bool     seen;
    uint32_t flags;
    char     title[MAX_WORKSPACE_TITLE];
    uint32_t color;
} workspace_t;

/* Wayland globals required by one command family. */
enum command_kind {
    COMMAND_WORKSPACE,
    COMMAND_WINDOW_LIST,
    COMMAND_WINDOW_ACTIVATE,
};

/* Cached state for one ext-foreign-toplevel handle. */
typedef struct {
    struct wl_list                         link;
    struct ext_foreign_toplevel_handle_v1 *handle;
    char                                  *title;
    char                                  *app_id;
    char                                   identifier[33];
    bool                                   closed;
} window_t;

static struct wl_display                   *display;
static struct swm_workspace_manager_v1     *manager;
static struct ext_foreign_toplevel_list_v1 *toplevel_list;
static struct swm_toplevel_manager_v1      *toplevel_control;
static struct wl_list                       windows;
static workspace_t                          workspaces[MAX_WORKSPACES];
static bool                                 subscribing;
static bool                                 waybar;
static bool                                 toplevel_list_finished;
static enum command_kind                    requested_kind;

[[noreturn]] static void usage(void) {
    fprintf(
        stderr,
        "usage:\n"
        "  swmctl workspace list\n"
        "  swmctl workspace get N (title | color | selected)\n"
        "  swmctl workspace set N (title TITLE | color COLOR)\n"
        "  swmctl workspace clear N (title | color)\n"
        "  swmctl workspace subscribe [--format=waybar]\n"
        "  swmctl window list\n"
        "  swmctl window activate IDENTIFIER\n"
    );
    exit(EXIT_FAILURE);
}

/* Print a color in its shortest lossless notation. */
static void print_color(uint32_t color) {
    if ((color & 0xff) == 0xff)
        printf("#%06" PRIx32, color >> 8);
    else
        printf("#%08" PRIx32, color);
}

/* Write a string with JSON quoting. */
static void json_string(const char *value) {
    const unsigned char *p;

    putchar('"');
    for (p = (const unsigned char *)value; *p; p++) {
        if (*p == '"' || *p == '\\')
            printf("\\%c", *p);
        else if (*p == '\n')
            fputs("\\n", stdout);
        else if (*p == '\r')
            fputs("\\r", stdout);
        else if (*p == '\t')
            fputs("\\t", stdout);
        else if (*p < 0x20)
            printf("\\u%04x", *p);
        else
            putchar(*p);
    }
    putchar('"');
}

/* Print one workspace as a JSON object. */
static void print_json_workspace(uint32_t number, workspace_t *ws) {
    printf("{\"workspace\":%" PRIu32 ",\"title\":", number);
    if (ws->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_TITLE)
        json_string(ws->title);
    else
        fputs("null", stdout);
    fputs(",\"color\":", stdout);
    if (ws->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_COLOR) {
        putchar('"');
        print_color(ws->color);
        putchar('"');
    } else {
        fputs("null", stdout);
    }
    printf(
        ",\"selected\":%s}",
        ws->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_SELECTED ? "true" : "false"
    );
}

/* Write a JSON string safe for use as Pango markup in Waybar. */
static void json_markup_string(const char *value, bool has_color, uint32_t color) {
    const unsigned char *p;

    putchar('"');
    if (has_color) {
        printf("<span foreground=\\\"#%06" PRIx32 "\\\"", color >> 8);
        if ((color & 0xff) != 0xff)
            printf(" foreground_alpha=\\\"%" PRIu32 "\\\"", (color & 0xff) * 257);
        putchar('>');
    }
    for (p = (const unsigned char *)value; *p; p++) {
        if (*p == '&')
            fputs("&amp;", stdout);
        else if (*p == '<')
            fputs("&lt;", stdout);
        else if (*p == '>')
            fputs("&gt;", stdout);
        else if (*p == '"')
            fputs("&quot;", stdout);
        else if (*p == '\'')
            fputs("&apos;", stdout);
        else if (*p == '\\')
            fputs("\\\\", stdout);
        else if (*p == '\n')
            fputs("\\n", stdout);
        else if (*p == '\r')
            fputs("\\r", stdout);
        else if (*p == '\t')
            fputs("\\t", stdout);
        else if (*p < 0x20)
            printf("\\u%04x", *p);
        else
            putchar(*p);
    }
    if (has_color)
        fputs("</span>", stdout);

    putchar('"');
}

/* Print the selected workspace label as a Waybar custom-module update. */
static void print_waybar(void) {
    workspace_t *selected = nullptr;
    uint32_t     number   = 0;
    bool         has_title;

    for (uint32_t i = 0; i < MAX_WORKSPACES; i++) {
        if (workspaces[i].seen &&
            workspaces[i].flags & SWM_WORKSPACE_MANAGER_V1_METADATA_SELECTED) {
            selected = &workspaces[i];
            number   = i + 1;
            break;
        }
    }
    has_title = selected && selected->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_TITLE;
    fputs("{\"text\":", stdout);
    json_markup_string(
        has_title ? selected->title : "",
        has_title && selected->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_COLOR,
        selected ? selected->color : 0
    );
    printf(",\"tooltip\":\"Workspace %" PRIu32 "\"}\n", number);
    fflush(stdout);
}

/* Print the complete cached workspace metadata as one JSON value. */
static void print_json_workspaces(void) {
    bool first = true;

    putchar('[');
    for (uint32_t i = 0; i < MAX_WORKSPACES; i++) {
        if (!workspaces[i].seen)
            continue;
        if (!first)
            putchar(',');
        print_json_workspace(i + 1, &workspaces[i]);
        first = false;
    }
    puts("]");
    fflush(stdout);
}

/* Cache one workspace metadata event. */
static void metadata(
    void                            *data,
    struct swm_workspace_manager_v1 *object,
    uint32_t                         number,
    uint32_t                         flags,
    const char                      *title,
    uint32_t                         color
) {
    workspace_t *ws;

    if (!number || number > MAX_WORKSPACES)
        return;
    ws        = &workspaces[number - 1];
    ws->seen  = true;
    ws->flags = flags;
    ws->color = color;

    snprintf(ws->title, sizeof(ws->title), "%s", title);
}

/* Finish a snapshot, printing it when following changes. */
static void done(void *data, struct swm_workspace_manager_v1 *object) {
    if (!subscribing)
        return;
    if (waybar)
        print_waybar();
    else
        print_json_workspaces();
}

static const struct swm_workspace_manager_v1_listener manager_listener = {
    .metadata = metadata,
    .done     = done,
};

/* Replace one cached toplevel string. */
static void replace_string(char **target, const char *value) {
    char *copy = strdup(value);

    if (!copy) {
        fprintf(stderr, "swmctl: out of memory\n");
        exit(EXIT_FAILURE);
    }
    free(*target);
    *target = copy;
}

/* Release a closed toplevel handle. */
static void window_closed(void *data, struct ext_foreign_toplevel_handle_v1 *handle) {
    window_t *window = data;

    ext_foreign_toplevel_handle_v1_destroy(handle);
    window->handle = nullptr;
    window->closed = true;
}

/* Accept a complete toplevel state update. */
static void window_done(void *data, struct ext_foreign_toplevel_handle_v1 *handle) {}

/* Cache a toplevel title. */
static void window_title(
    void *data, struct ext_foreign_toplevel_handle_v1 *handle, const char *title
) {
    window_t *window = data;

    replace_string(&window->title, title);
}

/* Cache a toplevel application ID. */
static void window_app_id(
    void *data, struct ext_foreign_toplevel_handle_v1 *handle, const char *app_id
) {
    window_t *window = data;

    replace_string(&window->app_id, app_id);
}

/* Cache a toplevel's stable identifier. */
static void window_identifier(
    void *data, struct ext_foreign_toplevel_handle_v1 *handle, const char *identifier
) {
    window_t *window = data;

    if (strlen(identifier) >= sizeof(window->identifier)) {
        fprintf(stderr, "swmctl: compositor sent an invalid toplevel identifier\n");
        exit(EXIT_FAILURE);
    }
    strcpy(window->identifier, identifier);
}

static const struct ext_foreign_toplevel_handle_v1_listener window_listener = {
    .closed     = window_closed,
    .done       = window_done,
    .title      = window_title,
    .app_id     = window_app_id,
    .identifier = window_identifier,
};

/* Cache a toplevel announced by ext-foreign-toplevel-list. */
static void window_announced(
    void                                  *data,
    struct ext_foreign_toplevel_list_v1   *list,
    struct ext_foreign_toplevel_handle_v1 *handle
) {
    window_t *window = calloc(1, sizeof(*window));

    if (!window) {
        fprintf(stderr, "swmctl: out of memory\n");
        exit(EXIT_FAILURE);
    }
    window->handle = handle;
    wl_list_insert(windows.prev, &window->link);
    ext_foreign_toplevel_handle_v1_add_listener(handle, &window_listener, window);
}

/* Record that the compositor accepted the list's stop request. */
static void window_list_finished(void *data, struct ext_foreign_toplevel_list_v1 *list) {
    toplevel_list_finished = true;
}

static const struct ext_foreign_toplevel_list_v1_listener window_list_listener = {
    .toplevel = window_announced,
    .finished = window_list_finished,
};

/* Bind swm's workspace and toplevel control globals. */
static void global(
    void *data, struct wl_registry *registry, uint32_t name, const char *interface, uint32_t version
) {
    if (requested_kind == COMMAND_WORKSPACE &&
        !strcmp(interface, swm_workspace_manager_v1_interface.name)) {
        manager = wl_registry_bind(registry, name, &swm_workspace_manager_v1_interface, 1);
        swm_workspace_manager_v1_add_listener(manager, &manager_listener, nullptr);
    } else if (
        requested_kind == COMMAND_WINDOW_LIST &&
        !strcmp(interface, ext_foreign_toplevel_list_v1_interface.name)
    ) {
        toplevel_list =
            wl_registry_bind(registry, name, &ext_foreign_toplevel_list_v1_interface, 1);
        ext_foreign_toplevel_list_v1_add_listener(toplevel_list, &window_list_listener, nullptr);
    } else if (
        requested_kind == COMMAND_WINDOW_ACTIVATE &&
        !strcmp(interface, swm_toplevel_manager_v1_interface.name)
    ) {
        toplevel_control = wl_registry_bind(registry, name, &swm_toplevel_manager_v1_interface, 1);
    }
}

static void global_remove(void *data, struct wl_registry *registry, uint32_t name) {}

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

/* Parse a one-based workspace number. */
static uint32_t parse_workspace(const char *value) {
    char         *end;
    unsigned long number;

    errno  = 0;
    number = strtoul(value, &end, 10);

    if (errno || *end || !number || number > MAX_WORKSPACES)
        usage();

    return (uint32_t)number;
}

/* Parse a color in #RRGGBB or #RRGGBBAA notation. */
static uint32_t parse_color(const char *value) {
    char         *end;
    unsigned long color;
    size_t        length = strlen(value);

    if ((length != 7 && length != 9) || value[0] != '#')
        usage();
    errno = 0;
    color = strtoul(value + 1, &end, 16);
    if (errno || *end || color > 0xffffffff)
        usage();

    return length == 7 ? (uint32_t)(color << 8) | 0xff : (uint32_t)color;
}

/* Print one field from a cached workspace. */
static void print_workspace_field(uint32_t number, const char *field) {
    workspace_t *ws = &workspaces[number - 1];

    if (!strcmp(field, "title")) {
        puts(ws->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_TITLE ? ws->title : "");
    } else if (!strcmp(field, "color")) {
        if (ws->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_COLOR)
            print_color(ws->color);
        putchar('\n');
    } else {
        puts(ws->flags & SWM_WORKSPACE_MANAGER_V1_METADATA_SELECTED ? "true" : "false");
    }
}

/* Print every cached workspace number. */
static void print_workspaces(void) {
    for (uint32_t i = 0; i < MAX_WORKSPACES; i++) {
        if (workspaces[i].seen)
            printf("%" PRIu32 "\n", i + 1);
    }
}

/* Write a field with line and column separators escaped. */
static void print_window_field(const char *value) {
    const unsigned char *p;

    for (p = (const unsigned char *)(value ? value : ""); *p; p++) {
        if (*p == '\\')
            fputs("\\\\", stdout);
        else if (*p == '\t')
            fputs("\\t", stdout);
        else if (*p == '\n')
            fputs("\\n", stdout);
        else if (*p == '\r')
            fputs("\\r", stdout);
        else
            putchar(*p);
    }
}

/* Print mapped toplevels as title, application ID, and identifier fields. */
static void print_windows(void) {
    window_t *window;

    wl_list_for_each(window, &windows, link) {
        if (window->closed || !window->identifier[0])
            continue;
        print_window_field(window->title);
        putchar('\t');
        print_window_field(window->app_id);
        putchar('\t');
        puts(window->identifier);
    }
}

/* Release cached foreign-toplevel state. */
static void destroy_windows(void) {
    window_t *window, *tmp;

    wl_list_for_each_safe(window, tmp, &windows, link) {
        if (window->handle)
            ext_foreign_toplevel_handle_v1_destroy(window->handle);
        wl_list_remove(&window->link);
        free(window->title);
        free(window->app_id);
        free(window);
    }
}

/* Stop the list, then release its child handles before the list object. */
static bool destroy_window_list(void) {
    if (!toplevel_list_finished) {
        ext_foreign_toplevel_list_v1_stop(toplevel_list);
        if (wl_display_roundtrip(display) < 0 || !toplevel_list_finished)
            return false;
    }
    destroy_windows();
    ext_foreign_toplevel_list_v1_destroy(toplevel_list);
    toplevel_list = nullptr;
    return wl_display_flush(display) >= 0;
}

/* Connect to swm and collect the requested initial state. */
static void connect_manager(enum command_kind kind) {
    struct wl_registry *registry;

    if (!(display = wl_display_connect(nullptr))) {
        fprintf(stderr, "swmctl: cannot connect to Wayland display\n");
        exit(EXIT_FAILURE);
    }
    wl_list_init(&windows);
    requested_kind = kind;
    registry       = wl_display_get_registry(display);
    wl_registry_add_listener(registry, &registry_listener, nullptr);

    if (wl_display_roundtrip(display) < 0 || (kind == COMMAND_WORKSPACE && !manager) ||
        (kind == COMMAND_WINDOW_LIST && !toplevel_list) ||
        (kind == COMMAND_WINDOW_ACTIVATE && !toplevel_control)) {
        fprintf(
            stderr,
            kind == COMMAND_WORKSPACE ? "swmctl: compositor does not support swm-workspace-v1\n"
            : kind == COMMAND_WINDOW_LIST
                ? "swmctl: compositor does not support ext-foreign-toplevel-list-v1\n"
                : "swmctl: compositor does not support swm-toplevel-v1\n"
        );
        exit(EXIT_FAILURE);
    }
    if (kind != COMMAND_WINDOW_ACTIVATE && wl_display_roundtrip(display) < 0) {
        fprintf(
            stderr,
            kind == COMMAND_WORKSPACE ? "swmctl: failed to read workspace metadata\n"
                                      : "swmctl: failed to read toplevels\n"
        );
        exit(EXIT_FAILURE);
    }
    wl_registry_destroy(registry);
}

int main(int argc, char **argv) {
    const char *identifier = nullptr;
    bool        window_command;
    const char *command;
    const char *field     = nullptr;
    const char *value     = nullptr;
    uint32_t    workspace = 0;

    if (argc < 3)
        usage();
    window_command = !strcmp(argv[1], "window");
    command        = argv[2];
    if (window_command) {
        if (!strcmp(command, "list")) {
            if (argc != 3)
                usage();
        } else if (!strcmp(command, "activate")) {
            if (argc != 4)
                usage();
            identifier = argv[3];
        } else {
            usage();
        }
    } else if (strcmp(argv[1], "workspace")) {
        usage();
    }
    subscribing = !window_command && !strcmp(command, "subscribe");

    if (window_command) {
        ;
    } else if (!strcmp(command, "list")) {
        if (argc != 3)
            usage();
    } else if (!strcmp(command, "get")) {
        if (argc != 5)
            usage();
        workspace = parse_workspace(argv[3]);
        field     = argv[4];
        if (strcmp(field, "title") && strcmp(field, "color") && strcmp(field, "selected"))
            usage();
    } else if (!strcmp(command, "set")) {
        if (argc != 6)
            usage();
        workspace = parse_workspace(argv[3]);
        field     = argv[4];
        value     = argv[5];
        if (!strcmp(field, "title")) {
            if (strlen(value) >= MAX_WORKSPACE_TITLE)
                usage();
        } else if (!strcmp(field, "color")) {
            parse_color(value);
        } else {
            usage();
        }
    } else if (!strcmp(command, "clear")) {
        if (argc != 5)
            usage();
        workspace = parse_workspace(argv[3]);
        field     = argv[4];
        if (strcmp(field, "title") && strcmp(field, "color"))
            usage();
    } else if (subscribing) {
        if (argc == 4 && !strcmp(argv[3], "--format=waybar"))
            waybar = true;
        else if (argc != 3)
            usage();
    } else {
        usage();
    }
    connect_manager(
        !window_command            ? COMMAND_WORKSPACE
        : !strcmp(command, "list") ? COMMAND_WINDOW_LIST
                                   : COMMAND_WINDOW_ACTIVATE
    );

    if (window_command) {
        if (!strcmp(command, "list"))
            print_windows();
        else
            swm_toplevel_manager_v1_activate(toplevel_control, identifier);
    } else if (subscribing) {
        while (wl_display_dispatch(display) >= 0)
            ;
        return EXIT_FAILURE;
    }
    if (!window_command && !strcmp(command, "list")) {
        print_workspaces();
    } else if (!window_command) {
        if (!workspaces[workspace - 1].seen) {
            fprintf(stderr, "swmctl: workspace %" PRIu32 " does not exist\n", workspace);
            return EXIT_FAILURE;
        }
        if (!strcmp(command, "get")) {
            print_workspace_field(workspace, field);
        } else if (!strcmp(command, "set")) {
            if (!strcmp(field, "title"))
                swm_workspace_manager_v1_set_title(manager, workspace, value);
            else
                swm_workspace_manager_v1_set_color(manager, workspace, parse_color(value));
        } else {
            if (!strcmp(field, "title"))
                swm_workspace_manager_v1_set_title(manager, workspace, "");
            else
                swm_workspace_manager_v1_set_color(manager, workspace, 0);
        }
    }
    if (wl_display_flush(display) < 0 || wl_display_roundtrip(display) < 0) {
        fprintf(stderr, "swmctl: command failed\n");
        return EXIT_FAILURE;
    }
    if (requested_kind == COMMAND_WINDOW_LIST && !destroy_window_list()) {
        fprintf(stderr, "swmctl: failed to release toplevel list\n");
        return EXIT_FAILURE;
    }
    if (toplevel_control)
        swm_toplevel_manager_v1_destroy(toplevel_control);
    if (manager)
        swm_workspace_manager_v1_destroy(manager);
    wl_display_disconnect(display);

    return EXIT_SUCCESS;
}