// SPDX-License-Identifier: BSD-3-Clause
/*
 * fbcopy_pub.c - capture the vGPU framebuffer WITHOUT X and publish frames
 * into the kvmfr ivshmem BAR.
 *
 * The X server's screen has 1280x720 pixels in the emulated VRAM, which the
 * guest maps as the vGPU's BAR1 (resource1). The X SHM readback path can
 * stall behind the X server's own push wait; this tool reads the framebuffer
 * pixels straight out of the BAR1 mapping instead (host RAM backing, no GPU
 * or X round-trip), swizzles BGRA -> RGBA and writes the kvmfr frame
 * protocol the host reader already decodes.
 *
 * Usage (guest, as root):
 *   ./fbcopy_pub test                     one solid red frame (pipeline check)
 *   ./fbcopy_pub find                     motion-scan BAR1, print hot offsets
 *   ./fbcopy_pub pub <off> [n] [fps] [w] [h] [pitch]
 *                                         copy w*h from BAR1@off (pitch)
 *
 * Find/build on the host (same glibc as the guest): gcc -O2 -o fbcopy_pub
 * fbcopy_pub.c
 */
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <dirent.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <time.h>

#define FRAME_OFFSET  0x100000UL
#define HDR_MAGIC     0x4b564d46u	/* "FMVK" */
#define DEF_W		1280
#define DEF_H		720
#define DEF_PITCH	5120		/* xorg-vgpu.conf: 1280*4 rounded */

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;
}

static const char *find_vgpu_resource1(void)
{
	static char path[320];
	DIR *d = opendir("/sys/bus/pci/devices");
	struct dirent *de;

	if (!d)
		return NULL;
	while ((de = readdir(d))) {
		char p[340], buf[32];
		FILE *f;
		int n;

		if (de->d_name[0] == '.')
			continue;
		snprintf(p, sizeof(p), "/sys/bus/pci/devices/%s/vendor",
			 de->d_name);
		f = fopen(p, "r");
		if (!f)
			continue;
		n = fread(buf, 1, sizeof(buf) - 1, f);
		fclose(f);
		if (n <= 0)
			continue;
		buf[n] = 0;
		if (strtoul(buf, NULL, 16) != 0x10de)
			continue;
		snprintf(path, sizeof(path),
			 "/sys/bus/pci/devices/%s/resource1", de->d_name);
		closedir(d);
		return path;
	}
	closedir(d);
	return NULL;
}

/*
 * The vfio mdev region mmap is limited (256 MB EINVAL, 64 MB OK): open the
 * resource once and map windows no larger than this via map_window().
 */
#define CHUNK_MAX (64UL << 20)

static unsigned char *map_window(int fd, size_t off, size_t len)
{
	if (len > CHUNK_MAX)
		return MAP_FAILED;
	return mmap(NULL, len, PROT_READ, MAP_SHARED, fd, (off_t)off);
}

static double t_ms(const struct timespec *a, const struct timespec *b)
{
	return (double)(b->tv_sec - a->tv_sec) * 1e3 +
	       (double)(b->tv_nsec - a->tv_nsec) / 1e6;
}

int main(int argc, char **argv)
{
	const char *mode = argc > 1 ? argv[1] : "find";
	const char *bar_path = "/sys/bus/pci/devices/0000:00:03.0/resource2";
	const char *res1 = find_vgpu_resource1();
	struct stat st;
	size_t map_size, i;
	unsigned char *bar;
	struct frame_hdr *hdr;
	volatile uint32_t *seqp;
	int fd, vram_fd;

	setvbuf(stdout, NULL, _IONBF, 0);
	if (!res1) {
		printf("NO_VGPU\n");
		return 1;
	}
	printf("VGPU_RES1 %s\n", res1);
	if (stat(res1, &st) || !st.st_size) {
		printf("STAT_FAIL\n");
		return 1;
	}
	map_size = (size_t)st.st_size;
	printf("VRAM_SIZE 0x%zx (%zu MB)\n", map_size, map_size >> 20);
	vram_fd = open(res1, O_RDONLY | O_SYNC);
	if (vram_fd < 0) {
		perror(res1);
		return 1;
	}

	/* kvmfr BAR next: every mode needs it except find. */
	fd = open(bar_path, O_RDWR | O_SYNC);
	if (fd < 0) {
		perror(bar_path);
		return 1;
	}
	bar = (unsigned char *)mmap(NULL, 128UL << 20, PROT_READ | PROT_WRITE,
				    MAP_SHARED, fd, 0);
	close(fd);
	if (bar == MAP_FAILED) {
		perror("mmap kvmfr");
		return 1;
	}
	hdr = (struct frame_hdr *)(bar + FRAME_OFFSET);
	seqp = &hdr->seq;

	if (!strcmp(mode, "test")) {
		unsigned char *dst = bar + FRAME_OFFSET +
				     sizeof(struct frame_hdr);
		size_t n = (size_t)DEF_W * DEF_H * 4;

		hdr->magic = HDR_MAGIC;
		hdr->version = 1;
		hdr->width = DEF_W;
		hdr->height = DEF_H;
		hdr->stride = DEF_W * 4;
		hdr->format = 0;
		for (i = 0; i < n; i += 4) {
			dst[i + 0] = (unsigned char)(argc > 2 ? atoi(argv[2])
							      : 255);
			dst[i + 1] = (unsigned char)(argc > 3 ? atoi(argv[3])
							      : 0);
			dst[i + 2] = (unsigned char)(argc > 4 ? atoi(argv[4])
							      : 0);
			dst[i + 3] = 0xff;
		}
		hdr->timestamp_ns = now_ns();
		__sync_synchronize();
		*seqp = 1;
		printf("TEST_FRAME_PUBLISHED\n");
		return 0;
	}

	if (!strcmp(mode, "find")) {
		/*
		 * Sparse motion scan: sample 256 B of every 4 KB page twice,
		 * 300 ms apart; score pages by differing bytes and report the
		 * hottest pages plus per-16MB moving counts. The framebuffer
		 * is the contiguous ~4 MB block with the most motion. The
		 * region is mapped one <=64 MB window at a time (the vfio
		 * mdev mmap refuses the full 256 MB).
		 */
		size_t pages = map_size >> 12;
		size_t cpages = CHUNK_MAX >> 12;
		unsigned char *s0, *s1;
		uint32_t *score;
		size_t p, chunk;
		int top;

		s0 = malloc(cpages * 256);
		s1 = malloc(cpages * 256);
		score = calloc(pages, sizeof(uint32_t));
		if (!s0 || !s1 || !score)
			return 1;
		for (chunk = 0; chunk * CHUNK_MAX < map_size; chunk++) {
			size_t off = chunk * CHUNK_MAX;
			size_t len = map_size - off;
			size_t np = 0;
			struct timespec a, b;
			unsigned char *w;

			if (len > CHUNK_MAX)
				len = CHUNK_MAX;
			w = map_window(vram_fd, off, len);
			if (w == MAP_FAILED) {
				printf("CHUNK_MAP_FAIL off=0x%zx\n", off);
				return 1;
			}
			np = len >> 12;
			for (p = 0; p < np; p++)
				memcpy(s0 + p * 256, w + (p << 12), 256);
			clock_gettime(CLOCK_MONOTONIC, &a);
			usleep(300000);
			for (p = 0; p < np; p++)
				memcpy(s1 + p * 256, w + (p << 12), 256);
			clock_gettime(CLOCK_MONOTONIC, &b);
			munmap(w, len);
			printf("SCAN chunk=0x%zx pages=%zu dt=%.0fms\n",
			       off, np, t_ms(&a, &b));
			for (p = 0; p < np; p++) {
				unsigned char *x = s0 + p * 256;
				unsigned char *y = s1 + p * 256;
				uint32_t d = 0;
				int j;

				for (j = 0; j < 256; j++)
					if (x[j] != y[j])
						d++;
				score[(off >> 12) + p] = d;
			}
		}
		/* per-16MB summary */
		for (i = 0; i + 4096 <= pages; i += 4096) {
			uint64_t s = 0;
			size_t j;

			for (j = i; j < i + 4096; j++)
				s += score[j];
			if (s)
				printf("BLOCK off=0x%zx moving_bytes=%llu\n",
				       i << 12, (unsigned long long)s);
		}
		for (top = 0; top < 16; top++) {
			size_t best = 0;
			uint32_t bs = 0;

			for (p = 0; p < pages; p++)
				if (score[p] > bs) {
					bs = score[p];
					best = p;
				}
			if (!bs)
				break;
			printf("HOT off=0x%zx score=%u\n", best << 12, bs);
			score[best] = 0;
		}
		printf("FIND_DONE\n");
		return 0;
	}

	if (!strcmp(mode, "pub")) {
		unsigned long long off = argc > 2 ?
			strtoull(argv[2], NULL, 0) : 0;
		int frames = argc > 3 ? atoi(argv[3]) : 30;
		int fps = argc > 4 ? atoi(argv[4]) : 5;
		int w = argc > 5 ? atoi(argv[5]) : DEF_W;
		int h = argc > 6 ? atoi(argv[6]) : DEF_H;
		int pitch = argc > 7 ? atoi(argv[7]) : DEF_PITCH;
		unsigned long long need = off +
			(unsigned long long)(h - 1) * pitch +
			(unsigned long long)w * 4;
		size_t aligned = (size_t)off & ~0xfffUL;
		size_t win = (size_t)(need - aligned);
		unsigned char *wmap;
		int f, y, x;

		if (need > map_size) {
			printf("OFF_RANGE_FAIL need=0x%llx size=0x%zx\n",
			       need, map_size);
			return 1;
		}
		if (win > CHUNK_MAX) {
			printf("WIN_TOO_BIG 0x%zx\n", win);
			return 1;
		}
		wmap = map_window(vram_fd, aligned, win);
		if (wmap == MAP_FAILED) {
			printf("WIN_MAP_FAIL off=0x%zx len=0x%zx\n",
			       aligned, win);
			return 1;
		}
		printf("PUB off=0x%llx frames=%d fps=%d %dx%d pitch=%d\n",
		       off, frames, fps, w, h, pitch);
		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 (f = 0; f < frames; f++) {
			uint64_t t0 = now_ns();
			unsigned char *dst = bar + FRAME_OFFSET +
					     sizeof(struct frame_hdr);
			unsigned char *src = wmap + ((size_t)off - aligned);

			for (y = 0; y < h; y++) {
				unsigned char *s = src + (size_t)y * pitch;
				unsigned char *d = dst + (size_t)y * w * 4;

				for (x = 0; x < w; x++) {
					d[4 * x + 0] = s[4 * x + 2];
					d[4 * x + 1] = s[4 * x + 1];
					d[4 * x + 2] = s[4 * x + 0];
					d[4 * x + 3] = 0xff;
				}
			}
			hdr->timestamp_ns = now_ns();
			__sync_synchronize();
			*seqp = (uint32_t)(f + 1);
			__sync_synchronize();
			printf("FRAME seq=%d ms=%.1f\n", f + 1,
			       (double)(now_ns() - t0) / 1e6);
			if (fps > 0)
				usleep((useconds_t)(1000000 / fps));
		}
		printf("FRAME_DONE frames=%d\n", f);
		return 0;
	}

	printf("USAGE: fbcopy_pub test|find|pub <off> [n] [fps] [w] [h] [pitch]\n");
	return 1;
}