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

#include <png.h>

#include <GL/glew.h>

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



static unsigned frame_time = 0;


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);
}

static void savePNG(uint8_t *pixels, int width, int height, const char *output_dir) {
        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);
}

static void render(SDL_Window *window, const char *output_dir) {
        int width, height;
        SDL_GetWindowSize(window, &width, &height);

	glViewport(0, 0, (GLsizei)width, (GLsizei)height);

        glUniform1f(time_loc, frame_time);
        glUniform2f(res_loc, width, height);

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

	SDL_GL_SwapWindow(window);

        uint8_t pixels[width*height*4];
        glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_INT_8_8_8_8, pixels);
        savePNG(pixels, width, height, output_dir);

        frame_time += 20;
}


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

        if (argc != 3) {
                fprintf(stderr, "Usage: glslwrite <shader> <output directory>\n");
                exit(1);
        }

	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;
			}
		}

		render(window, argv[2]);
                //usleep(1);
	}

	return 0;
}