/* SPDX-License-Identifier: ISC */
/*
 * shadermeh.c
 *
 * Copyright (C) 2022 David Oberhollenzer <goliath@infraroot.at>
 */
#include "shadermeh.h"

static GLfloat vertex_buffer[] = {
	-1.0f, -1.0f, 0.0f,	/* lower left corner */
	 1.0f,  1.0f, 1.0f,

	+1.0f, -1.0f, 0.0f,	/* lower right corner */
	 1.0f,  1.0f, 1.0f,

	-1.0f,  1.0f, 0.0f,	/* uper left corner */
	 1.0f,  1.0f, 1.0f,

	+1.0f, +1.0f, 0.0f,	/* upper right corner */
	 1.0f,  1.0f, 1.0f,
};

static GLubyte audio_buffer[AUDIO_SAMPLES * AUDIO_CHANNELS];
size_t sample_pointer = 0;
size_t sample_data_pointer = 0;
size_t sample_rate = 0;
static float *audio_sample_data;
static float *audio_receive_data;
static double *fftw_in;
static fftw_complex *fftw_out;
static fftw_plan plan;

static int try_fetch_audio(float iTimeDelta)
{
	/* To avoid generating stale images, we keep our own sample buffer,
	 * which is then used to move a sliding window of data for the fft and
	 * wave samples. We need to do this, as otherwise we would set an upper
	 * limit of fps (20 at 4800kHz sample rate), which would not be good.
	 * The size of the window is set in the header file. The with our
	 * approach is that the buffer allows for drifting to occur within the
	 * buffer limits. If you buffer is 3s long the delay can grow to 3s.
	 * Choose your buffer size wisely for your application.
	 */
	size_t i;
	ssize_t ret = 0;
	memset(audio_receive_data, 0, AUDIO_BUFFER_SIZE * 
		sizeof(*audio_receive_data));
	sample_pointer += (sample_rate * iTimeDelta);

	for (;;) {
		ret = read(STDIN_FILENO, (char *)audio_receive_data,
			sizeof(*audio_receive_data)*AUDIO_BUFFER_SIZE);

		if (ret < 0) {
			if (errno == EINTR)
				continue;
			if (errno == EAGAIN || errno == EWOULDBLOCK)
				break;
			perror("stdin");
			return -1;
		}

		if (ret == 0 || ret % sizeof(float) != 0){
			break;
		}

		ret /= 4;
		if((ret + sample_pointer) > AUDIO_BUFFER_SIZE){
			/* Not enough storage space to store all new audio data,
			 * will override not output data with new one */
			 memset(audio_sample_data, 0, 
				AUDIO_BUFFER_SIZE * sizeof(*audio_sample_data));
			 memcpy(audio_sample_data, audio_receive_data, 
				ret * sizeof(*audio_sample_data));
			 sample_pointer = 0;
			 sample_data_pointer = ret;
		}else{
			memmove(audio_sample_data,
				&audio_sample_data[sample_pointer], 
				(AUDIO_BUFFER_SIZE - sample_pointer)*
				sizeof(*audio_sample_data));
			if(sample_data_pointer <= sample_pointer){
				sample_data_pointer = 0;

			}else{
				sample_data_pointer -= sample_pointer;

			}
			sample_pointer = 0;
			size_t len = ret;
			if((ret + sample_data_pointer) >= AUDIO_BUFFER_SIZE){
				len = AUDIO_BUFFER_SIZE - sample_data_pointer;
			}
			memcpy(&audio_sample_data[sample_data_pointer], 
				audio_receive_data, len * sizeof(float));
			sample_data_pointer += len;
			break;
		}

	}
	if((sample_pointer+AUDIO_FFT_SIZE) >= sample_data_pointer){
		fprintf(stderr, "shadermeh input to slow %zu > %zu! wrapping around!\n", sample_pointer+AUDIO_FFT_SIZE, sample_data_pointer);
		sample_pointer = 0;
	}

	memset(fftw_in, 0, sizeof(*fftw_in) * AUDIO_BUFFER_SIZE);
	memset(fftw_out, 0, sizeof(*fftw_out) * AUDIO_BUFFER_SIZE);

	for (i = 0; i < AUDIO_FFT_SIZE; ++i)
		fftw_in[i] = audio_sample_data[sample_pointer+i];

	fftw_execute(plan);

	for (i = 0; i < AUDIO_SAMPLES; ++i) {
		float a = cabs(fftw_out[i]);

		audio_buffer[i + AUDIO_SAMPLES] = audio_sample_data[sample_pointer+i] * 127.0f	+ 127.0f;
		audio_buffer[i] = log(fabsf(a)+1) * 50;
	}

	return 0;
}

static double diff_timespec(const struct timespec *time1,
			    const struct timespec *time0)
{
	return (time1->tv_sec - time0->tv_sec)
		+ (time1->tv_nsec - time0->tv_nsec) / 1000000000.0;
}

static void convert_for_ffmpeg(const uint8_t *in, uint8_t *out,
			       size_t width, size_t height)
{
	size_t x, y;

	for (y = 0; y < height; ++y) {
		const uint8_t *src = in + y * width * 4;
		uint8_t *dst = out + (height - 1 - y) * width * 3;

		for (x = 0; x < width; ++x) {
			*(dst++) = *(src++);
			*(dst++) = *(src++);
			*(dst++) = *(src++);
			++src;
		}
	}
}

static int wait_fd_event(int fd, int events)
{
	struct pollfd pfd;
	int ret;

	for (;;) {
		pfd.fd = fd;
		pfd.events = events;
		pfd.revents = 0;

		ret = poll(&pfd, 1, -1);

		if (ret > 0) {
			if (pfd.revents & events)
				break;

			if (pfd.revents & (POLLERR | POLLHUP)) {
				fputs("poll reported error\n", stderr);
				return -1;
			}
		}

		if (ret < 0) {
			if (errno == EINTR)
				continue;
			perror("poll");
			return -1;
		}
	}

	return 0;
}

static int write_retry(int fd, const void *buffer, size_t size)
{
	while (size > 0) {
		int ret = write(fd, buffer, size);

		if (ret < 0) {
			if (errno == EINTR)
				continue;

			if (errno == EAGAIN) {
				if (wait_fd_event(fd, POLLOUT))
					return -1;
				continue;
			}

			perror("write");
			return -1;
		}

		if (ret == 0)
			return -1;

		size -= ret;
		buffer = (const char *)buffer + ret;
	}

	return 0;
}

static const struct option long_opts[] = {
	{ "width", required_argument, NULL, 'w' },
 	{ "height", required_argument, NULL, 'h' },
	{ "shader", required_argument, NULL, 's' },
	{ "to-stdout", no_argument, NULL, 'S' },
	{ "stdin-audio", required_argument, NULL, 'a' },
	{ NULL, 0, NULL, 0 },
};

static const char *short_opts = "w:a:h:s:S";

static const char *usage_str =
"shadermeh OPTIONS...\n"
"\n"
"Possible options:\n"
"\n"
"  --width, -w <pixels>\n"
"  --height, -h <pixels>\n"
"\n"
"  --to-stdout, -S              Poop raw RGB24 frames to stdout (blocking)\n"
"  --stdin-audio, -a <sample rate>     Read raw PCM audio from stdin (non-blocking)\n"
"\n"
"  --shader, -s <shader file>\n"
"\n";

int main(int argc, char **argv)
{
	GLuint vao, vbo, fbo, fbo_tex, sound_tex, sampler_sound;
	GLuint u_iResolution, u_iTime, u_iTimeDelta, u_iFrame;
	struct timespec start, frame_start, frame_end;
	unsigned int width, height, iFrame = 0;
	void *fb32 = NULL, *fb24 = NULL;
	const char *shader_file = NULL;
	GLint major, minor, prog;
	float iTime, iTimeDelta = 0;
	bool have_audio = false;
	bool to_stdout = false;
	window *wnd;
	int i;

	/******************** parse options ************************/
	width = 800;
	height = 450;

	for (;;) {
		i = getopt_long(argc, argv, short_opts, long_opts, NULL);
		if (i == -1)
			break;

		switch (i) {
		case 'w':
			width = strtol(optarg, NULL, 10);
			break;
		case 'h':
			height = strtol(optarg, NULL, 10);
			break;
		case 's':
			shader_file = optarg;
			break;
		case 'S':
			to_stdout = true;
			break;
		case 'a':
			have_audio = true;
			sample_rate = strtol(optarg, NULL, 10);
			audio_sample_data = malloc(AUDIO_BUFFER_SIZE * 
				sizeof(float));
			audio_receive_data = malloc(AUDIO_BUFFER_SIZE * 
				sizeof(float));
			break;
		default:
			fputs(usage_str, stderr);
			return EXIT_FAILURE;
		}
	}

	if (have_audio) {
		int flags = fcntl(STDIN_FILENO, F_GETFL, 0);

		if (fcntl(STDIN_FILENO, F_SETFL, flags | O_NONBLOCK)) {
			perror("making stdin non-blocking");
			return EXIT_FAILURE;
		}
	}

	if (!shader_file) {
		fputs(usage_str, stderr);
		fputs("No shader file specified!\n", stderr);
		return EXIT_FAILURE;
	}

	if (to_stdout) {
		fb32 = calloc(1, width * height * 4);
		if (!fb32) {
			perror("allocating scratch framebuffer");
			return EXIT_FAILURE;
		}

		fb24 = calloc(1, width * height * 3);
		if (!fb24) {
			perror("allocating scratch framebuffer");
			free(fb32);
			return EXIT_FAILURE;
		}
	}

	/********** create window and make context current **********/
	wnd = window_create(width, height, "shader meh...");

	if (!wnd) {
		fputs("failed to create window", stderr);
		free(fb32);
		free(fb24);
		return EXIT_FAILURE;
	}

	window_make_current(wnd);
	window_set_vsync(wnd, 1);

	/******************** load entry points ********************/
	glewExperimental = GL_TRUE;

	if (glewInit() != GLEW_OK) {
		fputs("glewInit() error", stderr);
		window_destroy(wnd);
		free(fb32);
		free(fb24);
		return EXIT_FAILURE;
	}

	glGetIntegerv(GL_MAJOR_VERSION, &major);
	glGetIntegerv(GL_MINOR_VERSION, &minor);

	if (!to_stdout)
		printf("OpenGL version %d.%d\n", major, minor);

	/******************** initialization ********************/
	glViewport(0, 0, width, height);
	glClearColor(0.0f, 0.0f, 0.0f, 1.0f);

	/* vertex buffer object & vertex array object */
	glGenVertexArrays(1, &vao);
	glBindVertexArray(vao);

	glGenBuffers(1, &vbo);
	glBindBuffer(GL_ARRAY_BUFFER, vbo);

	glBufferData(GL_ARRAY_BUFFER, sizeof(vertex_buffer),
		     vertex_buffer, GL_STATIC_DRAW);

	glEnableVertexAttribArray(0);
	glEnableVertexAttribArray(1);

	glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, sizeof(GLfloat) * 6, 0);
	glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, sizeof(GLfloat) * 6,
			      (GLvoid *)(3 * sizeof(GLfloat)));

	/* shader */
	prog = shader_program_load(shader_file);

	if (!shader_program_get_build_status(prog)) {
		shader_program_print_info_log(prog);
		goto fail_vao;
	}

	glUseProgram(prog);

	/* uniforms */
	u_iResolution = glGetUniformLocation(prog, "iResolution");
	u_iTime = glGetUniformLocation(prog, "iTime");
	u_iTimeDelta = glGetUniformLocation(prog, "iTimeDelta");
	u_iFrame = glGetUniformLocation(prog, "iFrame");

	glUniform3f(u_iResolution, width, height, 0.0f);

	glUniform1i(glGetUniformLocation(prog, "iChannel0"), 0);

	/************************* textures *************************/
	glGenTextures(1, &sound_tex);

	glBindTexture(GL_TEXTURE_2D, sound_tex);
	glTexImage2D(GL_TEXTURE_2D, 0, GL_R8,
		     AUDIO_SAMPLES, AUDIO_CHANNELS, 0,
		     GL_RED, GL_UNSIGNED_BYTE, audio_buffer);
	glBindTexture(GL_TEXTURE_2D, 0);

	glGenSamplers(1, &sampler_sound);
	glSamplerParameteri(sampler_sound, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glSamplerParameteri(sampler_sound, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
	glBindSampler(0, sampler_sound);

	if (have_audio) {
		fftw_in = fftw_alloc_real(AUDIO_BUFFER_SIZE);
		fftw_out = fftw_alloc_complex(AUDIO_BUFFER_SIZE);
		if(fftw_in == NULL || fftw_out == NULL)
			goto fail_vao;
		plan = fftw_plan_dft_r2c_1d(AUDIO_BUFFER_SIZE, fftw_in, fftw_out,
			 FFTW_MEASURE);
	}

	/******************** framebuffer object ********************/
	if (to_stdout) {
		glGenFramebuffers(1, &fbo);
		glGenTextures(1, &fbo_tex);

		glBindTexture(GL_TEXTURE_2D, fbo_tex);
		glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, width, height, 0,
			     GL_RGBA, GL_UNSIGNED_BYTE, NULL);
		glBindTexture(GL_TEXTURE_2D, 0);

		glBindFramebuffer(GL_FRAMEBUFFER, fbo);
		glFramebufferTexture(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
				     fbo_tex, 0);

		glBindFramebuffer(GL_FRAMEBUFFER, 0);
	} else {
		window_show(wnd);
	}

	/******************** drawing loop ********************/
	clock_gettime(CLOCK_MONOTONIC_RAW, &start);

	while (to_stdout || window_handle_events()) {
		/* render image to FBO */
		clock_gettime(CLOCK_MONOTONIC_RAW, &frame_start);

		if (to_stdout)
			glBindFramebuffer(GL_FRAMEBUFFER, fbo);

		glClear(GL_COLOR_BUFFER_BIT);

		if (have_audio) {
			if (try_fetch_audio(iTimeDelta))
				break;

			glBindTexture(GL_TEXTURE_2D, sound_tex);

			glTexImage2D(GL_TEXTURE_2D, 0, GL_R8,
				     AUDIO_SAMPLES, AUDIO_CHANNELS, 0,
				     GL_RED, GL_UNSIGNED_BYTE, audio_buffer);
		}

		glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
		glBindTexture(GL_TEXTURE_2D, 0);

		if (to_stdout) {
			glFlush();
		} else {
			window_swap_buffers(wnd);
		}

		/* get image from FBO, dump to stdout */
		if (to_stdout) {
			glBindFramebuffer(GL_FRAMEBUFFER, 0);
			glBindTexture(GL_TEXTURE_2D, fbo_tex);

			glGetTexImage(GL_TEXTURE_2D, 0, GL_RGBA,
				      GL_UNSIGNED_BYTE, fb32);

			glBindTexture(GL_TEXTURE_2D, 0);

			convert_for_ffmpeg(fb32, fb24, width, height);

			if (write_retry(STDOUT_FILENO, fb24,
					width * height * 3)) {
				break;
			}
		}

		/* update timers */
		clock_gettime(CLOCK_MONOTONIC_RAW, &frame_end);

		iFrame += 1;
		iTimeDelta = diff_timespec(&frame_end, &frame_start);
		iTime = diff_timespec(&frame_end, &start);

		glUniform1f(u_iTimeDelta, iTimeDelta);
		glUniform1f(u_iTime, iTime);
		glUniform1ui(u_iFrame, iFrame);
	}

	/******************** cleanup ********************/
	glDeleteTextures(1, &sound_tex);
	glDeleteSamplers(1, &sampler_sound);

	if (to_stdout) {
		glBindFramebuffer(GL_FRAMEBUFFER, 0);
		glDeleteFramebuffers(1, &fbo);
		glDeleteTextures(1, &fbo_tex);
	}

	glUseProgram(0);
	glDeleteProgram(prog);

	if (have_audio) {
		fftw_destroy_plan(plan);
	}
fail_vao:
	glBindBuffer(GL_ARRAY_BUFFER, 0);
	glDeleteBuffers(1, &vbo);
	glBindVertexArray(0);
	glDeleteVertexArrays(1, &vao);

	window_make_current(NULL);
	free(fb32);
	free(fb24);
	fftw_free(fftw_in);
	fftw_free(fftw_out);
	window_destroy(wnd);
	return EXIT_SUCCESS;
}