/*
 * Copyright 2022-2024 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>

#include "binary_reader.hpp"

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{0}
{
}

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;
}

}

namespace BluCat::QOI
{

Image::Image(const char *file_path, uint8_t channels):
  channels{channels}
{
  if(this->channels != 0 && this->channels != 3 && this->channels != 4)
  {
    throw std::invalid_argument{"invalid number of channels"};
  }

  BinaryReader input{file_path};
  if(input.size() < HEADER_SIZE + (int)sizeof(PADDING))
  {
    throw std::runtime_error{"invalid QOI file"};
  }

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

  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(0, 0, 0, 255);
  int chunks_len = input.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(input.pointer() < chunks_len)
    {
      int b1 = input.read_ui8();

      if (b1 == OP_RGB)
      {
        pixel.colors.red = input.read_ui8();
        pixel.colors.green = input.read_ui8();
        pixel.colors.blue = input.read_ui8();
      }
      else if (b1 == OP_RGBA)
      {
        pixel.colors.red = input.read_ui8();
        pixel.colors.green = input.read_ui8();
        pixel.colors.blue = input.read_ui8();
        pixel.colors.alpha = input.read_ui8();
      }
      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 = input.read_ui8();
        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;
  }
}

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

}