#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

#define STB_IMAGE_IMPLEMENTATION
#define STBI_FAILURE_USERMSG
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

typedef struct image_t {
    char* filename;
    unsigned char* image_data;
    int id, width, height, channels;
} image_t;

typedef struct tex_rect_t {
    image_t* img;
    int x,y,w,h;
} tex_rect_t;

int main(int argc, char* argv[]) {
    if (argc == 1) {
        fprintf(stderr, "Must provide image files to pack");
        exit(-1);
    }
    if (argc == 2) {
        fprintf(stderr, "Please provide at least 2 images to pack");
        exit(-1);
    }
    image_t* images = calloc(argc - 1, sizeof(image_t));
    if (images == NULL) {
        fprintf(stderr, "Failed to alloc image data array");
        exit(-1);
    }
    //-------------------------------------
    // First get the image info, make sure the file is an image and it exists
    //-------------------------------------
    int image_count = argc - 1;
    for (int i = 0; i < argc - 1; i++) {
        image_t* img = &images[i];
        img->filename = argv[i+1];

        int result = stbi_info(img->filename, &img->width, &img->height, &img->channels);
        if (result != 1) {
            char *errorMsg = "Cannot load image '%s': %s\n";
            fprintf(stderr, errorMsg, img->filename, stbi_failure_reason());
            exit(-1);
        }
    }
    //-------------------------------------
    // Now run the packing algorithm. This is super inefficient
    // texture packing but I just want to get it working first
    //-------------------------------------

    tex_rect_t* tex_rects = calloc(image_count, sizeof(tex_rect_t));
    int totalh = 0, maxw = 0;
    for (int i = 0; i < image_count; i++) {
        tex_rects[i] = (tex_rect_t){
            &images[i],
            0, // Put them all on the 1st column for now
            totalh,
            images[i].width,
            images[i].height,
        };
        totalh += images[i].height;
        if (images[i].width > maxw) {
            maxw = images[i].width;
        }
    }
    int max = totalh > maxw ? totalh : maxw;
    int exp = (int)log2(max) + 1;
    int size = (int)pow(2, exp);

    //-------------------------------------
    // Now load the data into memory, for now, let's just load all images in one go,
    // however, it would be more optimized if we loaded each image sequentially or
    // better yet in parallel so that we can write to the different parts of the
    // texture atlas concurrently
    //-------------------------------------

    for (int i = 0; i < image_count; i++) {
        int w,h,c;
        images[i].image_data = (unsigned char*)stbi_load(images[i].filename, &w, &h, &c, 4);
        if (images[i].image_data == NULL) {
            char *errorMsg = "Could not load data for image '%s': %s";
            fprintf(stderr, errorMsg, images[i].filename, stbi_failure_reason());
            // QUESTION: Should we clean up here?
            exit(-1);
        }
    }

    //-------------------------------------
    // Just write the bytes directly for now, we don't have any "rect" data, which
    // we need to make this work, since we're "hardcoding" the positions of the textures
    //-------------------------------------

    unsigned char* buf = calloc(size*size*4, sizeof(char));
    int channels = images[0].channels;
    for (int i = 0; i < image_count; i++) {
        tex_rect_t rect = tex_rects[i];
        int buf_row = rect.y;
        int buf_col = rect.x;
        int img_row = 0;
        int img_col = 0;
        while (img_row < rect.h) {
            int buf_idx = (buf_row * maxw + buf_col) * channels;
            int img_idx = (img_row * rect.w + img_col) * channels;
            for (int j = 0; j < channels; j++) {
                buf[buf_idx + j] = images[i].image_data[img_idx + j];
            }
            img_col++; buf_col++;
            if (img_col > rect.w) {
                img_col = buf_col = 0;
                img_row++; buf_row++;
            }
        }
    }

    stbi_write_png("test.png", max, max, 4, buf, maxw * 4);

    //-------------------------------------
    // Free them all, but eventually we would free after we finished 1
    //-------------------------------------
    for (int i = 0; i < image_count; i++) {
        stbi_image_free(images[i].image_data);
    }
    return 0;
}