// SPDX-License-Identifier: BSD-3-Clause
/*
 * vk_frame.c - headless Vulkan render probe for the vGPU guest.
 *
 * Creates a real Vulkan device on the NVIDIA ICD, runs a compute shader that
 * paints an animated plasma into an RGBA8 storage image, copies the image to
 * a host-visible buffer, saves a PPM, and (optionally) publishes each frame
 * into the kvmfr ivshmem BAR in the frame_producer layout so the host can
 * display it. No window system involved (pure headless), so it works even
 * while the GL/GLX walls are being fixed.
 *
 * Usage: ./vk_frame [frames] [width] [height] [fps] [outdir] [publish]
 * Build: gcc -O2 -o vk_frame vk_frame.c -lvulkan   (same glibc as the guest)
 */
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <time.h>
#include <sys/mman.h>
#include <vulkan/vulkan.h>

#include "vk_plasma_spv.h"

#define FRAME_OFFSET  0x100000UL
#define HDR_MAGIC     0x4b564d46u

struct frame_hdr {
	uint32_t magic;
	uint32_t version;
	uint32_t width;
	uint32_t height;
	uint32_t stride;
	uint32_t format;
	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;
}

#define VKC(x)                                                            \
	do {                                                              \
		VkResult r_ = (x);                                        \
		if (r_ != VK_SUCCESS) {                                   \
			printf("VKFAIL %s -> %d (line %d)\n", #x, r_,    \
			       __LINE__);                                 \
			return 1;                                         \
		}                                                         \
	} while (0)

int main(int argc, char **argv)
{
	int frames = argc > 1 ? atoi(argv[1]) : 60;
	uint32_t W = argc > 2 ? (uint32_t)atoi(argv[2]) : 640;
	uint32_t H = argc > 3 ? (uint32_t)atoi(argv[3]) : 480;
	int fps = argc > 4 ? atoi(argv[4]) : 10;
	const char *outdir = argc > 5 ? argv[5] : NULL;
	int publish = argc > 6 ? atoi(argv[6]) : 1;
	VkApplicationInfo app = {0};
	VkInstanceCreateInfo ici = {0};
	VkInstance inst;
	uint32_t ndev = 0;
	VkPhysicalDevice devs[16];
	VkPhysicalDevice phys = VK_NULL_HANDLE;
	char picked[256] = "";
	uint32_t nfam = 0, fam = ~0u;
	VkQueueFamilyProperties fams[16];
	VkDevice dev;
	VkQueue queue;
	VkImage image;
	VkDeviceMemory img_mem;
	VkImageView view;
	VkBuffer buf;
	VkDeviceMemory buf_mem;
	VkShaderModule shader;
	VkDescriptorSetLayout dsl;
	VkPipelineLayout pl;
	VkPipeline pipe;
	VkDescriptorPool pool;
	VkDescriptorSet dset;
	VkCommandPool cpool;
	VkCommandBuffer cb;
	VkFence fence;
	uint32_t i;
	float prio = 1.0f;
	VkDeviceSize imgsz = (VkDeviceSize)W * H * 4;
	unsigned char *pix;
	struct frame_hdr *hdr = NULL;
	volatile uint32_t *seqp = NULL;
	unsigned char *bar = NULL;
	int bfd = -1;
	int saved = 0;
	uint64_t t_first = 0, t_last = 0;

	setvbuf(stdout, NULL, _IONBF, 0);

	app.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
	app.pApplicationName = "vk_frame";
	app.pEngineName = "vk_frame";
	app.apiVersion = VK_API_VERSION_1_2;
	ici.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
	ici.pApplicationInfo = &app;
	VKC(vkCreateInstance(&ici, NULL, &inst));
	printf("VKINSTANCE ok\n");

	VKC(vkEnumeratePhysicalDevices(inst, &ndev, NULL));
	if (ndev > 16)
		ndev = 16;
	VKC(vkEnumeratePhysicalDevices(inst, &ndev, devs));
	printf("VKDEVS %u\n", ndev);
	for (i = 0; i < ndev; i++) {
		VkPhysicalDeviceProperties pp;

		vkGetPhysicalDeviceProperties(devs[i], &pp);
		printf("VKDEV %u: %s vendor=0x%04x\n", i, pp.deviceName,
		       pp.vendorID);
		if (pp.vendorID == 0x10de || strstr(pp.deviceName, "NVIDIA")) {
			phys = devs[i];
			snprintf(picked, sizeof(picked), "%s",
				 pp.deviceName);
		}
	}
	if (!phys) {
		printf("VK_NO_NVIDIA\n");
		return 2;
	}
	printf("VKPICK %s\n", picked);
	vkGetPhysicalDeviceQueueFamilyProperties(phys, &nfam, NULL);
	if (nfam > 16)
		nfam = 16;
	vkGetPhysicalDeviceQueueFamilyProperties(phys, &nfam, fams);
	for (i = 0; i < nfam; i++)
		if ((fams[i].queueFlags & VK_QUEUE_COMPUTE_BIT) &&
		    fams[i].queueCount > 0) {
			fam = i;
			break;
		}
	if (fam == ~0u) {
		printf("VK_NO_COMPUTE_QUEUE\n");
		return 3;
	}
	{
		VkDeviceQueueCreateInfo qci = {0};
		VkDeviceCreateInfo dci = {0};

		qci.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
		qci.queueFamilyIndex = fam;
		qci.queueCount = 1;
		qci.pQueuePriorities = &prio;
		dci.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
		dci.queueCreateInfoCount = 1;
		dci.pQueueCreateInfos = &qci;
		VKC(vkCreateDevice(phys, &dci, NULL, &dev));
	}
	printf("VKDEVICE ok family=%u\n", fam);
	vkGetDeviceQueue(dev, fam, 0, &queue);

	/* storage image */
	{
		VkImageCreateInfo ci = {0};
		VkMemoryRequirements mr;
		VkMemoryAllocateInfo ma = {0};

		ci.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
		ci.imageType = VK_IMAGE_TYPE_2D;
		ci.format = VK_FORMAT_R8G8B8A8_UNORM;
		ci.extent.width = W;
		ci.extent.height = H;
		ci.extent.depth = 1;
		ci.mipLevels = 1;
		ci.arrayLayers = 1;
		ci.samples = VK_SAMPLE_COUNT_1_BIT;
		ci.tiling = VK_IMAGE_TILING_OPTIMAL;
		ci.usage = VK_IMAGE_USAGE_STORAGE_BIT |
			   VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
		ci.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
		VKC(vkCreateImage(dev, &ci, NULL, &image));
		vkGetImageMemoryRequirements(dev, image, &mr);
		ma.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
		ma.allocationSize = mr.size;
		/* find a device-local memory type */
		{
			VkPhysicalDeviceMemoryProperties mp;
			uint32_t k;

			vkGetPhysicalDeviceMemoryProperties(phys, &mp);
			for (k = 0; k < mp.memoryTypeCount; k++)
				if ((mr.memoryTypeBits & (1u << k)) &&
				    (mp.memoryTypes[k].propertyFlags &
				     VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)) {
					ma.memoryTypeIndex = k;
					break;
				}
		}
		VKC(vkAllocateMemory(dev, &ma, NULL, &img_mem));
		VKC(vkBindImageMemory(dev, image, img_mem, 0));
	}
	{
		VkImageViewCreateInfo ci = {0};

		ci.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
		ci.image = image;
		ci.viewType = VK_IMAGE_VIEW_TYPE_2D;
		ci.format = VK_FORMAT_R8G8B8A8_UNORM;
		ci.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		ci.subresourceRange.levelCount = 1;
		ci.subresourceRange.layerCount = 1;
		VKC(vkCreateImageView(dev, &ci, NULL, &view));
	}
	/* host-visible readback buffer */
	{
		VkBufferCreateInfo ci = {0};
		VkMemoryRequirements mr;
		VkMemoryAllocateInfo ma = {0};

		ci.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
		ci.size = imgsz;
		ci.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT;
		ci.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
		VKC(vkCreateBuffer(dev, &ci, NULL, &buf));
		vkGetBufferMemoryRequirements(dev, buf, &mr);
		ma.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
		ma.allocationSize = mr.size;
		{
			VkPhysicalDeviceMemoryProperties mp;
			uint32_t k;

			vkGetPhysicalDeviceMemoryProperties(phys, &mp);
			for (k = 0; k < mp.memoryTypeCount; k++)
				if ((mr.memoryTypeBits & (1u << k)) &&
				    (mp.memoryTypes[k].propertyFlags &
				     (VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
				      VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))) {
					ma.memoryTypeIndex = k;
					break;
				}
		}
		VKC(vkAllocateMemory(dev, &ma, NULL, &buf_mem));
		VKC(vkBindBufferMemory(dev, buf, buf_mem, 0));
	}
	/* shader + pipeline */
	{
		VkShaderModuleCreateInfo ci = {0};
		VkDescriptorSetLayoutBinding b = {0};
		VkDescriptorSetLayoutCreateInfo dci = {0};
		VkPushConstantRange pcr = {0};
		VkPipelineLayoutCreateInfo plci = {0};
		VkComputePipelineCreateInfo pci = {0};

		ci.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
		ci.codeSize = vk_plasma_spv_len;
		ci.pCode = (const uint32_t *)vk_plasma_spv;
		VKC(vkCreateShaderModule(dev, &ci, NULL, &shader));
		b.binding = 0;
		b.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
		b.descriptorCount = 1;
		b.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
		dci.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
		dci.bindingCount = 1;
		dci.pBindings = &b;
		VKC(vkCreateDescriptorSetLayout(dev, &dci, NULL, &dsl));
		pcr.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
		pcr.size = 16;
		plci.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
		plci.setLayoutCount = 1;
		plci.pSetLayouts = &dsl;
		plci.pushConstantRangeCount = 1;
		plci.pPushConstantRanges = &pcr;
		VKC(vkCreatePipelineLayout(dev, &plci, NULL, &pl));
		pci.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
		pci.stage.sType =
			VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
		pci.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
		pci.stage.module = shader;
		pci.stage.pName = "main";
		pci.layout = pl;
		VKC(vkCreateComputePipelines(dev, VK_NULL_HANDLE, 1, &pci,
					     NULL, &pipe));
	}
	{
		VkDescriptorPoolSize ps = {0};
		VkDescriptorPoolCreateInfo pci = {0};
		VkDescriptorSetAllocateInfo ai = {0};
		VkDescriptorImageInfo ii = {0};
		VkWriteDescriptorSet w = {0};

		ps.type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
		ps.descriptorCount = 1;
		pci.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
		pci.maxSets = 1;
		pci.poolSizeCount = 1;
		pci.pPoolSizes = &ps;
		VKC(vkCreateDescriptorPool(dev, &pci, NULL, &pool));
		ai.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
		ai.descriptorPool = pool;
		ai.descriptorSetCount = 1;
		ai.pSetLayouts = &dsl;
		VKC(vkAllocateDescriptorSets(dev, &ai, &dset));
		ii.imageView = view;
		ii.imageLayout = VK_IMAGE_LAYOUT_GENERAL;
		w.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		w.dstSet = dset;
		w.dstBinding = 0;
		w.descriptorCount = 1;
		w.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
		w.pImageInfo = &ii;
		vkUpdateDescriptorSets(dev, 1, &w, 0, NULL);
	}
	{
		VkCommandPoolCreateInfo ci = {0};

		ci.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
		ci.queueFamilyIndex = fam;
		VKC(vkCreateCommandPool(dev, &ci, NULL, &cpool));
		VkCommandBufferAllocateInfo ai = {0};

		ai.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
		ai.commandPool = cpool;
		ai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
		ai.commandBufferCount = 1;
		VKC(vkAllocateCommandBuffers(dev, &ai, &cb));
	}
	{
		VkFenceCreateInfo ci = {0};

		ci.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
		VKC(vkCreateFence(dev, &ci, NULL, &fence));
	}
	if (publish) {
		const char *bp =
			"/sys/bus/pci/devices/0000:00:03.0/resource2";

		bfd = open(bp, O_RDWR | O_SYNC);
		if (bfd >= 0) {
			bar = mmap(NULL, 128UL << 20, PROT_READ | PROT_WRITE,
				   MAP_SHARED, bfd, 0);
			if (bar != MAP_FAILED) {
				hdr = (struct frame_hdr *)(bar +
							   FRAME_OFFSET);
				seqp = &hdr->seq;
				hdr->magic = HDR_MAGIC;
				hdr->version = 1;
				hdr->width = W;
				hdr->height = H;
				hdr->stride = W * 4;
				hdr->format = 0;
				printf("PUBLISH on\n");
			} else {
				bar = NULL;
			}
		}
		if (!bar)
			printf("PUBLISH off (BAR open/mmap failed)\n");
	}
	VKC(vkMapMemory(dev, buf_mem, 0, imgsz, 0, (void **)&pix));

	for (i = 0; i < (uint32_t)frames; i++) {
		VkCommandBufferBeginInfo bi = {0};
		VkImageMemoryBarrier ib = {0};
		VkBufferImageCopy bc = {0};
		VkBufferMemoryBarrier bb = {0};
		float pc[4];
		uint64_t t0 = now_ns();
		VkSubmitInfo si = {0};
		uint32_t gw = (W + 15) / 16, gh = (H + 15) / 16;

		VKC(vkResetCommandBuffer(cb, 0));
		bi.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
		VKC(vkBeginCommandBuffer(cb, &bi));
		ib.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
		ib.srcAccessMask = 0;
		ib.dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
		ib.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
		ib.newLayout = VK_IMAGE_LAYOUT_GENERAL;
		ib.image = image;
		ib.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		ib.subresourceRange.levelCount = 1;
		ib.subresourceRange.layerCount = 1;
		vkCmdPipelineBarrier(cb,
				     VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
				     VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0,
				     NULL, 0, NULL, 1, &ib);
		vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe);
		vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pl,
					0, 1, &dset, 0, NULL);
		pc[0] = (float)i * 0.1f;
		pc[1] = (float)W;
		pc[2] = (float)H;
		pc[3] = 0.0f;
		vkCmdPushConstants(cb, pl, VK_SHADER_STAGE_COMPUTE_BIT, 0, 16,
				   pc);
		vkCmdDispatch(cb, gw, gh, 1);
		ib.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
		ib.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
		ib.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
		ib.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
		vkCmdPipelineBarrier(cb,
				     VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
				     VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, NULL,
				     0, NULL, 1, &ib);
		bc.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		bc.imageSubresource.layerCount = 1;
		bc.imageExtent.width = W;
		bc.imageExtent.height = H;
		bc.imageExtent.depth = 1;
		vkCmdCopyImageToBuffer(cb, image,
				       VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
				       buf, 1, &bc);
		bb.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
		bb.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
		bb.dstAccessMask = VK_ACCESS_HOST_READ_BIT;
		bb.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
		bb.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
		bb.buffer = buf;
		bb.size = VK_WHOLE_SIZE;
		vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT,
				     VK_PIPELINE_STAGE_HOST_BIT, 0, 0, NULL, 1,
				     &bb, 0, NULL);
		VKC(vkEndCommandBuffer(cb));
		si.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
		si.commandBufferCount = 1;
		si.pCommandBuffers = &cb;
		VKC(vkResetFences(dev, 1, &fence));
		VKC(vkQueueSubmit(queue, 1, &si, fence));
		VKC(vkWaitForFences(dev, 1, &fence, VK_TRUE, 30000000000ull));
		vkUnmapMemory(dev, buf_mem);
		vkMapMemory(dev, buf_mem, 0, imgsz, 0, (void **)&pix);

		if (outdir) {
			char path[512];
			FILE *f;

			snprintf(path, sizeof(path), "%s/vk-%04d.ppm",
				 outdir, i + 1);
			f = fopen(path, "wb");
			if (f) {
				unsigned char *d = malloc((size_t)W * H * 3);
				uint32_t p;

				for (p = 0; p < W * H; p++) {
					d[p * 3 + 0] = pix[p * 4 + 0];
					d[p * 3 + 1] = pix[p * 4 + 1];
					d[p * 3 + 2] = pix[p * 4 + 2];
				}
				fprintf(f, "P6\n%u %u\n255\n", W, H);
				fwrite(d, 1, (size_t)W * H * 3, f);
				fclose(f);
				free(d);
				saved++;
			}
		}
		if (bar) {
			memcpy(bar + FRAME_OFFSET + sizeof(struct frame_hdr),
			       pix, imgsz);
			hdr->timestamp_ns = now_ns();
			__sync_synchronize();
			*seqp = (uint32_t)(i + 1);
			__sync_synchronize();
		}
		if (!t_first)
			t_first = t0;
		t_last = t0;
		if (i < 5 || (i % 30) == 29)
			printf("VKFRAME 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("VKFRAME_RATE fps=%.2f frames=%d\n",
		       (double)(i - 1) * 1e9 / (double)(t_last - t_first), i);
	printf("VKFRAME_DONE frames=%d %ux%u saved=%d\n", i, W, H, saved);
	vkDeviceWaitIdle(dev);
	printf("VKFRAME_IDLE ok\n");
	return 0;
}