/* 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 "farbfeld.h"
#include "util.h"

#include <CL/cl.h>
#include <CL/cl_platform.h>
#include <errno.h>
#include <limits.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

struct ClRuntime {
	cl_platform_id platform;
	cl_device_id device;
	cl_context context;
	cl_command_queue command_queue;
	cl_program program;
	cl_kernel k_hello;
};

static cl_platform_id *get_platforms(cl_uint *num_platforms);
static cl_device_id *get_devices(cl_platform_id platform, cl_uint *num_devices);
static cl_program compile_spirv_program(cl_context context, cl_device_id device,
                                        const void *spirv_start,
                                        size_t spirv_size);

static struct ClRuntime init(size_t platidx, size_t devidx);
static uint16_t *run_k_hello(struct ClRuntime *runtime, size_t *canvas_size_ret,
                             size_t image_width, size_t image_height);
static void clean(struct ClRuntime *runtime);

static long argtonum(const char *arg, long minval, long maxval);

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

static cl_platform_id *
get_platforms(cl_uint *const num_platforms)
{
	cl_int err;

	cl_platform_id *platforms;
	cl_uint len;

	err = clGetPlatformIDs(0, NULL, &len);
	if (err != CL_SUCCESS)
		die("clGetPlatformIDs: %s\n", cl_strerror(err));
	warn("number of platforms = %d\n", len);
	if (len < 1)
		die("clGetPlatformIDs: No OpenCL platforms\n");

	platforms = calloc(len, sizeof(*platforms));
	if (!platforms)
		die("calloc: Out of memory\n");
	err = clGetPlatformIDs(len, platforms, NULL);
	if (err != CL_SUCCESS)
		die("clGetPlatformIDs: %s\n", cl_strerror(err));

	*num_platforms = len;
	return platforms;
}

static cl_device_id *
get_devices(const cl_platform_id platform, cl_uint *const num_devices)
{
	cl_int err;

	cl_device_id *devices;
	cl_uint len;

	err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, 0, NULL, &len);
	if (err != CL_SUCCESS)
		die("clGetDeviceIDs: %s\n", cl_strerror(err));
	warn("number of devices in platform = %d\n", len);
	if (len < 1)
		die("clGetDeviceIDs: No OpenCL devices in platform\n");

	devices = calloc(len, sizeof(*devices));
	if (!devices)
		die("calloc: Out of memory\n");
	err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, len, devices, NULL);
	if (err != CL_SUCCESS)
		die("clGetDeviceIDs: %s\n", cl_strerror(err));

	*num_devices = len;
	return devices;
}

static cl_program
compile_spirv_program(const cl_context context, const cl_device_id device,
                      const void *const spirv_start, const size_t spirv_size)
{
	cl_int err, build_err;
	cl_program program;

	program = clCreateProgramWithIL(context, spirv_start, spirv_size, &err);
	if (err != CL_SUCCESS)
		die("clCreateProgramWithIL: %s\n", cl_strerror(err));

	build_err = clBuildProgram(program, 1, &device, NULL, NULL, NULL);
	if (build_err != CL_SUCCESS) {
		cl_char *log;
		size_t buffer_size;

		err = clGetProgramBuildInfo(program, device,
		                            CL_PROGRAM_BUILD_LOG, 0, NULL,
		                            &buffer_size);
		if (err != CL_SUCCESS)
			die("clGetProgramBuildInfo: %s\n", cl_strerror(err));
		if (buffer_size < 1)
			die("clGetProgramBuildInfo: "
			    "Build log buffer is empty\n");

		log = calloc(buffer_size, sizeof(*log));
		if (!log)
			die("calloc: Out of memory\n");
		err = clGetProgramBuildInfo(
		        program, device, CL_PROGRAM_BUILD_LOG,
		        sizeof(*log) * buffer_size, log, NULL);
		if (err != CL_SUCCESS)
			die("clGetProgramBuildInfo: %s\n", cl_strerror(err));

		efwrite(log, sizeof(*log), buffer_size, stderr);

		free(log);
		die("\nclBuildProgram: %s\n", cl_strerror(build_err));
	}

	return program;
}

static struct ClRuntime
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 ClRuntime 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_hello = clCreateKernel(runtime.program, "hello", &err);
	if (err != CL_SUCCESS)
		die("clCreateKernel: %s\n", cl_strerror(err));

	return runtime;
}

static uint16_t *
run_k_hello(struct ClRuntime *const runtime, size_t *const canvas_size_ret,
            const size_t image_width, const size_t image_height)
{
	cl_int err;

	size_t global_item_size[2];

	uint16_t *h_canvas;
	cl_mem d_canvas;
	size_t canvas_size;

	canvas_size = (sizeof("RGBA") - 1) * image_width * image_height;

	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_hello, 0, sizeof(d_canvas), &d_canvas);
	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_hello,
	                             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_canvas);

	*canvas_size_ret = canvas_size;
	return h_canvas;
}

static void
clean(struct ClRuntime *const runtime)
{
	clFlush(runtime->command_queue);
	clFinish(runtime->command_queue);

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

static long
argtonum(const char *const arg, const long minval, const long maxval)
{
	long ret;
	char *endptr;

	if (minval > maxval)
		die("argtonum: Invalid range\n");
	if (!arg)
		die("argtonum: Missing argument\n");

	errno = 0;
	ret = strtol(arg, &endptr, 10);
	if (errno)
		die("strtol: %s\n", strerror(errno));
	if (arg == endptr || *endptr)
		die("strtol: Invalid integer\n");

	if (ret < minval || ret > maxval)
		die("argtonum: Argument out of range\n");

	return ret;
}

int
main(int argc, char *argv[])
{
	long opt_platidx, opt_devidx, opt_width, opt_height;
	FILE *fd;

	struct ClRuntime runtime;

	uint16_t *h_canvas;
	size_t canvas_size;

	opt_platidx = 0;
	opt_devidx = 0;
	opt_width = 256;
	opt_height = 256;

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

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

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

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

		switch (flag) {
		case 'p':
			opt_platidx = argtonum(argv[0], 0, LONG_MAX);
			break;
		case 'd':
			opt_devidx = argtonum(argv[0], 0, LONG_MAX);
			break;
		case 'w':
			opt_width = argtonum(argv[0], 1, LONG_MAX);
			break;
		case 'h':
			opt_height = argtonum(argv[0], 1, LONG_MAX);
			break;
		case '-':
			goto end_flags;
		default:
			die("arg_parse: Unknown flag '-%c'\n", flag);
		}
	}

end_flags:
	if (argc > 1)
		die("arg_parse: Unexpected argument\n");

	fd = stdout;
	if (argv[0]) {
		errno = 0;
		fd = fopen(argv[0], "w");
		if (errno)
			die("fopen: %s\n", strerror(errno));
		if (!fd)
			die("fopen: Returned NULL\n");
	}

	runtime = init(opt_platidx, opt_devidx);

	h_canvas = run_k_hello(&runtime, &canvas_size, opt_width, opt_height);

	farbfeld_write(fd, h_canvas, canvas_size, opt_width, opt_height);
	fclose(fd);

	free(h_canvas);
	clean(&runtime);

	return 0;
}