/*
 * Copyright 2022-2023 Frederico de Oliveira Linhares
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "qoi.hpp"

#include <array>
#include <fstream>

namespace
{

const char MAGIC[]{"qoif"};
const int HEADER_SIZE{14};
const uint8_t PADDING[8]{0,0,0,0,0,0,0,1};
const unsigned int PIXELS_MAX{400000000};

const int OP_INDEX{0b00000000};
const int OP_DIFF {0b01000000};
const int OP_LUMA {0b10000000};
const int OP_RUN  {0b11000000};
const int OP_RGB  {0b11111110};
const int OP_RGBA {0b11111111};

const int MASK_2_BITS{0b11000000};

struct RGBA
{
  uint8_t red, green, blue, alpha;

  RGBA();
  RGBA(uint8_t red, uint8_t green, uint8_t blue, uint8_t alpha);
};

RGBA::RGBA():
  red{0},
  green{0},
  blue{0},
  alpha{255}
{
}

RGBA::RGBA(uint8_t red, uint8_t green, uint8_t blue, uint8_t alpha):
  red{red},
  green{green},
  blue{blue},
  alpha{alpha}
{
}

struct Pixel
{
  union
  {
    RGBA colors;
    uint32_t value;
  };

  Pixel();
  Pixel(uint8_t red, uint8_t green, uint8_t blue, uint8_t alpha);
};

Pixel::Pixel():
  colors{}
{
}

Pixel::Pixel(uint8_t red, uint8_t green, uint8_t blue, uint8_t alpha):
  colors(red, green, blue, alpha)
{
}

inline int
color_hash(const RGBA &colors)
{
  return colors.red*3 + colors.green*5 + colors.blue*7 + colors.alpha*11;
}

void
read_chars(uint8_t *data, int &p, char *str, int size)
{
  for(int i{0}; i < size; i++) str[i] = (char)data[p++];
}

void
read_ui8(uint8_t *data, int &p, uint8_t &num)
{
  num = data[p++];
}

void
read_ui32(uint8_t *data, int &p, uint32_t &num)
{
  uint8_t b1{data[p++]}, b2{data[p++]}, b3{data[p++]}, b4{data[p++]};

  num = b1 << 24 | b2 << 16 | b3 << 8 | b4;
}

}

namespace VK::QOI
{

Image::Image(const char *file_path, uint8_t channels):
  channels{channels}
{
  int data_p{0};
  int data_size;
  uint8_t *data;

  if(this->channels != 0 && this->channels != 3 && this->channels != 4)
  {
    throw std::invalid_argument{"invalid number of channels"};
  }

  { // Read data from file.
    std::ifstream file(file_path, std::ios::binary | std::ios::ate);
    if(!file)
    {
      throw std::runtime_error{"failed to open QOI file"};
    }

    data_size = file.tellg();
    if(data_size < HEADER_SIZE + (int)sizeof(PADDING))
    {
      throw std::runtime_error{"invalid QOI file"};
    }

    file.seekg(0);
    data = new uint8_t[data_size];
    file.read((char*)data, data_size);
  }

  read_chars(data, data_p, this->header.magic, 4);
  read_ui32(data, data_p, this->header.width);
  read_ui32(data, data_p, this->header.height);
  read_ui8(data, data_p, this->header.channels);
  read_ui8(data, data_p, this->header.colorspace);

  if(this->header.width == 0 || this->header.height == 0 ||
     this->header.channels < 3 || this->header.channels > 4 ||
     this->header.colorspace > 1 ||
     this->header.height >= PIXELS_MAX / this->header.width)
  {
    throw std::runtime_error{"QOI file have an invalid header"};
  }

  if(this->channels == 0) this->channels = this->header.channels;

  uint32_t num_pixels{this->header.width * this->header.height};
  this->pixels_len = num_pixels * this->channels;
  this->pixels = new uint8_t[this->pixels_len];

  std::array<Pixel, 64> index;
  Pixel pixel{};
  int chunks_len = data_size - (int)sizeof(PADDING);

  /*
    This algorithm is based on the original implementation that is in GitHub:
    https://github.com/phoboslab/qoi

    For information about the QOI image format, see: https://qoiformat.org/
   */ 
  int pixel_p{0};
  int run{0};
  for(uint32_t decoded_pixel{0}; decoded_pixel < num_pixels; decoded_pixel++)
  {
    if(run > 0)
    {
      run--;
    }
    else if(data_p < chunks_len)
    {
      uint8_t output;
      int b1;
      read_ui8(data, data_p, output);
      b1 = output;

      if (b1 == OP_RGB)
      {
        read_ui8(data, data_p, pixel.colors.red);
        read_ui8(data, data_p, pixel.colors.green);
        read_ui8(data, data_p, pixel.colors.blue);
      }
      else if (b1 == OP_RGBA)
      {
        read_ui8(data, data_p, pixel.colors.red);
        read_ui8(data, data_p, pixel.colors.green);
        read_ui8(data, data_p, pixel.colors.blue);
        read_ui8(data, data_p, pixel.colors.alpha);
      }
      else if ((b1 & MASK_2_BITS) == OP_INDEX)
      {
        pixel = index[b1];
      }
      else if ((b1 & MASK_2_BITS) == OP_DIFF)
      {
        pixel.colors.red += ((b1 >> 4) & 0x03) - 2;
        pixel.colors.green += ((b1 >> 2) & 0x03) - 2;
        pixel.colors.blue += (b1 & 0x03) - 2;
      }
      else if ((b1 & MASK_2_BITS) == OP_LUMA)
      {
        int b2;
        read_ui8(data, data_p, output);
        b2 = output;
        int vg = (b1 & 0x3f) - 32;
        pixel.colors.red += vg - 8 + ((b2 >> 4) & 0x0f);
        pixel.colors.green += vg;
        pixel.colors.blue += vg - 8 + (b2 & 0x0f);
      }
      else if ((b1 & MASK_2_BITS) == OP_RUN)
      {
        run = (b1 & 0x3f);
      }

      index[color_hash(pixel.colors) % 64] = pixel;
    }

    this->pixels[pixel_p++] = pixel.colors.red;
    this->pixels[pixel_p++] = pixel.colors.green;
    this->pixels[pixel_p++] = pixel.colors.blue;
    if(this->channels == 4) this->pixels[pixel_p++] = pixel.colors.alpha;
  }

  delete[] data;
}

Image::~Image()
{
  delete[] this->pixels;
}

}