// SPDX-License-Identifier: BSD-3-Clause
/*
 * xshm_pub.c - capture an X11 screen and publish frames into the kvmfr
 * ivshmem BAR so the host can display them (the Looking-Glass-style path:
 * the guest renders with its real graphics stack - X on the passed GPU /
 * GLX / Vulkan surface - and this tool copies the framebuffer into the
 * shared BAR; the host reader decodes it).
 *
 * Layout in the BAR (same as tools/cuda-ab/frame_producer.cu):
 *   offset 0x100000: struct frame_hdr { magic, version, width, height,
 *   stride, format, seq } then width*height*4 RGBA8 bytes.
 *
 * Usage (guest, with DISPLAY set):
 *   ./xshm_pub [frames] [fps] [x] [y] [w] [h]
 * Prints FRAME lines (first 5 and every 30th) then FRAME_DONE.
 * Build (host, same glibc): gcc -O2 -o xshm_pub xshm_pub.c -lX11 -lXext
 */
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <time.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/extensions/XShm.h>
#include <sys/ipc.h>
#include <sys/shm.h>

#define FRAME_OFFSET  0x100000UL
#define HDR_MAGIC     0x4b564d46u	/* "FMVK" */

struct frame_hdr {
	uint32_t magic;
	uint32_t version;
	uint32_t width;
	uint32_t height;
	uint32_t stride;
	uint32_t format;	/* 0 = RGBA8 */
	uint32_t seq;
	uint32_t reserved;
	uint64_t timestamp_ns;
};

static uint64_t now_ns(void)
{
	struct timespec ts;

	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (uint64_t)ts.tv_sec * 1000000000ull + ts.tv_nsec;
}

int main(int argc, char **argv)
{
	int frames = argc > 1 ? atoi(argv[1]) : 120;
	int fps = argc > 2 ? atoi(argv[2]) : 10;
	int gx = argc > 3 ? atoi(argv[3]) : 0;
	int gy = argc > 4 ? atoi(argv[4]) : 0;
	int gw = argc > 5 ? atoi(argv[5]) : 0;
	int gh = argc > 6 ? atoi(argv[6]) : 0;
	const char *bar_path =
		"/sys/bus/pci/devices/0000:00:03.0/resource2";
	Display *dpy;
	Window root;
	XImage *img = NULL;
	XShmSegmentInfo shm;
	int use_shm = 1, fd, i, w, h;
	unsigned char *bar;
	struct frame_hdr *hdr;
	volatile uint32_t *seqp;
	uint64_t t_first = 0;
	uint64_t t_last = 0;

	setvbuf(stdout, NULL, _IONBF, 0);

	dpy = XOpenDisplay(NULL);
	if (!dpy) {
		printf("XOPEN_FAIL DISPLAY=%s\n",
		       getenv("DISPLAY") ? getenv("DISPLAY") : "(unset)");
		return 1;
	}
	root = DefaultRootWindow(dpy);
	{
		XWindowAttributes wa;

		XGetWindowAttributes(dpy, root, &wa);
		w = gw > 0 ? gw : wa.width;
		h = gh > 0 ? gh : wa.height;
		printf("XSCREEN %dx%d root=%dx%d\n", w, h, wa.width,
		       wa.height);
	}
	fd = open(bar_path, O_RDWR | O_SYNC);
	if (fd < 0) {
		printf("OPEN_FAIL %s\n", bar_path);
		return 1;
	}
	bar = (unsigned char *)mmap(NULL, 128UL << 20, PROT_READ | PROT_WRITE,
				    MAP_SHARED, fd, 0);
	if (bar == MAP_FAILED) {
		printf("MMAP_FAIL\n");
		return 1;
	}
	hdr = (struct frame_hdr *)(bar + FRAME_OFFSET);
	seqp = &hdr->seq;

	memset(&shm, 0, sizeof(shm));
	if (use_shm) {
		img = XShmCreateImage(dpy, DefaultVisual(dpy,
					DefaultScreen(dpy)),
					DefaultDepth(dpy, DefaultScreen(dpy)),
					ZPixmap, NULL, &shm, w, h);
		if (!img) {
			printf("XSHM_CREATE_FAIL\n");
			return 1;
		}
		shm.shmid = shmget(IPC_PRIVATE,
				   (size_t)img->bytes_per_line * h,
				   IPC_CREAT | 0600);
		if (shm.shmid < 0) {
			printf("SHMGET_FAIL\n");
			return 1;
		}
		shm.shmaddr = img->data = shmat(shm.shmid, NULL, 0);
		if (shm.shmaddr == (char *)-1) {
			printf("SHMAT_FAIL\n");
			return 1;
		}
		shm.readOnly = False;
		if (!XShmAttach(dpy, &shm)) {
			printf("XSHM_ATTACH_FAIL (falling back)\n");
			XDestroyImage(img);
			img = NULL;
		} else {
			XSync(dpy, False);
			printf("XSHM OK %dx%d bpl=%d\n", w, h,
			       img->bytes_per_line);
		}
	}
	if (!img) {
		img = XGetImage(dpy, root, gx, gy, w, h, AllPlanes, ZPixmap);
		if (!img) {
			printf("XGETIMAGE_FAIL\n");
			return 1;
		}
		printf("XGETIMAGE OK %dx%d bpl=%d\n", w, h,
		       img->bytes_per_line);
	}

	hdr->magic = HDR_MAGIC;
	hdr->version = 1;
	hdr->width = (uint32_t)w;
	hdr->height = (uint32_t)h;
	hdr->stride = (uint32_t)w * 4;
	hdr->format = 0;

	for (i = 0; i < frames; i++) {
		uint64_t t0 = now_ns();
		unsigned char *dst = bar + FRAME_OFFSET +
				     sizeof(struct frame_hdr);
		int x, y;

		if (shm.shmaddr) {
			XShmGetImage(dpy, root, img, gx, gy, AllPlanes);
		} else {
			XImage *n = XGetImage(dpy, root, gx, gy, w, h,
					      AllPlanes, ZPixmap);

			if (!n)
				break;
			XDestroyImage(img);
			img = n;
		}
		/* X ZPixmap: 0x00RRGGBB little-endian -> RGBA bytes */
		for (y = 0; y < h; y++) {
			unsigned char *s = (unsigned char *)img->data +
					   (size_t)y * img->bytes_per_line;
			unsigned char *d = dst + (size_t)y * w * 4;

			for (x = 0; x < w; x++) {
				uint32_t p;

				memcpy(&p, s + 4 * x, 4);
				d[4 * x + 0] = (unsigned char)(p >> 16);
				d[4 * x + 1] = (unsigned char)(p >> 8);
				d[4 * x + 2] = (unsigned char)(p);
				d[4 * x + 3] = 0xff;
			}
		}
		hdr->timestamp_ns = now_ns();
		if (!t_first)
			t_first = t0;
		t_last = t0;
		__sync_synchronize();
		*seqp = (uint32_t)(i + 1);
		__sync_synchronize();
		if (i < 5 || (i % 30) == 29)
			printf("FRAME seq=%d ms=%.1f\n", i + 1,
			       (double)(now_ns() - t0) / 1e6);
		if (fps > 0)
			usleep((useconds_t)(1000000 / fps));
	}
	if (t_first && t_last > t_first)
		printf("XPUB_RATE fps=%.2f frames=%d\n",
		       (double)(i - 1) * 1e9 /
			       (double)(t_last - t_first), i);
	printf("FRAME_DONE frames=%d %dx%d\n", i, w, h);
	if (shm.shmaddr)
		XShmDetach(dpy, &shm);
	XCloseDisplay(dpy);
	return 0;
}