#include <rfb/rfb.h>
#include <fcntl.h>
#include <stdio.h>
#include <errno.h>
#include <sys/mman.h>

#include <xf86drm.h>
#include <libdrm/i915_drm.h>
#include <libdrm/amdgpu_drm.h>
#include <xf86drmMode.h>
#include <linux/uinput.h>
#include <xkbcommon/xkbcommon.h>
#include <math.h>


struct drm_nvidia_gem_map_offset_params {
    uint32_t handle;             /* IN Handle to gem object */
    uint32_t __pad;

    uint64_t offset;             /* OUT Fake offset */
};
#define DRM_NVIDIA_GEM_MAP_OFFSET 0x0a
#define DRM_IOCTL_NVIDIA_GEM_MAP_OFFSET                                 \
    DRM_IOWR((DRM_COMMAND_BASE + DRM_NVIDIA_GEM_MAP_OFFSET),            \
             struct drm_nvidia_gem_map_offset_params)

#define ARRAY_SIZE(x) (sizeof(x) / sizeof(x[0]))
/* 10 frames per second (if we can) */
#define PICTURE_TIMEOUT (1.0/10.0)
#define UINPUT_ABS_MAX INT16_MAX
#define UINPUT_MAX_KEY 256

struct Vec2d {
    int x;
    int y;
};
struct Vec2d resolution;

/*
 * throttle camera updates
*/
static int TimeToTakePicture() {
    static struct timeval now={0,0}, then={0,0};
    double elapsed, dnow, dthen;

    gettimeofday(&now,NULL);

    dnow  = now.tv_sec  + (now.tv_usec /1000000.0);
    dthen = then.tv_sec + (then.tv_usec/1000000.0);
    elapsed = dnow - dthen;

    if (elapsed > PICTURE_TIMEOUT)
        memcpy((char *)&then, (char *)&now, sizeof(struct timeval));
    return elapsed > PICTURE_TIMEOUT;
}

static void convert_bgrx_to_rgb(const char* in, int width, int height, char* buff) {
    for (int y = 0; y < height; y++)
    {
        for(int x = 0; x < width; x++)
        {
            buff[(y*width+x)*4] = in[(y*width+x)*4 + 2];
            buff[(y*width+x)*4 + 1] = in[(y*width+x)*4 + 1];
            buff[(y*width+x)*4 + 2] = in[(y*width+x)*4];
        }
    }
}

#define XSTRIPE 16
#define YSTRIPE 128
char *nv_stripe_buf = NULL;
static void convert_nv_stripe(const char* in, int width, int height, char* buff) {
    if (width % XSTRIPE) {
        return;
    }
    if (!nv_stripe_buf) nv_stripe_buf = malloc(width * height * 4);
    convert_bgrx_to_rgb(in, width, height, nv_stripe_buf);
    for (int i = 0; i < height*width; i++)
    {
        int sno = i / (XSTRIPE * YSTRIPE);
        int ord = i % (XSTRIPE * YSTRIPE);
        int base_x = sno % (width / XSTRIPE) * XSTRIPE;
        int base_y = sno / (width / XSTRIPE) * YSTRIPE;
        int sx = ord % XSTRIPE + base_x;
        int sy = ord / XSTRIPE + base_y;
        int offset = sy * width + sx;
        if (offset < height*width) {
            memcpy(buff + offset*4, nv_stripe_buf + i*4, 4);
        }
    }
}

struct vnc_xkb {
    struct xkb_context *ctx;
	struct xkb_keymap *map;
};
struct vnc_xkb xkb;

int uinput_fd = 0;
char keystate[UINPUT_MAX_KEY];
static void xkb_init() {
    xkb.ctx = xkb_context_new(XKB_CONTEXT_NO_FLAGS);
	if (xkb.ctx == NULL) {
        fprintf(stderr, "Failed to create XKB context\n");
        exit(1);
	}
	struct xkb_rule_names names = {
		.rules = "",
		.model = "",
		.layout = "us",
		.variant = "",
		.options = ""
	};
	xkb.map = xkb_keymap_new_from_names(xkb.ctx, &names, 0);
    if (xkb.map == NULL) {
        fprintf(stderr, "Failed to create XKB map\n");
        exit(1);
    }
    //printf("xkb: keymap = %s\n", xkb_keymap_get_as_string(xkb.map, XKB_KEYMAP_USE_ORIGINAL_FORMAT));
}

static void uinput_init() {
    struct uinput_setup usetup;
    uinput_fd = open("/dev/uinput", O_WRONLY | O_NONBLOCK);
    if (uinput_fd <= 0) {
        fprintf(stderr, "Failed to open uinput\n");
        exit(1);
    }
    ioctl(uinput_fd, UI_SET_EVBIT, EV_KEY);
    ioctl(uinput_fd, UI_SET_EVBIT, EV_SYN);
    for (int i = 0; i < UINPUT_MAX_KEY; i++) {
        ioctl(uinput_fd, UI_SET_KEYBIT, i);
    }


    ioctl(uinput_fd, UI_SET_EVBIT, EV_ABS);
    ioctl(uinput_fd, UI_SET_ABSBIT, ABS_X);
    ioctl(uinput_fd, UI_SET_ABSBIT, ABS_Y);

	ioctl(uinput_fd, UI_SET_KEYBIT, BTN_LEFT);
	ioctl(uinput_fd, UI_SET_KEYBIT, BTN_MIDDLE);
	ioctl(uinput_fd, UI_SET_KEYBIT, BTN_RIGHT);
	ioctl(uinput_fd, UI_SET_EVBIT, EV_REL);
    ioctl(uinput_fd, UI_SET_RELBIT, REL_WHEEL);

	struct uinput_abs_setup abs;
	memset(&abs, 0, sizeof(abs));
	abs.absinfo.maximum = UINPUT_ABS_MAX;
	abs.absinfo.minimum = 0;
	abs.code = ABS_X;
	ioctl(uinput_fd, UI_ABS_SETUP, &abs);
	abs.code = ABS_Y;
	ioctl(uinput_fd, UI_ABS_SETUP, &abs);

    memset(&usetup, 0, sizeof(usetup));
    usetup.id.bustype = BUS_USB;
    usetup.id.vendor = 0x0011;
    usetup.id.product = 0x4514;
    strcpy(usetup.name, "kmsvnc-device");

    ioctl(uinput_fd, UI_DEV_SETUP, &usetup);
    ioctl(uinput_fd, UI_DEV_CREATE);

    memset(keystate, 0, UINPUT_MAX_KEY);
}

struct key_iter_search {
	xkb_keysym_t keysym;

	xkb_keycode_t keycode;
	xkb_level_index_t level;
};

static void key_iter(struct xkb_keymap *xkb, xkb_keycode_t key, void *data) {
	struct key_iter_search *search = data;
	if (search->keycode != XKB_KEYCODE_INVALID) {
		return;  // We are done
	}
	xkb_level_index_t num_levels = xkb_keymap_num_levels_for_key(xkb, key, 0);
	for (xkb_level_index_t i = 0; i < num_levels; i++) {
		const xkb_keysym_t *syms;
		int num_syms = xkb_keymap_key_get_syms_by_level(xkb, key, 0, i, &syms);
		for (int k = 0; k < num_syms; k++) {
			if (syms[k] == search->keysym) {
				search->keycode = key;
				search->level = i;
				break;
				goto end;
			}
		}
	}
end:
	return;
}

static void rfb_key_hook(rfbBool down, rfbKeySym keysym, rfbClientPtr cl) {
	struct key_iter_search search = {
		.keysym = keysym,
		.keycode = XKB_KEYCODE_INVALID,
		.level = 0,
	};
	xkb_keymap_key_for_each(xkb.map, key_iter, &search);
	if (search.keycode == XKB_KEYCODE_INVALID) {
		fprintf(stderr, "Keysym %04x not found in our keymap\n", keysym);
		return;
	}
    printf("key %s, keysym %04x, keycode %u\n", down ? "down" : "up", keysym, search.keycode);
    if (search.keycode >= UINPUT_MAX_KEY) {
		fprintf(stderr, "Keycode %d >= %d\n", search.keycode, UINPUT_MAX_KEY);
		return;
    }
    if (down != keystate[search.keycode]) {
        struct input_event ies[] = {
            {
                .type = EV_KEY,
                .code = search.keycode - 8, // magic
                .value = down,
                .time.tv_sec = 0,
                .time.tv_usec = 0,
            },
            {
                .type = EV_SYN,
                .code = SYN_REPORT,
                .value = 0,
            },
        };
        for (int i = 0; i < ARRAY_SIZE(ies); i++) {
            write(uinput_fd, &ies[i], sizeof(ies[0]));
        }

        keystate[search.keycode] = down;
    }
}

static void rfb_ptr_hook(int mask, int screen_x, int screen_y, rfbClientPtr cl) {
    printf("pointer to %d, %d\n", screen_x, screen_y);
	float global_x = (float)screen_x;
	float global_y = (float)screen_y;
	int touch_x = round(global_x / resolution.x * UINPUT_ABS_MAX);
	int touch_y = round(global_y / resolution.y * UINPUT_ABS_MAX);
    struct input_event ies1[] = {
		{
			.type = EV_ABS,
			.code = ABS_X,
			.value = touch_x,
		},
		{
			.type = EV_ABS,
			.code = ABS_Y,
			.value = touch_y,
		},
		{
			.type = EV_KEY,
			.code = BTN_LEFT,
			.value = !! (mask & 0b1)
		},
		{
			.type = EV_KEY,
			.code = BTN_MIDDLE,
			.value = !! (mask & 0b10)
		},
		{
			.type = EV_KEY,
			.code = BTN_RIGHT,
			.value = !! (mask & 0b100)
		},
		{
			.type = EV_SYN,
			.code = SYN_REPORT,
			.value = 0,
		},
    };
    for (int i = 0; i < ARRAY_SIZE(ies1); i++) {
        write(uinput_fd, &ies1[i], sizeof(ies1[0]));
    }
    if (mask & 0b11000) {
        struct input_event ies2[] = {
        	{
        		.type = EV_REL,
        		.code = REL_WHEEL,
        		.value = mask & 0b1000 ? -1 : 1,
        	},
        	{
        		.type = EV_SYN,
        		.code = SYN_REPORT,
        		.value = 0,
        	},
        };
        for (int i = 0; i < ARRAY_SIZE(ies2); i++) {
            write(uinput_fd, &ies2[i], sizeof(ies2[0]));
        }
	}
}


static void drm_wait_vblank_noop(int x) {
}
static void drm_wait_vblank(int drmfd) {
    int ioctl_err;
	union drm_wait_vblank vbl;
	memset(&vbl, 0, sizeof(vbl));
	vbl.request.type = DRM_VBLANK_RELATIVE;
	vbl.request.sequence = 1;
    if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_WAIT_VBLANK, &vbl))
        fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
}

int main(int argc, const char **argv)
{
    if (argc < 2) {
        printf("not enough arguments\n");
        return 1;
    }
    xkb_init();
    uinput_init();

    const char *card = argv[1];
    const int drmfd = open(card, O_RDONLY);
    if (drmfd < 0) {
        fprintf(stderr, "card %s open failed: %s\n", card, strerror(errno));
        return 1;
    }
    drmVersionPtr drm_ver = drmGetVersion(drmfd);
    int err;
    int source_plane = 0;
    int source_crtc = 0;
    err = drmSetClientCap(drmfd, DRM_CLIENT_CAP_UNIVERSAL_PLANES, 1);
    if (err < 0) {
        perror("Failed to set universal planes capability: primary planes will not be usable");
	}
    drmModePlane *plane = NULL;
    drmModePlaneRes *plane_res = NULL;
    drmModeFB *fb = NULL;
    if (source_plane > 0) {
        plane = drmModeGetPlane(drmfd, source_plane);
        if (!plane) {
            fprintf(stderr, "Failed to get plane %d: %s\n", source_plane, strerror(errno));
            goto cleanup;
        }
        if (plane->fb_id == 0) {
            fprintf(stderr, "Place %d does not have an attached framebuffer\n", source_plane);
        }
    }
    else {
        plane_res = drmModeGetPlaneResources(drmfd);
        if (!plane_res) {
            perror("Failed to get plane resources");
            goto cleanup;
        }
        int i;
        for (i = 0; i < plane_res->count_planes; i++) {
            plane = drmModeGetPlane(drmfd, plane_res->planes[i]);
            if (!plane) {
                fprintf(stderr, "Failed to get plane %u: %s\n", plane_res->planes[i], strerror(errno));
                continue;
            }
            printf("Plane %u CRTC %u FB %u\n", plane->plane_id, plane->crtc_id, plane->fb_id);
            if ((source_crtc > 0 && plane->crtc_id != source_crtc) || plane->fb_id == 0) {
                // Either not connected to the target source CRTC
                // or not active.
                drmModeFreePlane(plane);
                plane = NULL;
                continue;
            }
            break;
        }
        if (i == plane_res->count_planes) {
            if (source_crtc > 0) {
                fprintf(stderr, "No usable planes found on CRTC %d\n", source_crtc);
            } else {
                fprintf(stderr, "No usable planes found\n");
            }
            goto cleanup;
        }
        printf("Using plane %u to locate framebuffers\n", plane->plane_id);
    }
    uint32_t plane_id = plane->plane_id;

    fb = drmModeGetFB(drmfd, plane->fb_id);
    if (!fb) {
        fprintf(stderr, "Failed to get framebuffer %u: %s\n", plane->fb_id, strerror(errno));
        goto cleanup;
    }
    printf("Template framebuffer is %u: %ux%u %ubpp %ub depth %u pitch\n", fb->fb_id, fb->width, fb->height, fb->bpp, fb->depth, fb->pitch);

    if (fb->bpp != 32 || fb->depth != 24) {
        fprintf(stderr, "Unsupported pixfmt\n");
        goto cleanup;
    }

    if (!fb->handle) {
        fprintf(stderr, "No handle set on framebuffer: maybe you need some additional capabilities?\n");
        goto cleanup;
    }

    int ioctl_err = 0;
    struct drm_gem_flink flink;
	flink.handle = fb->handle;
	if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_GEM_FLINK, &flink)) {
        fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
	    goto cleanup;
	}
	struct drm_gem_open open_arg;
    open_arg.name = flink.name;
	if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_GEM_OPEN, &open_arg)) {
        fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
	    goto cleanup;
    }
    off_t mmap_offset = 0;
    size_t mmap_size = 0;
    printf("drm driver is %s\n", drm_ver->name);
    char *mapped = NULL;
    void (*convert_func)(const char*, int, int, char*) = &convert_bgrx_to_rgb;
    void (*vblank_func)(int) = &drm_wait_vblank;
    if (strcmp(drm_ver->name, "i915") == 0) {
    	struct drm_i915_gem_mmap_gtt mmap_arg;
    	mmap_arg.handle = open_arg.handle;
    	if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_I915_GEM_MMAP_GTT, &mmap_arg)) {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
    	    goto cleanup;
    	}
    	mmap_size = open_arg.size;
    	mmap_offset = mmap_arg.offset;
    }
    else if (strcmp(drm_ver->name, "amdgpu") == 0) {
    	union drm_amdgpu_gem_mmap mmap_arg;
    	memset(&mmap_arg, 0, sizeof(mmap_arg));
    	mmap_arg.in.handle = open_arg.handle;
    	if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_AMDGPU_GEM_MMAP, &mmap_arg)) {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
    	    goto cleanup;
    	}
    	mmap_size = open_arg.size;
    	mmap_offset = mmap_arg.out.addr_ptr;
    }
    else if (strcmp(drm_ver->name, "nvidia-drm") == 0) {
        convert_func = &convert_nv_stripe;
        vblank_func = &drm_wait_vblank_noop;
        struct drm_nvidia_gem_map_offset_params mmap_arg;
        memset(&mmap_arg, 0, sizeof(mmap_arg));
        mmap_arg.handle = open_arg.handle;
        if (ioctl_err = drmIoctl(drmfd, DRM_IOCTL_NVIDIA_GEM_MAP_OFFSET, &mmap_arg)) {
            fprintf(stderr, "DRM ioctl error %d on line %d\n", ioctl_err, __LINE__);
    	    goto cleanup;
    	}
    	mmap_size = open_arg.size;
        mmap_offset = mmap_arg.offset;
    }
    else {
        fprintf(stderr, "Unimplemented drm driver\n");
        goto cleanup;
    }
    if (!mapped) {
        printf("mapping with size = %d, offset = %d, drmfd = %d\n", mmap_size, mmap_offset, drmfd);
	    mapped = mmap(NULL, mmap_size, PROT_READ, MAP_SHARED, drmfd, mmap_offset);
        if (mapped == MAP_FAILED) {
            perror("mmap");
            goto cleanup;
        }
    }

    size_t buflen = fb->width * fb->height * 32 / 8;
    char *buf = malloc(buflen);
    memset(buf, 0, buflen);
    char *buf2 = malloc(buflen);
    memset(buf2, 0, buflen);

    resolution.x = fb->width;
    resolution.y = fb->height;
    rfbScreenInfoPtr server=rfbGetScreen(0,NULL,fb->width,fb->height,8,3,32/8);
    if(!server)
        return 1;
    server->desktopName = "kmsvnc";
    server->frameBuffer = buf;
    server->port = 5900;
    //server->listenInterface = inet_addr("127.0.0.1");
    server->ipv6port = 0;
    server->listen6Interface = NULL;
    server->alwaysShared = (1==1);
    server->kbdAddEvent = rfb_key_hook;
    server->ptrAddEvent = rfb_ptr_hook;
    rfbInitServer(server);
    long usec;
    while (rfbIsActive(server)) {
        if (server->clientHead && TimeToTakePicture()) {
            vblank_func(drmfd);
            // still tearing
            memcpy(buf2, mapped, buflen);
            convert_func(buf2, fb->width, fb->height, buf);
            rfbMarkRectAsModified(server, 0, 0, fb->width, fb->height);
        }
        usec = server->deferUpdateTime*1000;
        rfbProcessEvents(server,usec);
    }

    if (drm_ver != NULL) drmFreeVersion(drm_ver);
    if (fb != NULL) drmModeFreeFB(fb);
    if (uinput_fd > 0) {
        ioctl(uinput_fd, UI_DEV_DESTROY);
        close(uinput_fd);
    }
    close(drmfd);
    return 0;

    cleanup:
        if (drm_ver != NULL) drmFreeVersion(drm_ver);
        if (fb != NULL) drmModeFreeFB(fb);
        if (uinput_fd > 0) {
            ioctl(uinput_fd, UI_DEV_DESTROY);
            close(uinput_fd);
        }
        close(drmfd);
        return 1;
}