/*
 * tilesplit
 * Simple tool to split a PNG into a tileset + tilemap.
 */

#include <assert.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <png.h>


#ifdef _MSC_VER

#define FATAL_ERROR(format, ...)          \
do {                                      \
    fprintf(stderr, format, __VA_ARGS__); \
    exit(1);                              \
} while (0)

#else

#define FATAL_ERROR(format, ...)            \
do {                                        \
    fprintf(stderr, format, ##__VA_ARGS__); \
    exit(1);                                \
} while (0)

#endif // _MSC_VER


struct Color
{
    unsigned char r;
    unsigned char g;
    unsigned char b;
};

struct Image
{
    unsigned int width;
    unsigned int height;
    unsigned int colorType;
    unsigned int bitDepth;
    unsigned int rowBytes;
    unsigned char *pixels;
    unsigned int numColors;
    struct Color palette[256];
};

struct Tile
{
    unsigned int row;
    unsigned int col;
    unsigned int hash;
};

unsigned int tileWidth = 8;
unsigned int outputWidth = 8;
unsigned int mapBitsPerEntry = 8;

struct Image inputImage;

void *tileMap;
struct Tile *tiles;
unsigned int tilesCount;

void ReadPngPalette(struct Image *image, png_structp pngReader, png_infop pngInfo)
{
    png_colorp colors;
    int numColors;

    if (png_get_PLTE(pngReader, pngInfo, &colors, &numColors) != PNG_INFO_PLTE)
        FATAL_ERROR("Failed to retrieve palette\n");

    if (numColors > 256)
        FATAL_ERROR("Images with more than 256 colors are not supported.\n");

    image->numColors = numColors;
    for (int i = 0; i < numColors; i++)
    {
        image->palette[i].r = colors[i].red;
        image->palette[i].g = colors[i].green;
        image->palette[i].b = colors[i].blue;
    }
}

void WritePngPalette(struct Image *image, png_structp pngReader, png_infop pngInfo)
{
    png_colorp colors = malloc(image->numColors * sizeof(*colors));

    if (colors == NULL)
        FATAL_ERROR("Failed to allocate PNG palette.\n");

    for (int i = 0; i < image->numColors; i++)
    {
        colors[i].red = image->palette[i].r;
        colors[i].green = image->palette[i].g;
        colors[i].blue = image->palette[i].b;
    }

    png_set_PLTE(pngReader, pngInfo, colors, image->numColors);

    free(colors);
}

static void ReadImageFromFile(struct Image *image, const char *path)
{
    unsigned char sig[8];
    png_structp pngReader;
    png_infop pngInfo;
    FILE *file;
    png_byte **rowPointers;
    unsigned int i;

    file = fopen(path, "rb");
    if (file == NULL)
        FATAL_ERROR("Failed to open \"%s\" for reading.\n", path);

    if (fread(sig, 8, 1, file) != 1)
        FATAL_ERROR("Failed to read PNG signature from \"%s\".\n", path);

    if (png_sig_cmp(sig, 0, 8))
        FATAL_ERROR("\"%s\" does not have a valid PNG signature.\n", path);

    pngReader = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (pngReader == NULL)
        FATAL_ERROR("Failed to create PNG read struct.\n");

    pngInfo = png_create_info_struct(pngReader);
    if (pngInfo == NULL)
        FATAL_ERROR("Failed to create PNG info struct.\n");

    if (setjmp(png_jmpbuf(pngReader)))
        FATAL_ERROR("Failed to init I/O for reading \"%s\".\n", path);

    png_init_io(pngReader, file);
    png_set_sig_bytes(pngReader, 8);
    png_read_info(pngReader, pngInfo);

    // Get info
    image->width = png_get_image_width(pngReader, pngInfo);
    image->height = png_get_image_height(pngReader, pngInfo);
    image->colorType = png_get_color_type(pngReader, pngInfo);
    image->bitDepth = png_get_bit_depth(pngReader, pngInfo);

    if (image->colorType != PNG_COLOR_TYPE_PALETTE && image->colorType != PNG_COLOR_TYPE_GRAY)
        FATAL_ERROR("error: only paletted and grayscale images are supported");

    printf("tile width: %u\n", tileWidth);
    printf("image dimensions: %ux%u\n", image->width, image->height);
    printf("image bit depth: %u\n", image->bitDepth);

    // Read palette
    if (image->colorType == PNG_COLOR_TYPE_PALETTE)
        ReadPngPalette(image, pngReader, pngInfo);

    // Read pixels
    image->rowBytes = png_get_rowbytes(pngReader, pngInfo);
    image->pixels = malloc(image->height * image->rowBytes);
    rowPointers = malloc(image->height * sizeof(*rowPointers));

    for (i = 0; i < image->height; i++)
        rowPointers[i] = image->pixels + (i * image->rowBytes);

    if (setjmp(png_jmpbuf(pngReader)))
        FATAL_ERROR("Error reading image\n");

    png_read_image(pngReader, rowPointers);

    free(rowPointers);

    fclose(file);
}

static void WriteImageToFile(struct Image *image, const char *path)
{
    FILE *fp = fopen(path, "wb");

    if (fp == NULL)
        FATAL_ERROR("Failed to open \"%s\" for writing.\n", path);

    png_structp pngReader = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);

    if (!pngReader)
        FATAL_ERROR("Failed to create PNG write struct.\n");

    png_infop pngInfo = png_create_info_struct(pngReader);

    if (!pngInfo)
        FATAL_ERROR("Failed to create PNG info struct.\n");

    if (setjmp(png_jmpbuf(pngReader)))
        FATAL_ERROR("Failed to init I/O for writing \"%s\".\n", path);

    png_init_io(pngReader, fp);

    if (setjmp(png_jmpbuf(pngReader)))
        FATAL_ERROR("Error writing header for \"%s\".\n", path);

    png_set_IHDR(pngReader, pngInfo, image->width, image->height,
        image->bitDepth, image->colorType, PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    if (image->colorType == PNG_COLOR_TYPE_PALETTE)
        WritePngPalette(image, pngReader, pngInfo);

    png_write_info(pngReader, pngInfo);

    png_bytepp row_pointers = malloc(image->height * sizeof(png_bytep));

    if (row_pointers == NULL)
        FATAL_ERROR("Failed to allocate row pointers.\n");

    int rowbytes = png_get_rowbytes(pngReader, pngInfo);

    for (int i = 0; i < image->height; i++)
        row_pointers[i] = (png_bytep)(image->pixels + (i * rowbytes));

    if (setjmp(png_jmpbuf(pngReader)))
        FATAL_ERROR("Error writing \"%s\".\n", path);

    png_write_image(pngReader, row_pointers);

    if (setjmp(png_jmpbuf(pngReader)))
        FATAL_ERROR("Error ending write of \"%s\".\n", path);

    png_write_end(pngReader, NULL);

    fclose(fp);

    png_destroy_write_struct(&pngReader, &pngInfo);
    free(row_pointers);
}

static unsigned char *GetTileStartAddr(struct Image *image, unsigned int row, unsigned int col)
{
    unsigned int x = col * tileWidth;
    unsigned int y = row * tileWidth;

    return image->pixels + (x + y * image->width) * image->bitDepth / 8;
}

static void CreateTile(struct Tile *tile, unsigned int row, unsigned int col)
{
    const size_t srcRowSize = inputImage.width * inputImage.bitDepth / 8;
    const size_t tileRowSize = tileWidth * inputImage.bitDepth / 8;
    unsigned int i, j;
    unsigned char *src = GetTileStartAddr(&inputImage, row, col);
    unsigned int hash = 0;

    for (i = 0; i < tileWidth; i++)
    {
        for (j = 0; j < tileRowSize; j++)
            hash += src[j];
        src += srcRowSize;
    }
    tile->row = row;
    tile->col = col;
    tile->hash = hash;
}

static int FindTile(struct Tile *tile)
{
    unsigned int currTile;

    for (currTile = 0; currTile < tilesCount; currTile++)
    {
        if (tile->hash == tiles[currTile].hash)
        {
            const size_t imageRowSize = inputImage.width * inputImage.bitDepth / 8;
            const size_t tileRowSize = tileWidth * inputImage.bitDepth / 8;
            unsigned int i;
            unsigned char *src = inputImage.pixels + tile->row * tileWidth * imageRowSize + tile->col * tileRowSize;
            unsigned char *dest = inputImage.pixels + tiles[currTile].row * tileWidth * imageRowSize + tiles[currTile].col * tileRowSize;

            for (i = 0; i < tileWidth; i++)
            {
                if (memcmp(src, dest, tileRowSize) != 0)
                    goto check_next;
                src += imageRowSize;
                dest += imageRowSize;
            }
            return currTile;
        }
      check_next:
        ;
    }
    return -1;
}

void ReadImageTiles(const char *fname)
{
    unsigned int row, col;
    unsigned int maxEntry = (1 << mapBitsPerEntry) - 1;

    ReadImageFromFile(&inputImage, fname);

    unsigned int numTilesX = inputImage.width / tileWidth;
    unsigned int numTilesY = inputImage.height / tileWidth;

    tileMap = malloc(numTilesX * numTilesY * mapBitsPerEntry / 8);
    tiles = malloc(sizeof(*tiles) * numTilesX * numTilesY);

    tilesCount = 0;

    for (row = 0; row < numTilesY; row++)
    {
        for (col = 0; col < numTilesX; col++)
        {
            struct Tile tile;

            CreateTile(&tile, row, col);
            int tileNum = FindTile(&tile);

            if (tileNum == -1)  // this is a new tile
            {
                tileNum = tilesCount;
                tiles[tilesCount++] = tile;
            }
            if (tileNum > maxEntry)
                FATAL_ERROR("Image has more than %u tiles\n", maxEntry);
            if (mapBitsPerEntry == 8)
                ((uint8_t *)tileMap)[row * numTilesX + col] = tileNum;
            else if (mapBitsPerEntry == 16)
                ((uint16_t *)tileMap)[row * numTilesX + col] = tileNum;
            else
                assert(0);
        }
    }
}

void OutputTileSet(char *fname)
{
    unsigned int i, j;
    const unsigned int outputTilesX = outputWidth;
    const unsigned int outputTilesY = (tilesCount + outputTilesX - 1) / outputTilesX;
    struct Image outputImage = inputImage;

    outputImage.width = outputTilesX * tileWidth;
    outputImage.height = outputTilesY * tileWidth;
    outputImage.rowBytes = outputImage.width * outputImage.bitDepth / 8;
    outputImage.pixels = malloc(outputImage.rowBytes * outputImage.height);

    for (i = 0; i < tilesCount; i++)
    {
        unsigned int row = i / outputTilesX;
        unsigned int col = i % outputTilesX;
        unsigned char *src = GetTileStartAddr(&inputImage, tiles[i].row, tiles[i].col);
        unsigned char *dest = GetTileStartAddr(&outputImage, row, col);

        for (j = 0; j < tileWidth; j++)
        {
            memcpy(dest, src, tileWidth * outputImage.bitDepth / 8);
            src += inputImage.rowBytes;
            dest += outputImage.rowBytes;
        }
    }

    WriteImageToFile(&outputImage, fname);
}

void OutputTileMap(char *fname)
{
    FILE *file = fopen(fname, "wb");
    unsigned int numEntries = (inputImage.width / tileWidth) * (inputImage.height / tileWidth);
    size_t size = numEntries * mapBitsPerEntry / 8;

    fwrite(tileMap, size, 1, file);
    fclose(file);
}

char *GetFileNameExtension(char *fname)
{
    char *ext = strrchr(fname, '.');

    if (ext != NULL)
        return ext;
    else
        return fname + strlen(fname);
}

int main(int argc, char **argv)
{
    int i;
    char *fname = NULL;
    char *outfname;
    char *ext;

    for (i = 1; i < argc; i++)
    {
        char *arg = argv[i];

        if (strcmp("--tile_width", arg) == 0)
        {
            i++;
            if (i >= argc || sscanf(argv[i], "%u", &tileWidth) != 1)
                FATAL_ERROR("error: expected number after --tile_width\n");
        }
        else if (strcmp("--map_bits", arg) == 0)
        {
            i++;
            if (i >= argc || sscanf(argv[i], "%u", &mapBitsPerEntry) != 1)
                FATAL_ERROR("error: expected number after --map_bits\n");
        }
        else if (strcmp("--output_width", arg) == 0)
        {
            i++;
            if (i >= argc || sscanf(argv[i], "%u", &outputWidth) != 1)
                FATAL_ERROR("error: expected number after --output_width\n");
        }
        else
        {
            fname = arg;
        }
    }
    if (fname == NULL)
        FATAL_ERROR("error: no input file specified\n");

    outfname = malloc(strlen(fname) + 11);
    strcpy(outfname, fname);
    ext = GetFileNameExtension(outfname);

    ReadImageTiles(fname);

    strcpy(ext, "_tiles.png");
    OutputTileSet(outfname);

    strcpy(ext, "_map.bin");
    OutputTileMap(outfname);

    return 0;
}