#include "window.h"
#include "registry.h"
#include "xdg-shell-client-protocol.h"
#include <wayland-client.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <stdatomic.h>

// Moving all window specific data into `struct wayland_window`.

static struct wl_callback_listener callback_listener;
static struct xdg_surface_listener surface_listener;
static struct xdg_toplevel_listener top_listener;
static struct xdg_wm_base_listener wm_base_listener;

static atomic_uint_fast64_t shm_counter = 0;

static int32_t alloc_shm(size_t size)
{
    char name[64];
    uint64_t id = atomic_fetch_add(&shm_counter, 1);
    snprintf(name, sizeof(name), "/wayland-shm-%d-%llu", getpid(), (unsigned long long)id);

    int32_t fd = shm_open(name, O_RDWR | O_CREAT, 0600);
    if (fd == -1)
        return fd;
    shm_unlink(name);
    if (ftruncate(fd, size) == -1)
    {
        close(fd);
        return -1;
    }
    return fd;
}

static void destroy_frame_callback(struct wayland_window *win)
{
    if (win->frame_callback)
    {
        wl_callback_destroy(win->frame_callback);
        win->frame_callback = NULL;
    }
}

static void resize_canvas(struct wayland_window *win)
{
    size_t stride = win->width * 4;
    size_t size = stride * win->height;
    int32_t fd = alloc_shm(size);
    if (fd == -1)
        return;

    if (win->buffer)
    {
        wl_buffer_destroy(win->buffer);
        win->buffer = NULL;
    }

    void *map = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (map == MAP_FAILED)
    {
        close(fd);
        return;
    }
    struct wl_shm *shm = registry_get_shm();
    struct wl_shm_pool *pool = wl_shm_create_pool(shm, fd, size);
    win->buffer = wl_shm_pool_create_buffer(pool, 0, win->width, win->height, stride, WL_SHM_FORMAT_ARGB8888);
    wl_shm_pool_destroy(pool);
    close(fd);

    win->data = map;
}

static void resize_new(struct wayland_window *win, int16_t w, int16_t h)
{
    if (!w || !h)
        return;
    if (win->width != w || win->height != h)
    {
        if (win->data)
        {
            munmap(win->data, win->width * win->height * 4);
            win->data = NULL;
        }
        win->width = w;
        win->height = h;
        resize_canvas(win);
    }
}

static void draw(struct wayland_window *win)
{
    size_t stride = win->width * 4;
    size_t size = stride * win->height;

    if (!win->data || !win->buffer)
        return;

    memset(win->data, win->color++, size);

    wl_surface_attach(win->wl_surface, win->buffer, 0, 0);
    wl_surface_damage_buffer(win->wl_surface, 0, 0, win->width, win->height);
    wl_surface_commit(win->wl_surface);
}

static void xdg_surface_conf(void *data, struct xdg_surface *xdg_surface_local, uint32_t serial)
{
    xdg_surface_ack_configure(xdg_surface_local, serial);
    struct wayland_window *win = data;
    if (!win->data)
        resize_canvas(win);

    draw(win);
}

static void xdg_toplevel_conf(void *data, struct xdg_toplevel *xdg_top, int32_t w, int32_t h, struct wl_array *states)
{
    (void)states;
    struct wayland_window *win = data;
    resize_new(win, w, h);
}

static void xdg_toplevel_cls(void *data, struct xdg_toplevel *xdg_top)
{
    struct wayland_window *win = data;
    (void)xdg_top;
    win->need_close = 1;
}

static void xdg_toplevel_bounds(void *data, struct xdg_toplevel *xdg_top, int32_t w, int32_t h)
{
    (void)data; (void)xdg_top; (void)w; (void)h;
}

static void xdg_toplevel_wm_caps(void *data, struct xdg_toplevel *xdg_top, struct wl_array *caps)
{
    (void)data; (void)xdg_top; (void)caps;
}

static void setup_frame_callback(struct wayland_window *win)
{
    struct wl_callback *cb = wl_surface_frame(win->wl_surface);
    win->frame_callback = cb;
    wl_callback_add_listener(cb, &callback_listener, win);
}

static void frame_new(void *data, struct wl_callback *cb, uint32_t _)
{
    struct wayland_window *win = data;
    if (win->frame_callback == cb)
        win->frame_callback = NULL;
    wl_callback_destroy(cb);
    setup_frame_callback(win);
    draw(win);
}

static void wm_base_ping(void *data, struct xdg_wm_base *wm_base, uint32_t serial)
{
    xdg_wm_base_pong(wm_base, serial);
}

int window_init(struct wl_display *display, struct wayland_window *win)
{
    (void)display;
    struct wl_compositor *compositor = registry_get_compositor();
    struct wl_shm *shm = registry_get_shm();
    struct xdg_wm_base *wm = registry_get_xdg_wm_base();
    if (!compositor || !shm || !wm)
        return -1;

    win->wl_surface = wl_compositor_create_surface(compositor);

    callback_listener.done = frame_new;
    win->frame_callback = NULL;
    setup_frame_callback(win);

    win->xdg_surface = xdg_wm_base_get_xdg_surface(wm, win->wl_surface);
    surface_listener.configure = xdg_surface_conf;
    xdg_surface_add_listener(win->xdg_surface, &surface_listener, win);

    win->xdg_toplevel = xdg_surface_get_toplevel(win->xdg_surface);

    top_listener.configure = xdg_toplevel_conf;
    top_listener.close = xdg_toplevel_cls;
    top_listener.configure_bounds = xdg_toplevel_bounds;
    top_listener.wm_capabilities = xdg_toplevel_wm_caps;
    xdg_toplevel_add_listener(win->xdg_toplevel, &top_listener, win);

    xdg_toplevel_set_title(win->xdg_toplevel, "Custom client");
    wl_surface_commit(win->wl_surface);

    wm_base_listener.ping = wm_base_ping;
    xdg_wm_base_add_listener(wm, &wm_base_listener, win);

    /* initialize state */
    win->width = 200;
    win->height = 100;
    win->data = NULL;
    win->buffer = NULL;
    win->frame_callback = NULL;
    win->need_close = 0;
    win->color = 0;

    return 0;
}

int window_should_close(struct wayland_window *win)
{
    return win->need_close;
}

void window_destroy(struct wayland_window *win)
{
    destroy_frame_callback(win);
    if (win->buffer)
        wl_buffer_destroy(win->buffer);
    if (win->xdg_toplevel)
        xdg_toplevel_destroy(win->xdg_toplevel);
    if (win->xdg_surface)
        xdg_surface_destroy(win->xdg_surface);
    if (win->wl_surface)
        wl_surface_destroy(win->wl_surface);
    if (win->data)
        munmap(win->data, win->width * win->height * 4);
}
