#include <stdio.h>
#include <string.h>
#include <stdint.h>

#include <syscall.h>
#include <unistd.h>
#include <sys/mman.h>
#include <poll.h>

#include <wayland-client.h>

// xdg-shell is generated from xdg-shell.xml by wayland-scanner (see compile.sh).
#include "xdg-shell-client-protocol.h"

struct wl_compositor *compositor;
struct wl_shm *shm;
struct xdg_wm_base *wm_base;
struct wl_seat *seat;

struct wl_surface *surface;
struct wl_buffer *buffer;
uint32_t *pixels;           // the mmap'd framebuffer (XRGB8888)

// ---- game geometry / state ------------------------------------------------

#define WIDTH   400
#define HEIGHT  300

#define PADDLE_W   10
#define PADDLE_H   60
#define PADDLE_X   15          // left paddle, fixed x
#define PADDLE_SPEED 6

#define BALL_SIZE  10

float paddle_y = (HEIGHT - PADDLE_H) / 2;
float ball_x = WIDTH / 2, ball_y = HEIGHT / 2;
float ball_vx = 3, ball_vy = 2;

// which movement keys are currently held (updated on key press/release)
int key_up = 0, key_down = 0;

// cleared when ESC is pressed, which ends the game loop
int running = 1;

// ---- drawing helpers ------------------------------------------------------

void fill_rect(int x, int y, int w, int h, uint32_t color)
{
    for (int j = y; j < y + h; j++) {
        if (j < 0 || j >= HEIGHT) continue;
        for (int i = x; i < x + w; i++) {
            if (i < 0 || i >= WIDTH) continue;
            pixels[j * WIDTH + i] = color;
        }
    }
}

void render(void)
{
    fill_rect(0, 0, WIDTH, HEIGHT, 0x00000000);              // clear to black
    fill_rect(PADDLE_X, (int)paddle_y, PADDLE_W, PADDLE_H, 0x00ffffff);
    fill_rect((int)ball_x, (int)ball_y, BALL_SIZE, BALL_SIZE, 0x0000ff00);

    wl_surface_attach(surface, buffer, 0, 0);
    wl_surface_damage(surface, 0, 0, WIDTH, HEIGHT);
    wl_surface_commit(surface);
}

// ---- game logic -----------------------------------------------------------

void update(void)
{
    // move paddle while a key is held, clamp to the window
    if (key_up)   paddle_y -= PADDLE_SPEED;
    if (key_down) paddle_y += PADDLE_SPEED;
    if (paddle_y < 0) paddle_y = 0;
    if (paddle_y > HEIGHT - PADDLE_H) paddle_y = HEIGHT - PADDLE_H;

    // move ball
    ball_x += ball_vx;
    ball_y += ball_vy;

    // bounce off top / bottom
    if (ball_y < 0) { ball_y = 0; ball_vy = -ball_vy; }
    if (ball_y > HEIGHT - BALL_SIZE) { ball_y = HEIGHT - BALL_SIZE; ball_vy = -ball_vy; }

    // bounce off the right wall
    if (ball_x > WIDTH - BALL_SIZE) { ball_x = WIDTH - BALL_SIZE; ball_vx = -ball_vx; }

    // paddle collision on the left
    if (ball_x <= PADDLE_X + PADDLE_W &&
        ball_x + BALL_SIZE >= PADDLE_X &&
        ball_y + BALL_SIZE >= paddle_y &&
        ball_y <= paddle_y + PADDLE_H &&
        ball_vx < 0) {
        ball_x = PADDLE_X + PADDLE_W;
        ball_vx = -ball_vx;
    }

    // ball missed the paddle: reset to center
    if (ball_x < 0) {
        ball_x = WIDTH / 2;
        ball_y = HEIGHT / 2;
        ball_vx = 3;
        printf("miss! ball reset\n");
    }
}

// ---- keyboard -------------------------------------------------------------

// Linux evdev keycodes (as delivered by wl_keyboard, offset by 8 vs. XKB).
#define KEY_W     17
#define KEY_S     31
#define KEY_UP    103
#define KEY_DOWN  108
#define KEY_ESC   1

void keyboard_key_handler
(
    void *data,
    struct wl_keyboard *keyboard,
    uint32_t serial,
    uint32_t time,
    uint32_t key,
    uint32_t state
) {
    int pressed = (state == WL_KEYBOARD_KEY_STATE_PRESSED);
    if (key == KEY_UP   || key == KEY_W) key_up   = pressed;
    if (key == KEY_DOWN || key == KEY_S) key_down = pressed;
    if (key == KEY_ESC  && pressed) running = 0;
}

void keyboard_noop() {}

const struct wl_keyboard_listener keyboard_listener = {
    .keymap      = (void *)keyboard_noop,
    .enter       = (void *)keyboard_noop,
    .leave       = (void *)keyboard_noop,
    .key         = keyboard_key_handler,
    .modifiers   = (void *)keyboard_noop,
    .repeat_info = (void *)keyboard_noop,
};

// ---- wayland boilerplate --------------------------------------------------

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

const struct xdg_wm_base_listener wm_base_listener = {
    .ping = wm_base_ping_handler
};

void registry_global_handler
(
    void *data,
    struct wl_registry *registry,
    uint32_t name,
    const char *interface,
    uint32_t version
) {
    if (strcmp(interface, "wl_compositor") == 0) {
        compositor = wl_registry_bind(registry, name, &wl_compositor_interface, 3);
    } else if (strcmp(interface, "wl_shm") == 0) {
        shm = wl_registry_bind(registry, name, &wl_shm_interface, 1);
    } else if (strcmp(interface, "xdg_wm_base") == 0) {
        wm_base = wl_registry_bind(registry, name, &xdg_wm_base_interface, 1);
        xdg_wm_base_add_listener(wm_base, &wm_base_listener, NULL);
    } else if (strcmp(interface, "wl_seat") == 0) {
        seat = wl_registry_bind(registry, name, &wl_seat_interface, 1);
    }
}

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

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

void xdg_surface_configure_handler
(
    void *data,
    struct xdg_surface *xdg_surface,
    uint32_t serial
) {
    xdg_surface_ack_configure(xdg_surface, serial);
}

const struct xdg_surface_listener xdg_surface_listener = {
    .configure = xdg_surface_configure_handler
};

int main(void)
{
    struct wl_display *display = wl_display_connect(NULL);
    struct wl_registry *registry = wl_display_get_registry(display);
    wl_registry_add_listener(registry, &registry_listener, NULL);
    wl_display_roundtrip(display);

    if (!compositor || !shm || !wm_base) {
        fprintf(stderr, "missing required global: compositor=%p shm=%p wm_base=%p\n",
            (void *)compositor, (void *)shm, (void *)wm_base);
        return 1;
    }

    surface = wl_compositor_create_surface(compositor);
    struct xdg_surface *xdg_surface = xdg_wm_base_get_xdg_surface(wm_base, surface);
    xdg_surface_add_listener(xdg_surface, &xdg_surface_listener, NULL);
    struct xdg_toplevel *xdg_toplevel = xdg_surface_get_toplevel(xdg_surface);
    xdg_toplevel_set_title(xdg_toplevel, "pong");

    wl_surface_commit(surface);
    wl_display_roundtrip(display);

    int stride = WIDTH * 4;
    int size = stride * HEIGHT;

    int fd = syscall(SYS_memfd_create, "buffer", 0);
    ftruncate(fd, size);
    pixels = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);

    struct wl_shm_pool *pool = wl_shm_create_pool(shm, fd, size);
    buffer = wl_shm_pool_create_buffer(pool, 0, WIDTH, HEIGHT, stride, WL_SHM_FORMAT_XRGB8888);

    if (seat) {
        struct wl_keyboard *keyboard = wl_seat_get_keyboard(seat);
        wl_keyboard_add_listener(keyboard, &keyboard_listener, NULL);
    }

    // Game loop: pump wayland events, advance the game, draw.
    //
    // non-blocking pump each frame:
    //   prepare_read            -> arm a read (drain queue first if it fails)
    //   poll(fd, 16ms)          -> wait briefly for input, but don't block the
    //                              ball forever when nothing is pressed
    //   read_events / cancel    -> consume from fd, or cancel if nothing came
    //   dispatch_pending        -> run the handlers (e.g. the key handler)
    int fd_display = wl_display_get_fd(display);
    while (running) {
        while (wl_display_prepare_read(display) != 0)
            wl_display_dispatch_pending(display);
        wl_display_flush(display);

        struct pollfd pfd = { .fd = fd_display, .events = POLLIN };
        if (poll(&pfd, 1, 16) > 0 && (pfd.revents & POLLIN))
            wl_display_read_events(display);    // data waiting: read it
        else
            wl_display_cancel_read(display);    // nothing: undo prepare_read

        wl_display_dispatch_pending(display);   // invoke handlers

        update();
        render();
        wl_display_flush(display);
    }

    // ESC was pressed: tear everything down cleanly.
    xdg_toplevel_destroy(xdg_toplevel);
    xdg_surface_destroy(xdg_surface);
    wl_surface_destroy(surface);
    wl_buffer_destroy(buffer);
    wl_shm_pool_destroy(pool);
    munmap(pixels, size);
    close(fd);
    wl_display_disconnect(display);
    return 0;
}