#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <unistd.h>

#ifdef GLSLWRITE
# include <png.h>
#endif

#include <GL/glew.h>

#include <SDL2/SDL.h>
#include <SDL2/SDL_opengl.h>


#ifdef GLSLWRITE
static unsigned frame_time = 0;
#endif


static GLint time_loc;
static GLint res_loc;


static char * readfile(const char *filename) {
	FILE *f = fopen(filename, "r");
	if (!f)
		return NULL;

	size_t size = 1024;
	char *buffer = malloc(size+1);
	size_t count = 0, r;

	do {
		if (count == size) {
			size *= 2;
			buffer = realloc(buffer, size+1);
		}

		r = fread(buffer+count, 1, size-count, f);
		count += r;
	} while (r);

	fclose(f);

	buffer[count] = 0;

	return buffer;
}

static void printShaderInfoLog(GLuint obj) {
	GLint length = 0;

	glGetShaderiv(obj, GL_INFO_LOG_LENGTH, &length);

	if(length > 0) {
		char log[length];

		glGetShaderInfoLog(obj, sizeof(log), NULL, log);
		fprintf(stderr, "%s\n", log);
	}
}

static void printProgramInfoLog(GLuint obj) {
	GLint length = 0;

	glGetProgramiv(obj, GL_INFO_LOG_LENGTH, &length);

	if(length > 0) {
		char log[length];

		glGetProgramInfoLog(obj, sizeof(log), NULL, log);
		fprintf(stderr, "%s\n", log);
	}
}

static void init(char *filename) {
	glewInit();

	glDisable(GL_DEPTH_TEST);

	char *shader = readfile(filename);
	if (!shader) {
		fprintf(stderr, "Error: unable to read '%s'\n", filename);
		exit(1);
	}

	GLuint frag = glCreateShader(GL_FRAGMENT_SHADER);
	glShaderSource(frag, 1, (const GLchar **)&shader, NULL);
	glCompileShader(frag);
	printShaderInfoLog(frag);

	GLuint program = glCreateProgram();
	glAttachShader(program, frag);

	glLinkProgram(program);
	printProgramInfoLog(program);

	glBindFragDataLocation(program, 0, "fragColor");
	time_loc = glGetUniformLocation(program, "time");
	res_loc = glGetUniformLocation(program, "res");

	glUseProgram(program);
}

#ifdef GLSLWRITE
static void savePNG(int width, int height, const char *output_dir) {
	uint8_t pixels[width*height*4];
	glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_INT_8_8_8_8, pixels);

	char filename[strlen(output_dir) + 20];
	snprintf(filename, sizeof(filename), "%s/%010u.png", output_dir, frame_time);

	FILE *f = fopen(filename, "wb");
	if (!f) {
		fprintf(stderr, "unable to open PNG file\n");
		exit(1);
	}

	png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (!png_ptr) {
		fprintf(stderr, "unable to open PNG file\n");
		exit(1);
	}

	png_set_swap_alpha(png_ptr);

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr) {
		fprintf(stderr, "unable to create PNG info struct\n");
		exit(1);
	}

	if (setjmp(png_jmpbuf(png_ptr))) {
		fprintf(stderr, "unable to write PNG file\n");
		exit(1);
	}

	png_init_io(png_ptr, f);

	png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGB_ALPHA,
		     PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

	uint8_t *row_pointers[height];
	for (size_t i = 0; i < height; i++)
		row_pointers[i] = &pixels[4*i*width];

	png_set_rows(png_ptr, info_ptr, row_pointers);
	png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

	png_destroy_write_struct(&png_ptr, &info_ptr);
	fclose(f);
}
#endif

static void render(int width, int height) {
	glViewport(0, 0, (GLsizei)width, (GLsizei)height);

#ifdef GLSLWRITE
	glUniform1f(time_loc, frame_time);
#else
	glUniform1f(time_loc, SDL_GetTicks());
#endif
	glUniform2f(res_loc, width, height);

	glBegin(GL_QUADS);
	glVertex2f(-1, -1);
	glVertex2f(1, -1);
	glVertex2f(1, 1);
	glVertex2f(-1, 1);
	glEnd();

#ifdef GLSLWRITE
	frame_time += 20;
#endif
}


int main(int argc, char *argv[]) {
	bool running = true;

#ifdef GLSLWRITE
	if (argc != 3) {
		fprintf(stderr, "Usage: glslwrite <shader> <output directory>\n");
		exit(1);
	}
#else
	if (argc != 2) {
		fprintf(stderr, "Usage: glslview <shader> <output directory>\n");
		exit(1);
	}
#endif

	SDL_Init(SDL_INIT_VIDEO);

	SDL_Window *window = SDL_CreateWindow("glslwrite", 0, 0, 800, 800, SDL_WINDOW_OPENGL|SDL_WINDOW_RESIZABLE);
	SDL_GL_CreateContext(window);

	init(argv[1]);

	while (running) {
		SDL_Event e;

		while( SDL_PollEvent(&e)) {
			if(e.type == SDL_QUIT) {
				running = false;
				break;
			}
		}

		int width, height;
		SDL_GetWindowSize(window, &width, &height);

		render(width, height);
		SDL_GL_SwapWindow(window);

#ifdef GLSLWRITE
		savePNG(width, height, argv[2]);
#endif
	}

	return 0;
}