#include "cairo.h"
#include "rect.h"
#include "timeutil.h"
#include "window.h"
#include "node.h"
#include "defs.h"
#include <xkbcommon/xkbcommon.h>
#ifdef _UI_DEBUG
#define PROF
#endif
#include "prof.c"
#include "time.h"
#include <stdint.h>
#include <stdlib.h>
#include <sys/poll.h>
#include <time.h>
#include <xcb/xcb.h>
#include <xcb/xproto.h>
#include <xcb/xkb.h>
#include <xkbcommon/xkbcommon-x11.h>

static void window_process_xcb_event(UIWindow *window, xcb_generic_event_t *event);
static int window_select_xkb_events_for_device(UIWindow *window, int32_t device_id);
static int window_keyboard_update_keymap(UIWindow *window, UIWindowXKBKeyboard *kb);
static int window_keyboard_init(UIWindow *window, UIWindowXKBKeyboard *kb, uint8_t first_xkb_event, int32_t device_id, struct xkb_context *ctx);
static int window_init_core_xkb_keyboard(UIWindow *window);
static void window_keyboard_process_event(UIWindow *window, UIWindowXKBKeyboard *kb, xcb_generic_event_t *gevent);
static void window_keyboard_deinit(UIWindowXKBKeyboard *kb);

int window_invalidate_node(UIWindow *window, UINode *node)
{
    if (node->flags & UI_NODE_COMPONENT) {
        if (!node->parent) {
            return -1;
        }
        node = node->parent;
    }
    UIRect local_rect = (UIRect){
        .x = node->rect.x + node->window_rel_x,
        .y = node->rect.y + node->window_rel_y,
        .w = node->rect.w,
        .h = node->rect.h
    };
    window->invalid_region = ui_rect_united(&window->invalid_region, &local_rect);
    return 0;
}

void window_free(UIWindow *window)
{
    if (!window) return;

    if (window->_cairo_surface) cairo_surface_destroy(window->_cairo_surface);
    if (window->drw) cairo_destroy(window->drw);
    if (window->_xcb_connection) xcb_disconnect(window->_xcb_connection);
    window_keyboard_deinit(window->xkb_core_keyboard);

    node_free(window->root);
    free(window);
}

UIWindow *window_new(int width, int height)
{
    bool fail = false;
    UIWindow *window = NULL;

    window = calloc(1, sizeof(UIWindow));
    if (!window) {
        fail = true;
        goto done;
    }
    window->hovered = NULL;
    window->pressed = NULL;
    window->root = NULL;
    window->invalid_region = (UIRect){ 0, 0, 0, 0 };
    window->_cairo_surface = NULL;
    window->_xcb_connection = NULL;
    window->_xcb_window_id = -1;
    window->drw = NULL;

    window->_xcb_connection = xcb_connect(NULL, NULL);
    if (xcb_connection_has_error(window->_xcb_connection) != 0) {
        fail = true;
        goto done;
    }

    const xcb_setup_t *setup = xcb_get_setup(window->_xcb_connection);
    xcb_screen_iterator_t iter = xcb_setup_roots_iterator(setup);
    xcb_screen_t *screen = iter.data;

    xcb_visualtype_t *visualtype = NULL;

    for (xcb_depth_iterator_t iter_depth = xcb_screen_allowed_depths_iterator(screen); iter_depth.rem; xcb_depth_next(&iter_depth)) {
        for (xcb_visualtype_iterator_t iter_visualtype = xcb_depth_visuals_iterator(iter_depth.data); iter_visualtype.rem; xcb_visualtype_next(&iter_visualtype)) {
            if (iter_visualtype.data->visual_id == screen->root_visual) {
                visualtype = iter_visualtype.data;
                break;
            }
        }
        if (visualtype) break;
    }

    if (!visualtype) {
        fail = true;
        goto done;
    }

    window->_xcb_window_id = xcb_generate_id(window->_xcb_connection);
    if (window->_xcb_window_id == (unsigned)-1) {
        fail = true;
        goto done;
    }
    xcb_create_window(
        window->_xcb_connection, // connection
        XCB_COPY_FROM_PARENT, // depth
        window->_xcb_window_id, // window id
        screen->root, // parent window
        0, // x
        0, // y
        width, // width
        height, // height
        0, // border_width
        XCB_WINDOW_CLASS_INPUT_OUTPUT, // class
        screen->root_visual, // visual
        XCB_CW_EVENT_MASK, // value mask
        (uint32_t[]){ XCB_EVENT_MASK_EXPOSURE | XCB_EVENT_MASK_STRUCTURE_NOTIFY | XCB_EVENT_MASK_POINTER_MOTION | XCB_EVENT_MASK_BUTTON_PRESS | XCB_EVENT_MASK_BUTTON_RELEASE | XCB_EVENT_MASK_KEY_PRESS } // value list
    );
    xcb_map_window(window->_xcb_connection, window->_xcb_window_id);
    if (xcb_flush(window->_xcb_connection) <= 0) {
        fail = true;
        goto done;
    }

    window->_cairo_surface = cairo_xcb_surface_create(window->_xcb_connection, window->_xcb_window_id, visualtype, width, height);
    if (cairo_surface_status(window->_cairo_surface) != CAIRO_STATUS_SUCCESS) {
        fail = true;
        goto done;
    }
    window->drw = cairo_create(window->_cairo_surface);
    if (cairo_status(window->drw) != CAIRO_STATUS_SUCCESS) {
        fail = true;
        goto done;
    }

    if (window_init_core_xkb_keyboard(window) < 0) {
        fprintf(stderr, "err: window_new: failed to initialize keyboard (_window_init_core_xkb_keyboard failed). Keyboard input may not be available.\n");
    }

done:
    if (fail) {
        window_free(window);
        return NULL;
    }
    return window;
}

bool window_flush_invalidated(UIWindow *window)
{
    if (ui_rect_null(&window->invalid_region)) {
        return false;
    }
    begin_clock("Flush invalidated widgets");

    cairo_push_group(window->drw);
    node_repaint(window->root, &window->invalid_region, true, 0, 0);
    window->invalid_region = (UIRect){ 0, 0, 0, 0 };
    cairo_pop_group_to_source(window->drw);
    cairo_paint(window->drw);
    cairo_surface_flush(window->_cairo_surface);

    end_clock();
    return true;
}

int window_attach_root(UIWindow *window, UINode *root)
{
    if (!window || !root || window->root || !window->drw) {
        return -1;
    }

    root->window = window;
    root->drw = window->drw;
    root->width_policy = UI_SIZE_POLICY_STATIC;
    root->height_policy = UI_SIZE_POLICY_STATIC;
    window->root = root;
    return 0;
}

// Thanks: https://github.com/xkbcommon/libxkbcommon/blob/003fdee1378382d4fef77089f5f4652cd2422c6a/tools/interactive-x11.c#L67
static int window_select_xkb_events_for_device(UIWindow *window, int32_t device_id)
{
    if (!window->_xcb_connection) {
        return -1;
    }

    enum {
        required_events =
            (XCB_XKB_EVENT_TYPE_NEW_KEYBOARD_NOTIFY |
                XCB_XKB_EVENT_TYPE_MAP_NOTIFY |
                XCB_XKB_EVENT_TYPE_STATE_NOTIFY),

        required_nkn_details =
            (XCB_XKB_NKN_DETAIL_KEYCODES),

        required_map_parts =
            (XCB_XKB_MAP_PART_KEY_TYPES |
                XCB_XKB_MAP_PART_KEY_SYMS |
                XCB_XKB_MAP_PART_MODIFIER_MAP |
                XCB_XKB_MAP_PART_EXPLICIT_COMPONENTS |
                XCB_XKB_MAP_PART_KEY_ACTIONS |
                XCB_XKB_MAP_PART_VIRTUAL_MODS |
                XCB_XKB_MAP_PART_VIRTUAL_MOD_MAP),

        required_state_details =
            (XCB_XKB_STATE_PART_MODIFIER_BASE |
                XCB_XKB_STATE_PART_MODIFIER_LATCH |
                XCB_XKB_STATE_PART_MODIFIER_LOCK |
                XCB_XKB_STATE_PART_GROUP_BASE |
                XCB_XKB_STATE_PART_GROUP_LATCH |
                XCB_XKB_STATE_PART_GROUP_LOCK),
    };

    static const xcb_xkb_select_events_details_t details = {
        .affectNewKeyboard = required_nkn_details,
        .newKeyboardDetails = required_nkn_details,
        .affectState = required_state_details,
        .stateDetails = required_state_details,
    };

    xcb_void_cookie_t cookie = xcb_xkb_select_events_aux_checked(
        window->_xcb_connection,
        device_id,
        required_events,    /* affectWhich */
        0,                  /* clear */
        0,                  /* selectAll */
        required_map_parts, /* affectMap */
        required_map_parts, /* map */
        &details            /* details */
    );

    xcb_generic_error_t *error = xcb_request_check(window->_xcb_connection, cookie);
    if (error) {
        free(error);
        return -1;
    }

    return 0;
}

static int window_keyboard_update_keymap(UIWindow *window, UIWindowXKBKeyboard *kb)
{
    struct xkb_keymap *keymap = xkb_x11_keymap_new_from_device(kb->ctx, window->_xcb_connection, kb->device_id, XKB_KEYMAP_COMPILE_NO_FLAGS);
    if (!keymap) {
        return -1;
    }

    struct xkb_state *state = xkb_x11_state_new_from_device(keymap, window->_xcb_connection, kb->device_id);
    if (!state) {
        xkb_keymap_unref(keymap);
        return -1;
    }

    xkb_state_unref(kb->state);
    xkb_keymap_unref(kb->keymap);
    kb->state = state;
    kb->keymap = keymap;

    return 0;
}

static int window_keyboard_init(UIWindow *window, UIWindowXKBKeyboard *kb, uint8_t first_xkb_event, int32_t device_id, struct xkb_context *ctx)
{
    kb->first_xkb_event = first_xkb_event;
    kb->device_id = device_id;
    kb->ctx = ctx;
    kb->keymap = NULL;
    kb->state = NULL;

    if (window_keyboard_update_keymap(window, kb) < 0) {
        return -1;
    }

    if (window_select_xkb_events_for_device(window, device_id) < 0) {
        xkb_state_unref(kb->state);
        xkb_keymap_unref(kb->keymap);
        return -1;
    }

    return 0;
}

static int window_init_core_xkb_keyboard(UIWindow *window)
{
    uint8_t first_xkb_event;
    int ret = 0;
    UIWindowXKBKeyboard *kb = NULL;
    struct xkb_context *ctx = NULL;

    if (!xkb_x11_setup_xkb_extension(
        window->_xcb_connection,
        XKB_X11_MIN_MAJOR_XKB_VERSION,
        XKB_X11_MIN_MINOR_XKB_VERSION,
        XKB_X11_SETUP_XKB_EXTENSION_NO_FLAGS,
        NULL, NULL, &first_xkb_event, NULL
    )) {
        fprintf(stderr, "err: _window_init_core_xkb_keyboard: failed to set up xkb extension (xkb_x11_setup_xkb_extension)\n");
        return -1;
    }

    kb = malloc(sizeof(UIWindowXKBKeyboard));
    if (!kb) {
        fprintf(stderr, "err: _window_init_core_xkb_keyboard: failed to malloc keyboard struct\n");
        ret = -1;
        goto defer_error;
    }

    ctx = xkb_context_new(XKB_CONTEXT_NO_FLAGS);
    if (!ctx) {
        fprintf(stderr, "err: _window_init_core_xkb_keyboard: failed to create xkb context\n");
        ret = -1;
        goto defer_error;
    }

    int32_t core_kbd_device_id = xkb_x11_get_core_keyboard_device_id(window->_xcb_connection);
    if (core_kbd_device_id < 0) {
        fprintf(stderr, "err: _window_init_core_xkb_keyboard: failed to find xkb core keyboard device_id\n");
        ret = -1;
        goto defer_error;
    }
    
    if (window_keyboard_init(window, kb, first_xkb_event, core_kbd_device_id, ctx) < 0) {
        fprintf(stderr, "err: _window_init_core_xkb_keyboard: failed to initialize keyboard (_window_keyboard_init)\n");
        ret = -1;
        goto defer_error;
    }

    goto success;

defer_error:
    free(kb);
    xkb_context_unref(ctx);
    return ret;
success:
    window->xkb_core_keyboard = kb;
    return ret;
}

static void window_keyboard_deinit(UIWindowXKBKeyboard *kb)
{
    if (!kb) {
        return;
    }

    xkb_state_unref(kb->state);
    xkb_keymap_unref(kb->keymap);
    free(kb);
}

static void window_keyboard_process_event(UIWindow *window, UIWindowXKBKeyboard *kb, xcb_generic_event_t *gevent)
{
    union xkb_event {
        struct {
            uint8_t response_type;
            uint8_t xkbType;
            uint16_t sequence;
            xcb_timestamp_t time;
            uint8_t deviceID;
        } any;
        xcb_xkb_new_keyboard_notify_event_t new_keyboard_notify;
        xcb_xkb_map_notify_event_t map_notify;
        xcb_xkb_state_notify_event_t state_notify;
    } *event = (union xkb_event *) gevent;

    if (event->any.deviceID != kb->device_id) {
        return;
    }

    switch (event->any.xkbType) {
        case XCB_XKB_NEW_KEYBOARD_NOTIFY: {
            if (event->new_keyboard_notify.changed && XCB_XKB_NKN_DETAIL_KEYCODES) {
                window_keyboard_update_keymap(window, kb);
            }
            break;
        }
        case XCB_XKB_MAP_NOTIFY: {
            window_keyboard_update_keymap(window, kb);
            break;
        }
        case XCB_XKB_STATE_NOTIFY: {
            xkb_state_update_mask(
                kb->state,
                event->state_notify.baseMods,
                event->state_notify.latchedMods,
                event->state_notify.lockedMods,
                event->state_notify.baseGroup,
                event->state_notify.latchedGroup,
                event->state_notify.lockedGroup
            );
            break;
        }
    }
}

static void window_process_xcb_event(UIWindow *window, xcb_generic_event_t *event)
{
    begin_clock("Process XCB event");
    uint8_t event_type = event->response_type & ~0x80;
    switch (event_type) {
        case XCB_EXPOSE: {
            begin_clock("XCB_EXPOSE");
            xcb_expose_event_t *ev = (xcb_expose_event_t*)event;
            UIRect region = { ev->x, ev->y, ev->width, ev->height };
            window->invalid_region = ui_rect_united(&window->invalid_region, &region);
            window_invalidate_node(window, window->root);
            end_clock();
            break;
        }

        case XCB_CONFIGURE_NOTIFY: {
            begin_clock("XCB_CONFIGURE_NOTIFY");
            xcb_configure_notify_event_t *ev = (xcb_configure_notify_event_t*)event;
            UIRect new_rect = { 0, 0, ev->width, ev->height };
            if (!ui_rect_equals(&window->root->rect, &new_rect)) {
                cairo_xcb_surface_set_size(window->_cairo_surface, ev->width, ev->height);
                window->root->rect.x = 0;
                window->root->rect.y = 0;
                window->root->rect.w = ev->width;
                window->root->rect.h = ev->height;
                node_dispatch(window->root, UI_EVENT_RELAYOUT, 0, NULL);
#ifdef _UI_DEBUG
                node_dump(window->root, 0);
#endif
            }
            end_clock();
            break;
        }

        case XCB_MOTION_NOTIFY: {
            begin_clock("XCB_MOTION_NOTIFY");
            xcb_motion_notify_event_t *ev = (xcb_motion_notify_event_t*)event;
            UINode *node = node_by_point(window->root, ev->event_x, ev->event_y);
            if (node && node != window->hovered) {
                UINode *previously_hovered = window->hovered;
                window->hovered = node;
                if (previously_hovered) {
                    node_dispatch(previously_hovered, UI_EVENT_UNHOVERED, 0, NULL);
                }
                node_dispatch(node, UI_EVENT_HOVERED, 0, NULL);
            }
            end_clock();
            break;
        }

        case XCB_KEY_PRESS: {
            begin_clock("XCB_KEY_PRESS");
            xcb_key_press_event_t *ev = (xcb_key_press_event_t*)event;

            if (window->xkb_core_keyboard) {
                UINode *n = window->pressed;
                while (n) {
                    if (node_dispatch(n, UI_EVENT_KEY_DOWN, ev->detail, window->xkb_core_keyboard->state) > 0) {
                        break;
                    }
                    n = n->parent;
                }
            }

            end_clock();
            break;
        }

        case XCB_BUTTON_RELEASE: /* through  */
        case XCB_BUTTON_PRESS: {
            begin_clock("XCB_BUTTON_PRESS");
            xcb_button_press_event_t *ev = (xcb_button_press_event_t*)event;
            UINode *node = node_by_point(window->root, ev->event_x, ev->event_y);
            if (node) {
                enum UIEvent ui_event = UI_EVENT_BUTTON_LEFT_UPDATE;
                bool state = event_type == XCB_BUTTON_PRESS ? true : false;
                double delta = 0.0;
                switch (ev->detail) {
                    case XCB_BUTTON_INDEX_3: {
                        ui_event = UI_EVENT_BUTTON_RIGHT_UPDATE;
                        break;
                    }
                    case XCB_BUTTON_INDEX_4: {
                        // scroll up
                        ui_event = UI_EVENT_SCROLL;
                        delta = 20.0;
                        break;
                    }
                    case XCB_BUTTON_INDEX_5: {
                        // scroll down
                        ui_event = UI_EVENT_SCROLL;
                        delta = -20.0;
                        break;
                    }
                }

                UINode *n = node;
                while (n) {
                    if (node_dispatch(n, ui_event, state, &delta) > 0) {
                        break;
                    }
                    n = n->parent;
                }

                if (ui_event == UI_EVENT_BUTTON_LEFT_UPDATE) {
                    if (state) {
                        UINode *previously_pressed = window->pressed;
                        window->pressed = node;
                        if (previously_pressed) {
                            node_dispatch(previously_pressed, UI_EVENT_UNPRESSED, 0, NULL);
                        }
                        node_dispatch(node, UI_EVENT_PRESSED, 0, NULL);
                    } else {
                        UINode *previously_pressed = window->pressed;
                        if (previously_pressed && !(previously_pressed->flags & UI_NODE_RETAINS_PRESS)) {
                            window->pressed = NULL;
                            node_dispatch(previously_pressed, UI_EVENT_UNPRESSED, 0, NULL);
                        }
                    }
                }
            }
            end_clock();
            break; 
        }
        
        default: {
            if (window->xkb_core_keyboard && event->response_type == window->xkb_core_keyboard->first_xkb_event) {
                window_keyboard_process_event(window, window->xkb_core_keyboard, event);
            }
            break;
        }
    }
    end_clock();
}

int window_turn(UIWindow *window)
{
    if (!window || !window->root || !window->_xcb_connection || !window->_cairo_surface || !window->drw) {
        return -1;
    }
    if (xcb_connection_has_error(window->_xcb_connection) != 0) {
        return -1;
    }
    if (cairo_surface_status(window->_cairo_surface) != CAIRO_STATUS_SUCCESS || cairo_status(window->drw) != CAIRO_STATUS_SUCCESS) {
        return -1;
    }

    bool has_looped_once = false;

    struct pollfd fds[UI_WINDOW_MAX_POLL] = {0};
    for (int i = 0; i < UI_WINDOW_MAX_POLL; i++) {
        fds[i].fd = -1;
        fds[i].events = 0;
        fds[i].revents = 0;
    }

    fds[0].fd = xcb_get_file_descriptor(window->_xcb_connection);
    fds[0].events = POLLIN;

    int64_t frame_peak_ms = 0;
    int64_t frame_peak_last_measurement_ms = 0;

    for (;;) {
        int64_t poll_timeout;
        /* compute `poll_timeout` based on active timers */
        {
            int64_t now = time_current_ms();
            int64_t lowest = INT64_MAX;
            bool has_timer = false;
            for (int i = 0; i < UI_WINDOW_MAX_TIMERS; i++) {
                UIWindowTimer *timer = &window->timers[i];
                if (timer->present) {
                    has_timer = true;
                    int64_t distance = (timer->started_at_ms + timer->duration_ms) - now;
                    if (distance < lowest) {
                        lowest = distance;
                    }
                }
            }
            if (lowest < 0) {
                lowest = 0;
            }

            poll_timeout = has_timer ? lowest : -1;
            if (!has_looped_once) {
                // Makes sure that we handle any pending events before entering the loop for the first time
                poll_timeout = 0;
            }
        }

        {
            if (poll(fds, UI_WINDOW_MAX_POLL, poll_timeout) < 0) {
                fprintf(stderr, "err: window_turn(): poll() failed\n");
                return -1;
            }
        }

        clear_summary();

        begin_clock("Frame");
        int64_t frame_start_ms = time_current_ms();
        for (int i = 0; i < UI_WINDOW_MAX_POLL; i++) {
            if (fds[i].revents & POLLNVAL) {
                fprintf(stderr, "err: window_turn(): poll got POLLNVAL\n");
                return -1;
            }
            if (fds[i].revents & POLLERR) {
                fprintf(stderr, "err: window_turn(): poll got POLLERR\n");
                return -1;
            }
            if (fds[i].revents & POLLHUP) {
                fprintf(stderr, "err: window_turn(): poll got POLLHUP\n");
                return -1;
            }

            if ((i == 0 && fds[i].revents & POLLIN) || !has_looped_once) {
                xcb_generic_event_t *xcb_event;
                while ((xcb_event = xcb_poll_for_event(window->_xcb_connection))) {
                    window_process_xcb_event(window, xcb_event);
                    free(xcb_event);
                }
            }
        }

        /* process finished timers */
        {
            for (int j = 0; j < UI_WINDOW_MAX_TIMERS; j++) {
                UIWindowTimer *timer = &window->timers[j];
                if (timer->present) {
                    int64_t now = time_current_ms();
                    int64_t distance = (timer->started_at_ms + timer->duration_ms) - now;
                    if (distance <= 0) {
                        int64_t dt = now - timer->started_at_ms;
#ifdef _UI_DEBUG
                        printf("timer jitter: %ld\n", dt - timer->duration_ms);
#endif
                        UINode *target = timer->target;
                        timer->present = false;
                        timer->started_at_ms = 0;
                        timer->duration_ms = 0;
                        timer->target = NULL;
                        if (target) {
                            node_dispatch(target, UI_EVENT_TIMER_END, 0, &dt);
                        }
                    }
                }
            }
        }

        if (window_flush_invalidated(window)) {
            begin_clock("Flush painting results to XCB");
            xcb_flush(window->_xcb_connection);
            end_clock();
        }

        has_looped_once = true;

#ifdef _UI_DEBUG
        int64_t frame_end_ms = time_current_ms();
        int64_t frame_delta_ms = frame_end_ms - frame_start_ms;
        if (frame_delta_ms > frame_peak_ms || frame_end_ms - frame_peak_last_measurement_ms >= 3000) {
            frame_peak_last_measurement_ms = frame_end_ms;
            frame_peak_ms = frame_delta_ms;
            printf("peak frametime: %ldms\n", frame_peak_ms);
        }
#else
        (void)frame_peak_last_measurement_ms;
        (void)frame_start_ms;
        (void)frame_peak_ms;
#endif

        end_clock();

        dump_summary(stdout);
    }
    return 0;
}

UIWindowTimer *window_sched_timer(UIWindow *window, UINode *target, int64_t duration_ms)
{
    for (int i = 0; i < UI_WINDOW_MAX_TIMERS; i++) {
        if (!window->timers[i].present) {
            UIWindowTimer *timer = &window->timers[i];
            timer->present = true;
            timer->duration_ms = duration_ms;
            timer->target = target;
            timer->started_at_ms = time_current_ms();
            return timer;
        }
    }
    return NULL;
}