#include <getopt.h>

#include "shorthand.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define TILE_SIZE 8
#define DEFAULT_PALETTE_SIZE 16
#define DEFAULT_BG_COLOR 0x0000

typedef u16 Color;

static inline Color
rgb15(u32 red, u32 green, u32 blue ) {
    return (blue << 10) | (green << 5) | red;
}

static inline
int quantize5(int color) {
    return color * 31 / 255;
}

// TODO: Assuming 4bpp tiles. The tile is an index to a palette.
// TODO: Write more documentation about this.
typedef u32 Tile[TILE_SIZE];

typedef struct Tiles {
    Tile *data;
    size_t size;
    size_t capacity;
} Tiles;

typedef struct Palette {
    Color *data;
    size_t size;
    size_t capacity;
} Palette;

typedef struct Image {
    unsigned char *data;
    int width;
    int height;
    int n_channels;
} Image;

u8
pack_1bb(u32 hex) {
    u32 low_nibble = hex & 0xFFFF;
    u32 high_nibble = (hex >> 16) & 0xFFFF;
    u32 low = (low_nibble & 0x1);
    low += 2 * ((low_nibble >> 4) & 0x1);
    low += 4 * ((low_nibble >> 8) & 0x1);
    low += 8 * ((low_nibble >> 12) & 0x1);
    u32 high = (high_nibble & 0x1);
    high += 2 * ((high_nibble >> 4) & 0x1);
    high += 4 * ((high_nibble >> 8) & 0x1);
    high += 8 * ((high_nibble >> 12) & 0x1);
    return (high << 4) | low;
}

void
export_c_file(Tiles *tiles, Palette *palette, FILE *file, bool pack_1bpp) {
    // Output tiles.
    if (pack_1bpp) {
        fprintf(file, "u32 packed_tiles[] = {\n    ");
        size_t counter = 0;
        for (size_t i = 0; i < tiles->capacity; ++i) {
            for (size_t j = 0; j < TILE_SIZE; j += 4, counter++) {
                if (counter == 4) {
                    counter = 0;
                    fprintf(file, "\n    ");
                }
                u32 tile_a = tiles->data[i][j];
                u32 tile_b = tiles->data[i][j + 1];
                u32 tile_c = tiles->data[i][j + 2];
                u32 tile_d = tiles->data[i][j + 3];
                u32 packed = pack_1bb(tile_d);
                packed |= pack_1bb(tile_c) << 8;
                packed |= pack_1bb(tile_b) << 16;
                packed |= pack_1bb(tile_a) << 24;
                fprintf(file, "0x%08x,", packed);
                if (counter < 3) {
                    fprintf(file, " ");
                }
            }
        }
        fprintf(file, "\n};\n");
        fprintf(file, "\n");
    } else {
        fprintf(file, "u32 tiles[%lu][%u] = {\n", tiles->capacity, TILE_SIZE);
        for (size_t i = 0; i < tiles->capacity; ++i) {
            fprintf(file, "    {");
            for (size_t j = 0; j < TILE_SIZE; ++j) {
                fprintf(file, "0x%08x", tiles->data[i][j]);
                if (j != (TILE_SIZE - 1)) {
                    fprintf(file, ", ");
                }
            }
            fprintf(file, "},\n");
        }
        fprintf(file, "};\n");
        fprintf(file, "\n");
    }

    // Output palette.
    fprintf(file, "u16 palette[%lu] = {\n", palette->capacity);
    fprintf(file, "    ");
    size_t counter = 0;
    for (size_t i = 0; i < palette->capacity; ++i, ++counter) {
            if (counter == 4) {
                counter = 0;
                fprintf(file, "\n    ");
            }
            fprintf(file, "0x%04x,", palette->data[i]);
            if (counter < 3) {
                fprintf(file, " ");
            }
    }
    fprintf(file, "\n");
    fprintf(file, "};\n");
}

void
extract_tile(Tiles *tiles, Palette *palette, Image *img,
        size_t offset_x, size_t offset_y) {
    for (size_t j = 0; j < TILE_SIZE; ++j) {
        u32 col_index = 0x00000000;
        for (size_t i = 0; i < TILE_SIZE; ++i) {
            size_t pal_index = palette->size;
            // Find the memory index for this pixel.
            int idx = (i + offset_x) + (j + offset_y) * img->width;
            idx *= img->n_channels;

            int red = img->data[idx];
            int green = img->data[idx + 1];
            int blue = img->data[idx + 2];

            // Quantize to 5 bits per channel.
            red = quantize5(red);
            green = quantize5(green);
            blue = quantize5(blue);

            Color clr = rgb15(red, green, blue);
            bool found = false;
            // Brute force search, but since the palette size will likely be
            // small, shouldn't be a big problem.
            if (palette->size < palette->capacity) {
                // If the palette is not full try to find if the color is
                // already stored.
                for (size_t p = 0; p < palette->capacity; ++p) {
                    if (clr == palette->data[p]) {
                        pal_index = p;
                        found = true;
                        break;
                    }
                }
            } else {
                // If the palette is full, find the closest perceived color
                // instead.
                size_t best_index = 0;
                double best_distance = 1e99;
                for (size_t p = 0; p < palette->capacity; ++p) {
                    Color pal_clr = palette->data[p];
                    int pal_red = pal_clr & 0x1F;
                    int pal_green = (pal_clr >> 5) & 0x1F;
                    int pal_blue = (pal_clr >> 10) & 0x1F;
                    double r = (pal_red - red);
                    double g = (pal_green - green);
                    double b = (pal_blue - blue);
                    double distance = r * r * 0.3 + g * g * 0.59 + b * b * 0.11;
                    if (distance < best_distance) {
                        best_distance = distance;
                        best_index = p;
                    }
                }
                pal_index = best_index;
                found = true;
            }
            if (!found) {
                palette->data[palette->size++] = clr;
            }
            col_index |= pal_index << i * 4; // TODO: FlipH?
        }
        tiles->data[tiles->size][j] = col_index;
    }
    tiles->size++;
}

void
print_usage(void) {
    printf("Usage: %s [options] <filename>\n", BIN_NAME);
    printf("\n");
    printf("\t-b <RRGGBB>\tSelects the background color used for transparency.\n");
    printf("\t-o <out_file.c>\tPath to the output file. If blank, stdout will be used.\n");
    printf("\t-x <0>\tStart x position (in pixels).\n");
    printf("\t-y <0>\tStart y position (in pixels).\n");
    printf("\t-W <8>\tSelect the sprite width in pixels. Must be a multiple of the tile size.\n");
    printf("\t-H <8>\tSelect the sprite height in pixels. Must be a multiple of the tile size.\n");
    printf("\t-1\tBit pack the output to 1bbp (4x size reduction).\n");
    printf("\t-w <0>\tWidth of region to extract from (x0,y0). Defaults to image width.\n");
    printf("\t-h <0>\tHeight of region to extract from (x0,y0). Defaults to image height.\n");
    printf("\n");
}

// TODO: Support multiple files at once?
int
main(int argc, char *argv[]) {
    // Initialize default parameters.
    u16 background_color = DEFAULT_BG_COLOR;
    u16 palette_size = DEFAULT_PALETTE_SIZE;
    size_t sprite_width = TILE_SIZE;
    size_t sprite_height = TILE_SIZE;
    char *out_file_path = NULL;
    bool pack_1bpp = false;
    size_t start_x = 0;
    size_t start_y = 0;
    size_t width = 0;
    size_t height = 0;

    int option;
    while ((option = getopt(argc, argv, "b:o:W:H:1x:y:w:h:")) != -1) {
        switch (option) {
            case 'b': {
                // Background color.
                char *bg = optarg;
                if (strlen(bg) != 6) {
                    fprintf(stderr, "%s: Invalid background color.\n", BIN_NAME);
                    print_usage();
                    return EXIT_FAILURE;
                }
                unsigned int red = 0;
                unsigned int green = 0;
                unsigned int blue = 0;

                // NOTE: Undefined behaviour, here we go! Also, no error checks,
                // so if an invalid number is given it will be parsed as 0.
                char r_str[3] = {bg[0], bg[1], '\0'};
                char g_str[3] = {bg[2], bg[3], '\0'};
                char b_str[3] = {bg[4], bg[5], '\0'};
                sscanf(r_str, "%x", &red);
                sscanf(g_str, "%x", &green);
                sscanf(b_str, "%x", &blue);

                // Quantize and transform to rgb15.
                red = quantize5(red);
                green = quantize5(green);
                blue = quantize5(blue);
                background_color = rgb15(red, green, blue);
            } break;
            case 'W': {
                // Sprite width.
                char *arg = optarg;
                sprite_width = 0;
                sprite_width = atoi(arg);

                // Check if a valid number is given.
                if (sprite_width == 0 || sprite_width % TILE_SIZE != 0) {
                    fprintf(stderr, "%s: Invalid sprite width.\n", BIN_NAME);
                    print_usage();
                    return EXIT_FAILURE;
                }
            } break;
            case 'H': {
                // Sprite height.
                char *arg = optarg;
                sprite_height = 0;
                sprite_height = atoi(arg);

                // Check if a valid number is given.
                if (sprite_height == 0 || sprite_height % TILE_SIZE != 0) {
                    fprintf(stderr, "%s: Invalid sprite height.\n", BIN_NAME);
                    print_usage();
                    return EXIT_FAILURE;
                }
            } break;
            case 'o': {
                // Output file.
                out_file_path = optarg;
            } break;
            case '1': {
                // Pack output to 1bpp.
                pack_1bpp = true;
            } break;
            case 'x': {
                // Start x position in pixels.
                start_x = atoi(optarg);
            } break;
            case 'y': {
                // Start y position in pixels.
                start_y = atoi(optarg);
            } break;
            case 'w': {
                // Width of region to extract from (x0,y0).
                width = atoi(optarg);
            } break;
            case 'h': {
                // Height of region to extract from (x0,y0).
                height = atoi(optarg);
            } break;
            default: {
                print_usage();
                return EXIT_FAILURE;
            } break;
        }
    }

    // Get the path to the file to be exported.
    if (optind != argc - 1) {
        fprintf(stderr, "%s: No filename given.\n", BIN_NAME);
        print_usage();
        return EXIT_FAILURE;
    }
    char *file_name = argv[optind];

    // Fill the palette with the background color if one was given.
    Palette palette = {
        .data = malloc(palette_size * sizeof(Color)),
        .size = 1,
        .capacity = palette_size,
    };
    for (size_t i = 0; i < palette_size; ++i) {
        palette.data[i] = background_color;
    }

    Image img = {0};

    // Open the given file.
    FILE *input_file = fopen(file_name, "rb");
    if (input_file == NULL) {
        fprintf(stderr, "%s: can't open input file: %s\n", BIN_NAME, file_name);
        return EXIT_FAILURE;
    }
    img.data = stbi_load_from_file(input_file, &img.width, &img.height, &img.n_channels, 0);
    fclose(input_file);

    // TODO: Implement support for different file inputs.
    if (img.n_channels != 3) {
        fprintf(stderr, "File format not supported. Only 3 channel files for now.\n");
        return EXIT_FAILURE;
    }

    if (width == 0) {
        width = img.width;
    }
    if (height == 0) {
        height = img.height;
    }
    size_t n_tiles_x = width / TILE_SIZE;
    size_t n_tiles_y = height / TILE_SIZE;
    size_t n_tiles = n_tiles_x * n_tiles_y;

    // Allocate memory for the tiles in this file, with zero-initialization.
    Tiles tiles = {
        .data = calloc(n_tiles, sizeof(Tile)),
        .size = 0,
        .capacity = n_tiles,
    };

    int offset_x = start_x;
    int offset_y = start_y;
    size_t n_sprites_x = width / sprite_width;
    size_t n_sprites_y = height / sprite_height;
    size_t n_sprite_tiles_x = sprite_width / TILE_SIZE;
    size_t n_sprite_tiles_y = sprite_height / TILE_SIZE;
    for (size_t p = 0; p < n_sprites_y; ++p) {
        for (size_t k = 0; k < n_sprites_x; ++k) {
            for (size_t j = 0; j < n_sprite_tiles_y; ++j) {
                int offset_x = start_x + k * sprite_width;
                int offset_y = start_y + p * sprite_height + j * TILE_SIZE;
                for (size_t i = 0; i < n_sprite_tiles_x; ++i) {
                    extract_tile(&tiles, &palette, &img, offset_x, offset_y);
                    offset_x += TILE_SIZE;
                }
            }
        }
    }

    if (out_file_path == NULL) {
        export_c_file(&tiles, &palette, stdout, pack_1bpp);
    } else {
        FILE *file = fopen(out_file_path, "w");
        if (!file) {
            fprintf(stderr, "%s: can't open output file: %s\n", BIN_NAME, out_file_path);
            return EXIT_FAILURE;
        }
        export_c_file(&tiles, &palette, file, pack_1bpp);
        fclose(file);
    }

    // Cleanup resources and exit.
    free(tiles.data);
    free(palette.data);
    stbi_image_free(img.data);
    return EXIT_SUCCESS;
}