/*
 * Copyright (C) 2021 - 2022 René Rebe, ExactCODE GmbH
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; version 2. A copy of the GNU General
 * Public License can be found in the file LICENSE.
 * 
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANT-
 * ABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
 * Public License for more details.
 *
 * Alternatively, commercial licensing options are available from the
 * copyright holder ExactCODE GmbH Germany.
 */

#include "qoi.hh"

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

#include <iostream>
#include <sstream>
#include <vector>
#include <algorithm>

#include <Endianess.hh>
using Exact::EndianessConverter;
using Exact::BigEndianTraits;

/* TODO:
 * lossy encoding mode?
 * try bitfields or so for some operands
 * some more error handling
 * delayed, on-demand decoding
 * proof-of-concept gray8
 */

typedef int8_t s8;
typedef uint8_t u8;
typedef EndianessConverter<uint32_t, BigEndianTraits> u32_be;

static const char QoiMagic[] = "qoif";
static const char QoiEof[] = "\0\0\0\0\0\0\0\1";

struct QoiHeader {
  //char magic[4];
  u32_be width;
  u32_be height;
  u8 channels;
  u8 colorspace;
}
#ifdef __GNUC__
__attribute__((packed))
#endif
;

enum QoiTag {
  QOI_OP_INDEX = 0b00000000,
  QOI_OP_DIFF = 0b01000000,
  QOI_OP_LUMA = 0b10000000,
  QOI_OP_RUN = 0b11000000,
  QOI_OP_RGB = 0b11111110,
  QOI_OP_RGBA = 0b11111111,
  QOI_OP_MASK = 0b00111111,
};

// OP_INDEX: u8 index
// OP_DIFF: u8 dr, dg, gb
// OP_LUMA: u16, dg, dr, db
// OP_RUN: u8 run + 1

struct QoiPixel {
  u8 r, g, b, a;
}
#ifdef __GNUC__
__attribute__((packed))
#endif
;

bool operator==(const QoiPixel& a, const QoiPixel& b) {
  // TODO: simply cast to u32?
  return a.r == b.r &&
    a.g == b.g &&
    a.b == b.b &&
    a.a == b.a; 
}

inline u8 qoiIndex(QoiPixel& p) {
  return (p.r * 3 + p.g * 5 + p.b * 7 + p.a * 11) & 0b111111;
}

int QOICodec::readImage(std::istream* stream, Image& image, const std::string& decompress)
{
  if (true) {
    // quick magic check
    std::string buf(4, '\0');
    stream->read(&buf[0], buf.size());
    if (buf != QoiMagic)
      return false;
  }
  
  QoiHeader header;
  
  // read file
  stream->read((char*)&header, sizeof(header));
  
  if (header.channels < 3 || header.channels > 4) {
    std::cerr << "unsupported channels: " << (unsigned)header.channels << std::endl;
    return false;
  }
  
  if (header.colorspace > 1) {
    std::cerr << "unsupported colorspace" << std::endl;
    return false;
  }
  
  image.bps = 8;
  image.spp = header.channels;
  image.resize(header.width, header.height);
  
  QoiPixel history[64] = {};
  QoiPixel pixel = {.a = 255};
 
  u8* begin = image.getRawData();
  u8* it = begin;
  const u8* end = begin + image.stride() * header.height;
  
  for (int tag = stream->get(); tag != EOF; tag = stream->get()) {
    s8 run = 0;
    switch (tag) {
    case QOI_OP_RGB:
      stream->read((char*)&pixel, 3);
      break;
      
    case QOI_OP_RGBA:
      stream->read((char*)&pixel, 4);
      break;
    
    default:
      switch (tag & 0b11000000) {
      case QOI_OP_INDEX:
	pixel = history[tag]; //  & QOI_OP_MASK]; // high bits already zero
	goto nohistory;
	break;
      
      case QOI_OP_RUN:
	run = tag & QOI_OP_MASK;
	goto nohistory;
	break;
      
      case QOI_OP_DIFF:
	pixel.r += ((tag >> 4) & 0b11) - 2;
	pixel.g += ((tag >> 2) & 0b11) - 2;
	pixel.b += ((tag >> 0) & 0b11) - 2;
	break;
      
      case QOI_OP_LUMA:
	{
	  s8 dg = (tag & 0b111111) - 32;
	  u8 d = stream->get(); // TODO: EOF
	  pixel.r += dg + (d >> 4) - 8;
	  pixel.g += dg;
	  pixel.b += dg + (d & 0b1111) - 8;
	}
	break;
      }
    }
    
    history[qoiIndex(pixel)] = pixel;
   
  nohistory:
    
    for (; run >= 0  && it < end; --run) {
      *it++ = pixel.r;
      *it++ = pixel.g;
      *it++ = pixel.b;
      if (header.channels > 3)
	*it++ = pixel.a;
    }
  }
  
  // TODO: check QoiEof padding trailer
  
  return true;
}

static inline void qoiWriteRun(std::ostream* stream, unsigned& run) {
  for (; run > 0; ) {
    u8 thisrun = std::min(run, (unsigned)62);
    stream->put(QOI_OP_RUN | (thisrun - 1));
    run -= thisrun;
  }
}

bool QOICodec::writeImage(std::ostream* stream, Image& image, int quality,
			     const std::string& compress)
{
  if (image.spp < 3 || image.spp > 4 ||
      image.bps != 8) {
    std::cerr << "QOICodec: unsuppported color format" << std::endl;
  }
  
  QoiHeader header = {};
  header.width = image.w;
  header.height = image.h;
  header.channels = image.spp;
  
  stream->write(QoiMagic, strlen(QoiMagic));
  stream->write((char*)&header, sizeof(header));

  QoiPixel history[64] = {};
  QoiPixel lpixel = {.a = 255};
  QoiPixel pixel;
  
  u8* begin = image.getRawData();
  unsigned run = 0;
  for (int y = 0; y < image.h; ++y) {
    u8* it = begin + y * image.stride();
    for (int x = 0; x < image.w; ++x) {
      pixel.r = *it++;
      pixel.g = *it++;
      pixel.b = *it++;
      
      if (header.channels > 3)
	pixel.a = *it++;
      
      if (pixel == lpixel) {
	++run;
      } else {
	qoiWriteRun(stream, run);
	
	u8 index = qoiIndex(pixel);
	if (history[index] == pixel) {
	  stream->put(QOI_OP_INDEX | index);
	} else {
	  s8 dr = pixel.r - lpixel.r;
	  s8 dg = pixel.g - lpixel.g;
	  s8 db = pixel.b - lpixel.b;
	  s8 da = pixel.a - lpixel.a;
	  
	  s8 drg = dr -dg;
	  s8 dbg = db -dg;
	  
	  if (da == 0 &&
	      dr >= -2 && dr <= 1 &&
	      dg >= -2 && dg <= 1 &&
	      db >= -2 && db <= 1) {
	    dr += 2;
	    dg += 2;
	    db += 2;
	    stream->put(QOI_OP_DIFF | dr << 4 | dg << 2 | db);
	  } else if (da == 0 &&
		     dg >= -32 && dg <= 31 &&
		     drg >= -8 && drg <= 7 &&
		     dbg >= -8 && dbg <= 7) {
	    dg += 32;
	    drg += 8;
	    dbg += 8;
	    stream->put(QOI_OP_LUMA | dg);
	    stream->put(drg << 4 | dbg);
	  } else {
	    stream->put(pixel.a == lpixel.a ? QOI_OP_RGB : QOI_OP_RGBA);
	    stream->put(pixel.r);
	    stream->put(pixel.g);      
	    stream->put(pixel.b);
	    if (pixel.a != lpixel.a)
	      stream->put(pixel.a);
	  }
	  
	  history[qoiIndex(pixel)] = pixel;
	}
	lpixel = pixel;
      }
    }
  }
  
  qoiWriteRun(stream, run);
  
  stream->write(QoiEof, sizeof(QoiEof) - 1);
  
  return true;
}

QOICodec qoi_loader;