#include <assert.h>
#include <setjmp.h>
#include <stdio.h>
#include <stdlib.h>
#include <png.h>
#include <string.h>
#include <mazemaker.h>
#include "grid.h"
#include "options.h"

#define LINE_THICKNESS 3
#define BLOCK_SIZE 32
#define MARGIN 20

typedef struct {
    char* data;
    int w, h;
} img_data_t;

struct mem_encode {
    uint8_t *buffer;
    size_t size;
};

static void setPixel(img_data_t* img, int x, int y, char v) {
    //fprintf(stderr, "setPixel(img, %d, %d, %d)\n", x, y, v);
    assert((x < img->w) && (y < img->h));
    img->data[y * img->w + x] = v;
}

static void setBox(img_data_t* img, int x, int y, int w, int h, char v) {
    //fprintf(stderr, "setBox(img, %d, %d, %d, %d, %d)\n", x, y, w, h, v);
    for (int i = 0; i < h; i++) {
        for (int j = 0; j < w; j++) {
            setPixel(img, x + j, y + i, v);
        }
    }
}

static void setBoxCoords(img_data_t* img, int x1, int y1, int x2, int y2, char v) {
    if (x1 > x2) setBoxCoords(img, x2, x1, y1, y2, v);
    else if (y1 > y2) setBoxCoords(img, x1, x2, y2, y1, v);
    else setBox(img, x1, y1, (x2 - x1 + 1), (y2 - y1 + 1), v);
}

#if 0
static void printImageData(img_data_t const* img) {
    FILE *ptr = fopen("output.txt", "wt");
    for (int i = 0; i < img->h; i++) {
        int flip_i = img->h - 1 - i;
        for (int j = 0; j < img->w; j++) {
            char v = img->data[img->w * flip_i + j];
            if (v)
                fprintf(ptr, "#");
            else
                fprintf(ptr, ".");
        }
        fprintf(ptr, "\n");
    }
    fclose(ptr);
}
#endif // 0

static int _box_x(int x) {
    return MARGIN + LINE_THICKNESS + BLOCK_SIZE * x;
}

static int _box_y(int y, int height) {
    return height - MARGIN - LINE_THICKNESS - BLOCK_SIZE * y;
}

void drawPathSegment(size_t x, size_t y, size_t new_x, size_t new_y, img_data_t* result) {
    size_t img_height = result->h;
    if ((x > new_x) || (y > new_y)) {
        drawPathSegment(new_x, new_y, x, y, result);
        return;
    }

    // note: it is ALWAYS the case that either x == new_x OR y == new_y
    // so the above should not result in infinite recursion
    
    if (y == new_y) {
        // horizontal line (right)
        setBox(result, _box_x(x) + BLOCK_SIZE / 2 - LINE_THICKNESS / 2,
                       _box_y(y, img_height) - BLOCK_SIZE / 2 - LINE_THICKNESS / 2,
                       BLOCK_SIZE + LINE_THICKNESS,
                       LINE_THICKNESS,
                       1);
    } else {
        // vertical line (up)
        setBox(result, _box_x(x) + BLOCK_SIZE / 2 - LINE_THICKNESS / 2,
                       _box_y(new_y, img_height) - BLOCK_SIZE / 2 - LINE_THICKNESS / 2,
                       LINE_THICKNESS,
                       BLOCK_SIZE + LINE_THICKNESS,
                       1);
    }
}

static img_data_t* pathToImageData(maze_path_t const* path) {
    size_t width = path->nodes[path->length - 1].x + 1,
           height = path->nodes[path->length - 1].y + 1,
           x, y;
    img_data_t* result = malloc(sizeof(img_data_t));
    result->w  = LINE_THICKNESS + width * BLOCK_SIZE + 2 * MARGIN;
    result->h = LINE_THICKNESS + height * BLOCK_SIZE + 2 * MARGIN;
    result->data = malloc(result->w * result->h);
    //fprintf(stderr, "#########\n");

    // set background
    memset(result->data, 0, result->w * result->h);

    // set start and end points
    x = y = 0;
    setBox(result, MARGIN, _box_y(0, result->h) - BLOCK_SIZE / 2 - LINE_THICKNESS / 2, LINE_THICKNESS + BLOCK_SIZE / 2, LINE_THICKNESS, 1);
    setBox(result, _box_x(width - 1) + BLOCK_SIZE / 2 - LINE_THICKNESS / 2, _box_y(height - 1, result->h) - BLOCK_SIZE / 2 - LINE_THICKNESS / 2, BLOCK_SIZE / 2 + LINE_THICKNESS, LINE_THICKNESS, 1);

    for (int i = 1; i < path->length; i++) {
        size_t new_x = path->nodes[i].x, new_y = path->nodes[i].y;
        drawPathSegment(x, y, new_x, new_y, result);
        x = new_x;
        y = new_y;
    }

    //printImageData(result);
    return result;
}

static img_data_t* gridToImageData(mazegrid_t const* g) {
    img_data_t* result = malloc(sizeof(img_data_t));
    result->w  = LINE_THICKNESS + g->width  * BLOCK_SIZE + 2 * MARGIN;
    result->h = LINE_THICKNESS + g->height * BLOCK_SIZE + 2 * MARGIN;
    result->data = malloc(result->w * result->h);
    //fprintf(stderr, "Box dimensions: %d x %d\n", result->w, result->h);

    // set background
    memset(result->data, 0, result->w * result->h);

    // draw basic outline
    setBox(result, MARGIN, MARGIN, result->w - 2 * MARGIN, LINE_THICKNESS, 1);
    setBox(result, MARGIN, MARGIN + LINE_THICKNESS, LINE_THICKNESS, (g->height - 1) * BLOCK_SIZE, 1);
    setBox(result, MARGIN, result->h - MARGIN - 1 - LINE_THICKNESS, result->w - 2 * MARGIN, LINE_THICKNESS, 1);
    setBox(result, result->w - MARGIN - LINE_THICKNESS, MARGIN + LINE_THICKNESS + BLOCK_SIZE, LINE_THICKNESS, (g->height - 1) * BLOCK_SIZE, 1);

    // partitions
    for (int i = 0; i < g->height; i++) {
        for (int j = 0; j < g->width; j++) {
            //fprintf(stderr, "Drawing (%d, %d) -> %d, %d\n", j, i, mazegrid_get_edge(g, j, i, EDGE_UP), mazegrid_get_edge(g, j, i, EDGE_RIGHT));
            int x = MARGIN + LINE_THICKNESS + BLOCK_SIZE * j, y = result->h - MARGIN - LINE_THICKNESS - BLOCK_SIZE * i;
            if (i < g->height - 1) {
                if (mazegrid_get_edge(g, j, i, EDGE_UP) == 0)
                    setBox(result, x - LINE_THICKNESS, y - BLOCK_SIZE, BLOCK_SIZE + LINE_THICKNESS, LINE_THICKNESS, 1);
            }
            if (j < g->width - 1) {
                if (mazegrid_get_edge(g, j, i, EDGE_RIGHT) == 0)
                    setBox(result, x + BLOCK_SIZE - LINE_THICKNESS, y - BLOCK_SIZE, LINE_THICKNESS, BLOCK_SIZE, 1);
            }
        }
    }

    return result;
}

static void writeImageData(img_data_t const* img, png_structp png_ptr, mazeoptions_t const* options) {
    for (int i = 0; i < img->h; i++) {
        png_byte row[3 * img->w];
        for (int j = 0; j < img->w; j++) {
            if (img->data[i * img->w + j]) {
                // wall color
                row[3 * j] = options->wall_color[0];
                row[3 * j + 1] = options->wall_color[1];
                row[3 * j + 2] = options->wall_color[2];
            } else {
                // background
                row[3 * j] = options->background_color[0];
                row[3 * j + 1] = options->background_color[1];
                row[3 * j + 2] = options->background_color[2];
            }
        }
        png_write_row(png_ptr, row);
    }
}

static void writePathImageData(img_data_t const* img, png_structp png_ptr, mazeoptions_t const* options) {
    for (int i = 0; i < img->h; i++) {
        png_byte row[4 * img->w];
        for (int j = 0; j < img->w; j++) {
            if (img->data[i * img->w + j]) {
                // wall color
                row[4 * j] = options->path_color[0];
                row[4 * j + 1] = options->path_color[1];
                row[4 * j + 2] = options->path_color[2];
                row[4 * j + 3] = 0xff;
            } else memset(row + 4 * j, 0, 4);
        }
        png_write_row(png_ptr, row);
    }
}

static void freeImageData(img_data_t* img) {
    free(img->data);
    free(img);
}

static int
prepare_png(png_structp* png_ptr, png_infop* info_ptr) {
    *png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (*png_ptr == NULL) {
        fprintf(stderr, "Could not allocate PNG write struct.\n");
        return 0;
    }

    *info_ptr = png_create_info_struct(*png_ptr);
    if (info_ptr == NULL) {
        fprintf(stderr, "Could not allocate PNG info struct.\n");
        return 0;
    }

    if (setjmp(png_jmpbuf(*png_ptr))) {
        fprintf(stderr, "Error during PNG creation.\n");
        return 0;
    }

    return 1;
}

static void
cleanup_png(png_structp png_ptr, png_infop info_ptr) {
    if (info_ptr != NULL) {
        png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
        png_destroy_info_struct(png_ptr, &info_ptr);
    }
    if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
}

int mazemaker_path_to_png_opt(maze_path_t const* path, char const* filename, mazeoptions_t const* options) {
    int code = 1;
    img_data_t* img = pathToImageData(path);
    FILE *f = NULL;
    png_structp png_ptr = NULL;
    png_infop info_ptr  = NULL;
    
    if (strcmp(filename, "-") == 0) {
        f = stdout;
    } else {
        f = fopen(filename, "wb");
        if (NULL == f) {
            fprintf(stderr, "Could not open %s for writing.\n", filename);
            code = 0;
            goto exit;
        }
    }

    if (prepare_png(&png_ptr, &info_ptr) == 0) {
        code = 0;
        goto exit;
    }

    png_init_io(png_ptr, f);

    // Write header (8 bit color depth)
    png_set_IHDR(png_ptr, info_ptr, img->w, img->h,
                 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_write_info(png_ptr, info_ptr);

    writePathImageData(img, png_ptr, options);

    png_write_end(png_ptr, NULL);

exit:
    if (strcmp(filename, "-") != 0) {
        if (f != NULL) fclose(f);
    }
    cleanup_png(png_ptr, info_ptr);
    if (img != NULL) freeImageData(img);
    return code;
}

int mazemaker_maze_to_png_opt(mazegrid_t const* g, char const* filename, mazeoptions_t const* options) {
    int code = 1;
    img_data_t* img = gridToImageData(g);
    FILE *f = NULL;
    png_structp png_ptr = NULL;
    png_infop info_ptr  = NULL;
    
    if (strcmp(filename, "-") == 0) {
        f = stdout;
    } else {
        f = fopen(filename, "wb");
        if (NULL == f) {
            fprintf(stderr, "Could not open %s for writing.\n", filename);
            code = 0;
            goto exit;
        }
    }

    if (prepare_png(&png_ptr, &info_ptr) == 0) {
        code = 0;
        goto exit;
    }


    png_init_io(png_ptr, f);

    // Write header (8 bit color depth)
    png_set_IHDR(png_ptr, info_ptr, img->w, img->h,
                 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_write_info(png_ptr, info_ptr);

    writeImageData(img, png_ptr, options);

    png_write_end(png_ptr, NULL);

exit:
    if (strcmp(filename, "-") != 0) {
        if (f != NULL) fclose(f);
    }
    cleanup_png(png_ptr, info_ptr);
    if (img != NULL) freeImageData(img);
    return code;
}

int mazemaker_maze_to_png(mazegrid_t const* g, char const* filename) {
    mazeoptions_t* options = mazemaker_options_new(); // use default options
    int result = mazemaker_maze_to_png_opt(g, filename, options);
    mazemaker_options_free(options);
    return result;
}

int mazemaker_path_to_png(maze_path_t const* path, char const* filename) {
    mazeoptions_t* options = mazemaker_options_new(); // use default options
    int result = mazemaker_path_to_png_opt(path, filename, options);
    mazemaker_options_free(options);
    return result;
}

static void append_png_data(png_structp png_ptr, png_bytep data, png_size_t length) {
    struct mem_encode* p = (struct mem_encode*)png_get_io_ptr(png_ptr);
    size_t nsize = p->size + length;
    if (p->buffer) 
        p->buffer = realloc(p->buffer, nsize);
    else
        p->buffer = malloc(nsize);
    if (!p->buffer)
        png_error(png_ptr, "Write error");
    memcpy(p->buffer + p->size, data, length);
    p->size += length;
}

static void png_no_flush(png_structp png_ptr) {
    /* noop */
}

int mazemaker_maze_to_png_mem_opt(mazegrid_t const* g, size_t* len, uint8_t** buf, mazeoptions_t const* options) {
    int code = 1;
    img_data_t* img = gridToImageData(g);
    png_structp png_ptr = NULL;
    png_infop info_ptr  = NULL;
    struct mem_encode menc = { .buffer = NULL, .size = 0 };
    
    if (prepare_png(&png_ptr, &info_ptr) == 0) {
        code = 0;
        goto exit;
    }

    png_set_write_fn(png_ptr, &menc, append_png_data, png_no_flush);

    // Write header (8 bit color depth)
    png_set_IHDR(png_ptr, info_ptr, img->w, img->h,
                 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_write_info(png_ptr, info_ptr);

    writeImageData(img, png_ptr, options);

    png_write_end(png_ptr, NULL);

    *len = menc.size;
    *buf = menc.buffer;

exit:
    if (!code && menc.buffer) {
        free(menc.buffer);
    }
    cleanup_png(png_ptr, info_ptr);
    if (img != NULL) freeImageData(img);
    return code;
}

int mazemaker_maze_to_png_mem(mazegrid_t const* g, size_t* len, uint8_t** buf) {
    mazeoptions_t* options = mazemaker_options_new(); // use default options
    int result = mazemaker_maze_to_png_mem_opt(g, len, buf, options);
    mazemaker_options_free(options);
    return result;
}

int mazemaker_path_to_png_mem_opt(maze_path_t const* path, size_t* len, uint8_t** buf, mazeoptions_t const* options) {
    int code = 1;
    img_data_t* img = pathToImageData(path);
    png_structp png_ptr = NULL;
    png_infop info_ptr  = NULL;
    struct mem_encode menc = { .buffer = NULL, .size = 0 };
    
    if (prepare_png(&png_ptr, &info_ptr) == 0) {
        code = 0;
        goto exit;
    }

    png_set_write_fn(png_ptr, &menc, append_png_data, png_no_flush);

    // Write header (8 bit color depth)
    png_set_IHDR(png_ptr, info_ptr, img->w, img->h,
                 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_write_info(png_ptr, info_ptr);

    writePathImageData(img, png_ptr, options);

    png_write_end(png_ptr, NULL);

    *len = menc.size;
    *buf = menc.buffer;

exit:
    if (!code && menc.buffer) {
        free(menc.buffer);
    }
    cleanup_png(png_ptr, info_ptr);
    if (img != NULL) freeImageData(img);
    return code;
}

int mazemaker_path_to_png_mem(maze_path_t const* path, size_t* len, uint8_t** buf) {
    mazeoptions_t* options = mazemaker_options_new(); // use default options
    int result = mazemaker_path_to_png_mem_opt(path, len, buf, options);
    mazemaker_options_free(options);
    return result;
}