#include <stdlib.h>
#include <stdio.h>
#include <byteswap.h>
#include <inttypes.h>
#include "qoi.h"
#include "image.h"

#define MONS_QOI_RGB 3
#define MONS_QOI_RGBA 4
#define MONS_QOI_SRGB_LINEAR_ALPHA 0
#define MONS_QOI_LINEAR 1
#define MONS_QOI_HEADER_LEN 14

#define MONS_QOI_OP_RGB 0b11111110
#define MONS_QOI_OP_RGBA 0b11111111
#define MONS_QOI_OP_INDEX 0b00000000
#define MONS_QOI_OP_DIFF 0b01000000
#define MONS_QOI_OP_LUMA 0b10000000
#define MONS_QOI_OP_RUN 0b11000000

typedef struct mons_qoi_header {
    char magic[4];
    uint32_t width;
    uint32_t height;
    uint8_t channels;
    uint8_t colorspace;
} mons_qoi_header;

unsigned char pixel_index_hash(mons_pixel pixel) {
    return (pixel.r * 3 + pixel.g * 5 + pixel.b * 7 + pixel.a * 11) % 64;
}

mons_image mons_load_qoi(FILE *stream) {
    mons_qoi_header header;
    if (fread(&header, MONS_QOI_HEADER_LEN, 1, stream) != 1) {
        fprintf(stderr, "Error reading QOI header");
        return (mons_image){0};
    }
    header.width = __bswap_32(header.width);
    header.height = __bswap_32(header.height);
    mons_pixel previous = {0, 0, 0, 255};
    mons_pixel previously_seen[64] = {0};
    mons_image output = {
        .width = header.width,
        .height = header.height,
        .data = calloc(4, header.width * header.height),
    };
    for (size_t i = 0; i < header.width * header.height; i++) {
        unsigned char op;
        if (fread(&op, 1, 1, stream) != 1) {
            fprintf(stderr, "Failed reading operation byte!");
            return (mons_image){0};
        }
        switch (op) {
        case MONS_QOI_OP_RGB: {
            unsigned char rgb[3];
            if (fread(&rgb, 1, 3, stream) != 3) {
                fprintf(stderr, "Failed reading RGB data!");
                return (mons_image){0};
            }
            output.data[i] = (mons_pixel){
                .r = rgb[0],
                .g = rgb[1],
                .b = rgb[2],
                .a = previous.a,
            };
            break;
        }
        case MONS_QOI_OP_RGBA: {
            unsigned char rgba[4];
            if (fread(&rgba, 1, 4, stream) != 4) {
                fprintf(stderr, "Failed reading RGBA data!");
                return (mons_image){0};
            }
            output.data[i] = (mons_pixel){
                .r = rgba[0],
                .g = rgba[1],
                .b = rgba[2],
                .a = rgba[3],
            };
            break;
        }
        default:
            switch (op & 0b11000000) {
            case MONS_QOI_OP_INDEX: {
                unsigned char index = op & 0b00111111;
                output.data[i] = previously_seen[index];
                break;
            }
            case MONS_QOI_OP_DIFF: {
                char r_diff = ((op & 0b00110000) >> 4) - 2;
                char g_diff = ((op & 0b00001100) >> 2) - 2;
                char b_diff = (op & 0b00000011) - 2;
                output.data[i] = (mons_pixel){
                    .r = previous.r + r_diff,
                    .g = previous.g + g_diff,
                    .b = previous.b + b_diff,
                    .a = previous.a,
                };
                break;
            }
            case MONS_QOI_OP_LUMA: {
                int g_diff = (op & 0b00111111) - 32;
                unsigned char second_byte;
                if (fread(&second_byte, 1, 1, stream) != 1) {
                    fprintf(stderr,
                            "Failed to read second byte of luma operation!");
                    return (mons_image){0};
                }
                int r_diff = g_diff - 8 + (((second_byte & 0b11110000) >> 4));
                int b_diff = g_diff - 8 + (second_byte & 0b00001111);
                output.data[i] = (mons_pixel){
                    .r = previous.r + r_diff,
                    .g = previous.g + g_diff,
                    .b = previous.b + b_diff,
                    .a = previous.a,
                };
                break;
            }
            case MONS_QOI_OP_RUN: {
                unsigned char run = (op & 0b00111111) + 1;
                for (int j = 0; j < run; j++) {
                    output.data[i++] = previous;
                }
                i--;
                break;
            }
            default:
                break;
            }
            break;
        }
        previous = output.data[i];
        unsigned int hash = pixel_index_hash(output.data[i]);
        previously_seen[hash] = output.data[i];
    }
    fclose(stream);

    return output;
}