/* Copyright (C) 2024 Aiden Gall

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#include "camera.h"
#include "util.h"
#include "util_cl.h"

#include <CL/cl.h>
#include <CL/cl_platform.h>
#include <limits.h>
#include <raylib.h>
#include <stdio.h>
#include <stdlib.h>

#define SPIRV_START _binary_src_cl_spirt_spv_start
#define SPIRV_END _binary_src_cl_spirt_spv_end
#define SPIRV_SIZE ((size_t)SPIRV_END - (size_t)SPIRV_START)

struct kernel_context {
	cl_platform_id platform;
	cl_device_id device;
	cl_context context;
	cl_command_queue command_queue;
	cl_program program;
	cl_kernel k_gen_rays;
	cl_kernel k_ray_colour;
};

struct opts {
	long width;
	long height;
	long platidx;
	long devidx;
};

static float *kernel_render(struct kernel_context *runtime, size_t image_width,
                            size_t image_height, struct camera cam);

static struct kernel_context kernel_context_init(size_t platidx, size_t devidx);
static void kernel_context_clean(struct kernel_context *runtime);

static struct opts get_args(int argc, char *argv[]);

extern const char *SPIRV_START[];
extern const char *SPIRV_END[];

static float *
kernel_render(struct kernel_context *const runtime, const size_t image_width,
              const size_t image_height, const struct camera cam)
{
	cl_int err;

	size_t global_item_size[2];

	float *h_canvas;
	cl_mem d_canvas, d_rays;
	size_t rays_size, canvas_size;

	canvas_size = (sizeof("RGB") - 1) * image_width * image_height;
	rays_size = 2 * image_width * image_height;

	d_rays = clCreateBuffer(runtime->context, CL_MEM_READ_WRITE,
	                        sizeof(cl_float3) * rays_size, NULL, &err);
	if (err != CL_SUCCESS)
		die("clCreateBuffer: %s\n", cl_strerror(err));

	h_canvas = calloc(canvas_size, sizeof(*h_canvas));
	if (!h_canvas)
		die("calloc: Out of memory\n");

	d_canvas = clCreateBuffer(runtime->context, CL_MEM_WRITE_ONLY,
	                          sizeof(*h_canvas) * canvas_size, NULL, &err);
	if (err != CL_SUCCESS)
		die("clCreateBuffer: %s\n", cl_strerror(err));

	err = clSetKernelArg(runtime->k_gen_rays, 0, sizeof(d_rays), &d_rays);
	err |= clSetKernelArg(runtime->k_gen_rays, 1, sizeof(cl_float3),
	                      &cam.centre);
	err |= clSetKernelArg(runtime->k_gen_rays, 2, sizeof(cl_float3),
	                      &cam.pixel_delta_u);
	err |= clSetKernelArg(runtime->k_gen_rays, 3, sizeof(cl_float3),
	                      &cam.pixel_delta_v);
	err |= clSetKernelArg(runtime->k_gen_rays, 4, sizeof(cl_float3),
	                      &cam.corner00);
	if (err != CL_SUCCESS)
		die("clSetKernelArg: %s\n", cl_strerror(err));

	global_item_size[0] = image_width;
	global_item_size[1] = image_height;
	err = clEnqueueNDRangeKernel(runtime->command_queue,
	                             runtime->k_gen_rays, 2, NULL,
	                             global_item_size, NULL, 0, NULL, NULL);
	if (err != CL_SUCCESS)
		die("clEnqueueNDRangeKernel: %s\n", cl_strerror(err));

	err = clSetKernelArg(runtime->k_ray_colour, 0, sizeof(d_canvas),
	                     &d_canvas);
	err |= clSetKernelArg(runtime->k_ray_colour, 1, sizeof(d_rays),
	                      &d_rays);
	if (err != CL_SUCCESS)
		die("clSetKernelArg: %s\n", cl_strerror(err));

	global_item_size[0] = image_width;
	global_item_size[1] = image_height;
	err = clEnqueueNDRangeKernel(runtime->command_queue,
	                             runtime->k_ray_colour, 2, NULL,
	                             global_item_size, NULL, 0, NULL, NULL);
	if (err != CL_SUCCESS)
		die("clEnqueueNDRangeKernel: %s\n", cl_strerror(err));

	err = clEnqueueReadBuffer(runtime->command_queue, d_canvas, CL_TRUE, 0,
	                          sizeof(*h_canvas) * canvas_size, h_canvas, 0,
	                          NULL, NULL);
	if (err != CL_SUCCESS)
		die("clEnqueueReadBuffer: %s\n", cl_strerror(err));

	clReleaseMemObject(d_rays);
	clReleaseMemObject(d_canvas);

	return h_canvas;
}

static struct kernel_context
kernel_context_init(const size_t platidx, const size_t devidx)
{
	cl_int err;

	cl_platform_id *platforms;
	cl_device_id *devices;
	cl_uint num_platforms;
	cl_uint num_devices;

	struct kernel_context runtime;

	platforms = get_platforms(&num_platforms);
	if (platidx >= num_platforms)
		die("get_platforms: Index out of range\n");
	runtime.platform = platforms[platidx];
	free(platforms);

	devices = get_devices(runtime.platform, &num_devices);
	if (devidx >= num_devices)
		die("get_devices: Index out of range\n");
	runtime.device = devices[devidx];
	free(devices);

	runtime.context =
	        clCreateContext(NULL, 1, &runtime.device, NULL, NULL, &err);
	if (err != CL_SUCCESS)
		die("clCreateContext: %s\n", cl_strerror(err));

	runtime.command_queue = clCreateCommandQueueWithProperties(
	        runtime.context, runtime.device, NULL, &err);
	if (err != CL_SUCCESS)
		die("clCreateCommandQueueWithProperties: %s\n",
		    cl_strerror(err));

	runtime.program = compile_spirv_program(runtime.context, runtime.device,
	                                        SPIRV_START, SPIRV_SIZE);

	runtime.k_gen_rays = clCreateKernel(runtime.program, "gen_rays", &err);
	if (err != CL_SUCCESS)
		die("clCreateKernel: %s\n", cl_strerror(err));
	runtime.k_ray_colour =
	        clCreateKernel(runtime.program, "ray_colour", &err);
	if (err != CL_SUCCESS)
		die("clCreateKernel: %s\n", cl_strerror(err));

	return runtime;
}

static void
kernel_context_clean(struct kernel_context *const runtime)
{
	clFlush(runtime->command_queue);
	clFinish(runtime->command_queue);

	clReleaseKernel(runtime->k_gen_rays);
	clReleaseKernel(runtime->k_ray_colour);
	clReleaseProgram(runtime->program);
	clReleaseCommandQueue(runtime->command_queue);
	clReleaseContext(runtime->context);
}

static struct opts
get_args(int argc, char *argv[])
{
	struct opts opts;

	opts.platidx = 0;
	opts.devidx = 0;
	opts.width = 1024;
	opts.height = 576;

	if (argv[0])
		argv++, argc--;

	for (; argv[0] && argv[0][0] == '-' && argv[0][1]; argv++, argc--) {
		char flag;

		if (argv[0][2])
			die("get_args: Invalid flag '%s'\n", argv[0]);

		flag = argv[0][1];
		argv++, argc--;

		switch (flag) {
		case 'p':
			opts.platidx = argtonum(argv[0], 0, LONG_MAX);
			break;
		case 'd':
			opts.devidx = argtonum(argv[0], 0, LONG_MAX);
			break;
		case 'w':
			opts.width = argtonum(argv[0], 1, LONG_MAX);
			break;
		case 'h':
			opts.height = argtonum(argv[0], 1, LONG_MAX);
			break;
		case '-':
			goto end_flags;
		default:
			die("get_args: Unknown flag '-%c'\n", flag);
		}
	}

end_flags:
	if (argc > 0)
		die("get_args: Unexpected argument\n");

	return opts;
}

int
main(int argc, char *argv[])
{
	struct opts opts;

	struct kernel_context runtime;
	float *h_canvas;
	struct camera cam;

	Image img;
	Texture2D texture;

	opts = get_args(argc, argv);

	runtime = kernel_context_init(opts.platidx, opts.devidx);

	cam = camera_init(opts.width, opts.height, 1.0f);

	h_canvas = kernel_render(&runtime, opts.width, opts.height, cam);

	img.data = h_canvas;
	img.width = opts.width;
	img.height = opts.height;
	img.format = PIXELFORMAT_UNCOMPRESSED_R32G32B32;
	img.mipmaps = 1;

	InitWindow(opts.width, opts.height, "spirt");

	SetTargetFPS(60);

	texture = LoadTextureFromImage(img);

	while (!WindowShouldClose()) {
		BeginDrawing();
		ClearBackground(BLACK);
		DrawTexture(texture, 0, 0, WHITE);
		EndDrawing();
	}

	UnloadTexture(texture);
	CloseWindow();

	free(h_canvas);
	kernel_context_clean(&runtime);

	return 0;
}