#include <stdbool.h>
#include <dpa-qoi.h>

enum qoi_op {
  QOI_OP_INDEX,
  QOI_OP_DIFF,
  QOI_OP_LUMA,
  QOI_OP_RUN,
  QOI_OP_RGB,
  QOI_OP_RGBA,
};

static inline enum qoi_op qoi_get_op(const unsigned char symbol){
  if(symbol == 0xFE){
    return QOI_OP_RGB;
  }else if(symbol == 0xFF){
    return QOI_OP_RGBA;
  }
  return symbol >> 6;
}

static inline unsigned char qoi_color_to_index(const struct qoi_color_rgba c){
  return (c.color[0] * 3 + c.color[1] * 5 + c.color[2] * 7 + c.color[3] * 11) % 64;
}

static inline int qoi_output_store(unsigned char*restrict out, struct qoi_color_rgba c){
  // Feel free to change this as needed
  out[0] = c.color[2];
  out[1] = c.color[1];
  out[2] = c.color[0];
  out[3] = c.color[3];
  return 4;
}

void qoi_decode(struct qoi_decoder*restrict state, size_t count, unsigned char*restrict output){
  if(!count || !state->input)
    return;
  if(state->run)
    goto skip_store;
  do {
    state->run = 1;
    {
      const unsigned char ch = *state->input++;
      switch(qoi_get_op(ch)){
        case QOI_OP_INDEX: {
          if(!ch && !*state->input){
            // Failsafe to prevent buffer overflows
            state->input = 0;
            count = 1;
          }
          state->current_color = state->table[ch];
        } goto skip_store;
        case QOI_OP_DIFF: {
          state->current_color.color[0] += ((ch>>4) & 3) - 2;
          state->current_color.color[1] += ((ch>>2) & 3) - 2;
          state->current_color.color[2] += ( ch     & 3) - 2;
        } break;
        case QOI_OP_LUMA: {
          const unsigned char rb = *state->input++;
          int8_t dg = (ch & 0x3F) - 32;
          state->current_color.color[0] += (rb >>   4) - 8 + dg;
          state->current_color.color[1] += dg;
          state->current_color.color[2] += (rb & 0x0F) - 8 + dg;
        } break;
        case QOI_OP_RUN: {
          state->run = (ch & 0x3F) + 1;
        } goto skip_store;
        case QOI_OP_RGB: {
          state->current_color.color[0] = *state->input++;
          state->current_color.color[1] = *state->input++;
          state->current_color.color[2] = *state->input++;
        } break;
        case QOI_OP_RGBA: {
          state->current_color.color[0] = *state->input++;
          state->current_color.color[1] = *state->input++;
          state->current_color.color[2] = *state->input++;
          state->current_color.color[3] = *state->input++;
        } break;
      }
    }
    state->table[qoi_color_to_index(state->current_color)] = state->current_color;
    skip_store: {
      if(!output){
        if(state->run > count){
          state->run -= count;
          count = 0;
        }else{
          count -= state->run;
        }
        if(!count)
          break;
      }else{
        for(; count && state->run--; count--){
          output += qoi_output_store(output, state->current_color);
        }
      }
    }
  } while(count);
}