#include "screen.h"
#undef Screen

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>
#include <signal.h>

#define XLIB_ILLEGAL_ACCESS
#include <X11/X.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>

static Display* dpy;
static int scr;
static Window root;

static int count = 0;
static size_t bit_depth = sizeof(uint32_t) * 8;
static int key_count;

static bool* key_state;
static bool* key_checked;

typedef struct {
    Window window;
    GC graphics_ctx;
    Pixmap pixel_map;
    XImage** images;
    pthread_t update_thread;
    pthread_mutex_t recreate_lock;
} WindowState;

static void init_x() {
    if ((dpy = XOpenDisplay(NULL)) == NULL) {
        printf("error: failed to open x11 display\n");
        exit(EXIT_FAILURE);
    }

    scr = DefaultScreen(dpy);
    root = RootWindow(dpy, scr);

    key_count = dpy->max_keycode - dpy->min_keycode;
    key_state = malloc(sizeof(bool) * key_count);
    key_checked = malloc(sizeof(bool) * key_count);

    memset(key_state, 0, sizeof(bool) * key_count);
    memset(key_checked, 0, sizeof(bool) * key_count);
}

static Window create_window(int x, int y, int w, int h, int d, const char* t) {
    
    XVisualInfo vis_info;
    if(!XMatchVisualInfo(dpy, scr, d, TrueColor, &vis_info)) {
        printf("error: %d depth not supported\n", d);
        exit(EXIT_FAILURE);    
    }

    Visual* visual = vis_info.visual;
    
    XSetWindowAttributes xwa;
    xwa.background_pixel = 0;
    xwa.border_pixel = 0;
    xwa.event_mask = KeyPressMask | KeyReleaseMask |  StructureNotifyMask;
    xwa.colormap = XCreateColormap(dpy, root, visual, AllocNone);

    long wm = CWBackPixel | CWColormap | CWBorderPixel | CWEventMask;

    Window window = XCreateWindow(dpy, root, x, y, w, h, 0, d, InputOutput, visual, wm, &xwa);
    XStoreName(dpy, window, t);
    XSetIconName(dpy, window, t);

    return window;
}

static XImage* create_image(int width, int height, void* data) {
    return XCreateImage(
        dpy, 
        CopyFromParent, 
        bit_depth, 
        ZPixmap, 
        0, 
        data,
        width, 
        height, 
        bit_depth, 
        0
    );
}

static Swapchain create_swapchain(int width, int height, int frames, int bit_depth, WindowState* state) {
    Swapchain swapchain;
    swapchain.width = width;
    swapchain.height = height;
    swapchain.image_count = frames;
    swapchain.image_front = 0;
    swapchain.image_recent = 0;
    swapchain.image_current = 0;
    swapchain.images = malloc(sizeof(uint32_t*) * frames);
    
    state->images = malloc(sizeof(XImage*) * frames);
    state->pixel_map = XCreatePixmap(dpy, state->window, width, height, bit_depth);
    state->graphics_ctx = XCreateGC(dpy, state->pixel_map, 0, NULL);
    
    size_t pixel_count = width * height;
    for (int i = 0; i < frames; i++) {
        uint32_t* data = malloc(pixel_count * bit_depth / 8);
        memset(data, 0, pixel_count * bit_depth / 8);

        swapchain.images[i] = data;
        state->images[i] = create_image(width, height, data);

    }

    pthread_mutex_init(&state->recreate_lock, NULL);
    pthread_mutex_init(&swapchain.lock, NULL);

    return swapchain;
}

static void free_swapchain(Swapchain* swapchain, WindowState* state) {
    pthread_mutex_lock(&state->recreate_lock);
    for (uint8_t i = 0; i < swapchain->image_count; i++) {
        XDestroyImage(state->images[i]); // also frees pixel buffer :3
    }
    
    free(swapchain->images);
    free(state->images);
    swapchain->images = NULL;
    
    XFreePixmap(dpy, state->pixel_map);
    XFreeGC(dpy, state->graphics_ctx);

    pthread_mutex_unlock(&state->recreate_lock);
}

static void recreate_swapchain(int width, int height, int bit_depth, Swapchain* swapchain, WindowState* state) {
    pthread_mutex_lock(&state->recreate_lock);

    swapchain->width = width;
    swapchain->height = height;
    
    XFreePixmap(dpy, state->pixel_map);
    XFreeGC(dpy, state->graphics_ctx);

    state->pixel_map = XCreatePixmap(dpy, state->window, width, height, bit_depth);
    state->graphics_ctx = XCreateGC(dpy, state->pixel_map, 0, NULL);
   
    size_t pixel_count = width * height;
    for (uint8_t i = 0; i < swapchain->image_count; i++) {
        XDestroyImage(state->images[i]); // also frees pixel buffer :3
        
        uint32_t* data = malloc(pixel_count * bit_depth / 8);
        memset(data, 0, pixel_count * bit_depth / 8);

        swapchain->images[i] = data;
        state->images[i] = create_image(width, height, data);
    }

    pthread_mutex_unlock(&state->recreate_lock);
}

void swapchain_next(Swapchain* swapchain) {
    pthread_mutex_lock(&swapchain->lock);
    uint8_t next = swapchain->image_current + 1;
    next %= swapchain->image_count;

    if (next == swapchain->image_front) {
        next += 1;
        next %= 3;
    }

    swapchain->image_current = next;
    pthread_mutex_unlock(&swapchain->lock);
}

void swapchain_submit(Swapchain* swapchain) {
    pthread_mutex_lock(&swapchain->lock);
    swapchain->image_recent = swapchain->image_current;
    pthread_mutex_unlock(&swapchain->lock);
}

static void* swapchain_thread(void* arg) {
    struct Screen* screen = (struct Screen*) arg;
    WindowState* state = (WindowState*) screen->internal;

    while (1) {

        if (screen->swapchain.images == NULL) {
            return NULL;
        }
        
        pthread_mutex_lock(&state->recreate_lock);

        pthread_mutex_lock(&screen->swapchain.lock);
        int index = screen->swapchain.image_recent;
        screen->swapchain.image_front = index;
        pthread_mutex_unlock(&screen->swapchain.lock);
        
        XPutImage(
            dpy,
            state->window,
            state->graphics_ctx,
            state->images[index],
            0, 0, 0, 0,
            screen->swapchain.width,
            screen->swapchain.height
        );
        XSync(dpy, 0);
        
        pthread_mutex_unlock(&state->recreate_lock);

        usleep(1000 * 1000 / 60);
    }

    return NULL;
}

#define FRAMES 3

void init_screen(struct Screen* screen, uint16_t width, uint16_t height, const char* title) {

    if (dpy == NULL) {
        init_x();
    }
    
    Window window = create_window(0, 0, width, height, bit_depth, title);
    XMapWindow(dpy, window);

    WindowState* state = malloc(sizeof(WindowState));
    state->window = window;

    screen->swapchain = create_swapchain(width, height, FRAMES, bit_depth, state);
    pthread_create(&state->update_thread, NULL, &swapchain_thread, screen);

    screen->delta = 0;
    screen->internal = state;
    count++;
}

static long get_time() {
    struct timespec spec;
    clock_gettime(CLOCK_MONOTONIC, &spec);
    return round(spec.tv_nsec / 1.0e6) + spec.tv_sec * 1000;
}

static long last = 0;

static void update_delta(struct Screen* screen) {
    long now = get_time();
    if (last == 0) {
        screen->delta = 0;
    } else {
        screen->delta = (now - last) / 1000.0;
    }
    last = now;
}

static void handle_event(struct Screen* screen) {
    XEvent event;
    XNextEvent(dpy, &event);
    switch(event.type) {
        case KeyPress: {
            XKeyEvent xkpe = event.xkey;
            uint32_t code = xkpe.keycode - dpy->min_keycode;
            key_state[code] = true;
            key_checked[code] = false;
            break;
        }

        case KeyRelease: {
            XKeyEvent xkre = event.xkey;
            uint32_t code = xkre.keycode - dpy->min_keycode;
            key_state[code] = false;
            break;
        }

        case ConfigureNotify: {
            XConfigureEvent xce = event.xconfigure;
            if (xce.width != screen->swapchain.width || xce.height != screen->swapchain.height)
                recreate_swapchain(xce.width, xce.height, bit_depth, &screen->swapchain, (WindowState*) screen->internal);
            break;
        }

        default: break;
    }
}

bool poll_screen(struct Screen* screen) {
    
    while (XPending(dpy))
        handle_event(screen);

    update_delta(screen);
    return true;
}

void free_screen(struct Screen* screen) {
    WindowState* state = (WindowState*) screen->internal;

    free_swapchain(&screen->swapchain, state);
    pthread_exit(&state->update_thread);
    pthread_mutex_destroy(&screen->swapchain.lock);
    pthread_mutex_destroy(&state->recreate_lock);

    XUnmapWindow(dpy, state->window);
    XDestroyWindow(dpy, state->window);
    free(screen->internal);

    count--;

    if(count == 0) {
        XCloseDisplay(dpy);
        free(key_state);
        free(key_checked);
    }
}

bool key_pressed(int keycode) {
    if (key_checked[keycode]) return false;
    if (!key_state[keycode]) return false;
    key_checked[keycode] = true;
    return true;
}

bool key_down(int keycode) {
    return key_state[keycode];
}