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

#define MAX_DICT_SIZE 4096

namespace gif
{
    typedef struct {
        int prefix;
        uint8_t character;
    } DictEntry;

    typedef struct {
        uint8_t *data;
        size_t size;
        size_t capacity;
        int bit_pos;
        uint32_t bit_buffer;
    } BitStream;

    void bitstream_init(BitStream *bs) {
        bs->capacity = 256;
        bs->size = 0;
        bs->data = (uint8_t*)malloc(bs->capacity);
        bs->bit_pos = 0;
        bs->bit_buffer = 0;
    }

    void bitstream_write(BitStream *bs, uint16_t code, int code_size) {
        bs->bit_buffer |= ((uint32_t)code) << bs->bit_pos;
        bs->bit_pos += code_size;

        while (bs->bit_pos >= 8) {
            if (bs->size >= bs->capacity) {
                bs->capacity *= 2;
                bs->data = (uint8_t*)realloc(bs->data, bs->capacity);
            }
            bs->data[bs->size++] = bs->bit_buffer & 0xFF;
            bs->bit_buffer >>= 8;
            bs->bit_pos -= 8;
        }
    }

    void bitstream_flush(BitStream *bs) {
        if (bs->bit_pos > 0) {
            if (bs->size >= bs->capacity) {
                bs->capacity *= 2;
                bs->data = (uint8_t*)realloc(bs->data, bs->capacity);
            }
            bs->data[bs->size++] = bs->bit_buffer & 0xFF;
        }
    }

    uint8_t *lzw_compress(uint8_t *input, int width, int height, int min_code_size, size_t *out_size) {
        int clear_code = 1 << min_code_size;
        int end_code = clear_code + 1;
        int next_code = end_code + 1;
        int code_size = min_code_size + 1;

        DictEntry dict[MAX_DICT_SIZE];
        int dict_len = next_code;

        BitStream bs;
        bitstream_init(&bs);
        bitstream_write(&bs, clear_code, code_size);

        int prefix = input[0];
        for (int i = 1; i < width * height; i++) {
            uint8_t c = input[i];

            // Search for prefix + c in dictionary
            int found = -1;
            for (int j = end_code + 1; j < dict_len; j++) {
                if (dict[j].prefix == prefix && dict[j].character == c) {
                    found = j;
                    break;
                }
            }

            if (found != -1) {
                prefix = found; // Extend prefix
            } else {
                bitstream_write(&bs, prefix, code_size); // Emit current prefix
                if (dict_len < MAX_DICT_SIZE) {
                    if (dict_len == (1 << code_size) && code_size < 12) code_size++;
                    dict[dict_len].prefix = prefix;
                    dict[dict_len].character = c;
                    dict_len++;
                } else {
                    bitstream_write(&bs, clear_code, code_size);
                    dict_len = end_code + 1;
                    code_size = min_code_size + 1;
                }
                prefix = c; // Start new prefix
            }
        }

        // Emit final prefix and end code
        bitstream_write(&bs, prefix, code_size);
        bitstream_write(&bs, end_code, code_size);
        bitstream_flush(&bs);

        *out_size = bs.size;
        return bs.data;
    }

    void write_gif(const char *filename, uint8_t *pixels, int width, int height, uint8_t *palette, int palette_size) {
        FILE *f = fopen(filename, "wb");
        if (!f) {
            perror("Failed to open file");
            return;
        }

        // Header
        fwrite("GIF89a", 1, 6, f);

        // Determine min_code_size from palette_size
        int palette_depth = 0;
        while ((1 << palette_depth) < palette_size) palette_depth++;
        const int min_code_size = palette_depth < 2 ? 2 : palette_depth; // GIF spec requires at least 2

        printf("min_code_size: %i\n", palette_depth);

        // Logical Screen Descriptor
        uint8_t packed_field = palette ? (0x80 | ((palette_depth - 1) << 4) | (palette_depth - 1)) : 0x00;
        uint8_t screen_desc[] = {
            uint8_t(width & 0xFF), uint8_t((width >> 8) & 0xFF),
            uint8_t(height & 0xFF), uint8_t((height >> 8) & 0xFF),
            packed_field, // GCT flag + color resolution + size
            0x00,         // Background color index
            0x00          // Pixel aspect ratio
        };
        fwrite(screen_desc, 1, sizeof(screen_desc), f);

        // Global Color Table (if provided)
        if (palette) {
            int gct_size = 1 << palette_depth;
            fwrite(palette, 1, gct_size * 3, f);
        }

        // Image Descriptor
        uint8_t image_desc[] = {
            0x2C, 0, 0, 0, 0,
            uint8_t(width & 0xFF), uint8_t((width >> 8) & 0xFF),
            uint8_t(height & 0xFF), uint8_t((height >> 8) & 0xFF),
            0x00 // No local color table
        };
        fwrite(image_desc, 1, sizeof(image_desc), f);

        // LZW-compressed image data
        fwrite(&min_code_size, 1, 1, f);

        size_t compressed_size;
        uint8_t *compressed = lzw_compress(pixels, width, height, min_code_size, &compressed_size);

        // Write as sub-blocks
        size_t offset = 0;
        while (offset < compressed_size) {
            uint8_t block_size = (compressed_size - offset > 255) ? 255 : (uint8_t)(compressed_size - offset);
            fwrite(&block_size, 1, 1, f);
            fwrite(compressed + offset, 1, block_size, f);
            offset += block_size;
        }
        fputc(0x00, f); // Block terminator

        // Trailer
        fputc(0x3B, f);

        free(compressed);
        fclose(f);
    }
}