#define TINYEXR_USE_MINIZ 0
#include "zlib.h"

#define TINYEXR_OTMOD_IMPLEMENTATION
#include "tinyexr_otmod.h"

#include "tiio_exr.h"
#include "tpixel.h"

#include <QMap>

namespace {
inline unsigned char ftouc(float f, float gamma = 2.2f) {
  int i = static_cast<int>(255.0f * powf(f, 1.0f / gamma));
  if (i > 255) i = 255;
  if (i < 0) i = 0;

  return static_cast<unsigned char>(i);
}
inline float uctof(unsigned char uc, float gamma = 2.2f) {
  return powf(static_cast<float>(uc) / 255.0f, gamma);
}

inline unsigned short ftous(float f, float gamma = 2.2f) {
  int i = static_cast<int>(65535.0f * powf(f, 1.0f / gamma));
  if (i > 65535) i = 65535;
  if (i < 0) i = 0;

  return static_cast<unsigned short>(i);
}
inline float ustof(unsigned short us, float gamma = 2.2f) {
  return powf(static_cast<float>(us) / 65535.0f, gamma);
}

const QMap<int, std::wstring> ExrCompTypeStr = {
    {TINYEXR_COMPRESSIONTYPE_NONE, L"None"},
    {TINYEXR_COMPRESSIONTYPE_RLE, L"RLE"},
    {TINYEXR_COMPRESSIONTYPE_ZIPS, L"ZIPS"},
    {TINYEXR_COMPRESSIONTYPE_ZIP, L"ZIP"},
    {TINYEXR_COMPRESSIONTYPE_PIZ, L"PIZ"}};

const std::wstring EXR_STORAGETYPE_SCANLINE = L"Store Image as Scanlines";
const std::wstring EXR_STORAGETYPE_TILE     = L"Store Image as Tiles";
}  // namespace

//**************************************************************************
//    ExrReader  implementation
//**************************************************************************

class ExrReader final : public Tiio::Reader {
  float* m_rgbaBuf;
  int m_row;
  EXRHeader* m_exr_header;
  FILE* m_fp;

public:
  ExrReader();
  ~ExrReader();

  void open(FILE* file) override;
  Tiio::RowOrder getRowOrder() const override;
  bool read16BitIsEnabled() const override;
  int skipLines(int lineCount) override;
  ;
  void readLine(char* buffer, int x0, int x1, int shrink) override;
  void readLine(short* buffer, int x0, int x1, int shrink) override;
  void loadImage();
};

ExrReader::ExrReader() : m_rgbaBuf(nullptr), m_row(0), m_exr_header(nullptr) {}

ExrReader::~ExrReader() {
  if (m_rgbaBuf) free(m_rgbaBuf);
  if (m_exr_header) FreeEXRHeader(m_exr_header);
}

void ExrReader::open(FILE* file) {
  m_fp         = file;
  m_exr_header = new EXRHeader();
  const char* err;
  {
    int ret = LoadEXRHeaderFromFileHandle(*m_exr_header, file, &err);
    if (ret != 0) {
      m_exr_header = nullptr;
      throw(std::string(err));
    }
  }
  m_info.m_lx =
      m_exr_header->data_window.max_x - m_exr_header->data_window.min_x + 1;
  m_info.m_ly =
      m_exr_header->data_window.max_y - m_exr_header->data_window.min_y + 1;

  m_info.m_samplePerPixel = m_exr_header->num_channels;

  int bps = 16;
  switch (m_exr_header->pixel_types[0]) {
  case TINYEXR_PIXELTYPE_UINT:
  case TINYEXR_PIXELTYPE_FLOAT:
    bps = 32;
    break;
  case TINYEXR_PIXELTYPE_HALF:
    bps = 16;
    break;
  }
  m_info.m_bitsPerSample = bps;
}

Tiio::RowOrder ExrReader::getRowOrder() const { return Tiio::TOP2BOTTOM; }

bool ExrReader::read16BitIsEnabled() const { return true; }

int ExrReader::skipLines(int lineCount) {
  m_row += lineCount;
  return lineCount;
}

void ExrReader::loadImage() {
  assert(!m_rgbaBuf);
  const char* err;
  {
    int ret =
        LoadEXRImageBufFromFileHandle(&m_rgbaBuf, *m_exr_header, m_fp, &err);
    if (ret != 0) {
      m_exr_header = nullptr;
      throw(std::string(err));
    }
  }
  // header memory is freed after loading image
  m_exr_header = nullptr;
}

void ExrReader::readLine(char* buffer, int x0, int x1, int shrink) {
  const int pixelSize = 4;
  if (m_row < 0 || m_row >= m_info.m_ly) {
    memset(buffer, 0, (x1 - x0 + 1) * pixelSize);
    m_row++;
    return;
  }

  if (!m_rgbaBuf) loadImage();

  TPixel32* pix = (TPixel32*)buffer;
  float* v      = m_rgbaBuf + m_row * m_info.m_lx * 4;

  pix += x0;
  v += x0 * 4;

  int width =
      (x1 < x0) ? (m_info.m_lx - 1) / shrink + 1 : (x1 - x0) / shrink + 1;

  for (int i = 0; i < width; i++) {
    pix->r = ftouc(v[0]);
    pix->g = ftouc(v[1]);
    pix->b = ftouc(v[2]);
    pix->m = ftouc(v[3]);

    v += shrink * 4;
    pix += shrink;
  }

  m_row++;
}

void ExrReader::readLine(short* buffer, int x0, int x1, int shrink) {
  const int pixelSize = 8;
  if (m_row < 0 || m_row >= m_info.m_ly) {
    memset(buffer, 0, (x1 - x0 + 1) * pixelSize);
    m_row++;
    return;
  }

  if (!m_rgbaBuf) loadImage();

  TPixel64* pix = (TPixel64*)buffer;
  float* v      = m_rgbaBuf + m_row * m_info.m_lx * 4;

  pix += x0;
  v += x0 * 4;

  int width =
      (x1 < x0) ? (m_info.m_lx - 1) / shrink + 1 : (x1 - x0) / shrink + 1;

  for (int i = 0; i < width; i++) {
    pix->r = ftous(v[0]);
    pix->g = ftous(v[1]);
    pix->b = ftous(v[2]);
    pix->m = ftous(v[3], 1.0f);

    v += shrink * 4;
    pix += shrink;
  }

  m_row++;
}

//============================================================

Tiio::ExrWriterProperties::ExrWriterProperties()
    : m_compressionType("Compression Type")
    , m_storageType("Storage Type")
    , m_bitsPerPixel("Bits Per Pixel") {
  m_bitsPerPixel.addValue(L"48(RGB)");
  m_bitsPerPixel.addValue(L"64(RGBA)");
  m_bitsPerPixel.setValue(L"64(RGBA)");

  m_compressionType.addValue(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_NONE));
  m_compressionType.addValue(ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_RLE));
  m_compressionType.addValue(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_ZIPS));
  m_compressionType.addValue(ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_ZIP));
  m_compressionType.addValue(ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_PIZ));

  m_compressionType.setValue(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_NONE));

  m_storageType.addValue(EXR_STORAGETYPE_SCANLINE);
  m_storageType.addValue(EXR_STORAGETYPE_TILE);
  m_storageType.setValue(EXR_STORAGETYPE_SCANLINE);

  bind(m_bitsPerPixel);
  bind(m_compressionType);
  bind(m_storageType);
}

void Tiio::ExrWriterProperties::updateTranslation() {
  m_bitsPerPixel.setQStringName(tr("Bits Per Pixel"));
  m_bitsPerPixel.setItemUIName(L"48(RGB)", tr("48(RGB Half Float)"));
  m_bitsPerPixel.setItemUIName(L"64(RGBA)", tr("64(RGBA Half Float)"));

  m_compressionType.setQStringName(tr("Compression Type"));
  m_compressionType.setItemUIName(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_NONE), tr("No compression"));
  m_compressionType.setItemUIName(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_RLE),
      tr("Run Length Encoding (RLE)"));
  m_compressionType.setItemUIName(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_ZIPS),
      tr("ZIP compression per Scanline (ZIPS)"));
  m_compressionType.setItemUIName(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_ZIP),
      tr("ZIP compression per scanline band (ZIP)"));
  m_compressionType.setItemUIName(
      ExrCompTypeStr.value(TINYEXR_COMPRESSIONTYPE_PIZ),
      tr("PIZ-based wavelet compression (PIZ)"));

  m_storageType.setQStringName(tr("Storage Type"));
  m_storageType.setItemUIName(EXR_STORAGETYPE_SCANLINE, tr("Scan-line based"));
  m_storageType.setItemUIName(EXR_STORAGETYPE_TILE, tr("Tile based"));
}

//============================================================

class ExrWriter final : public Tiio::Writer {
  std::vector<float> m_imageBuf[4];
  EXRHeader m_header;
  EXRImage m_image;
  int m_row;
  FILE* m_fp;
  int m_bpp;

public:
  ExrWriter();
  ~ExrWriter();

  void open(FILE* file, const TImageInfo& info) override;
  void writeLine(char* buffer) override;
  void writeLine(short* buffer) override;

  void flush() override;

  Tiio::RowOrder getRowOrder() const override { return Tiio::TOP2BOTTOM; }

  // m_bpp is set to "Bits Per Pixel" property value in the function open()
  bool writeAlphaSupported() const override { return m_bpp == 64; }
};

ExrWriter::ExrWriter() : m_row(0), m_bpp(64) {}

ExrWriter::~ExrWriter() {
  free(m_header.channels);
  free(m_header.pixel_types);
  free(m_header.requested_pixel_types);
}

void ExrWriter::open(FILE* file, const TImageInfo& info) {
  m_fp   = file;
  m_info = info;
  InitEXRHeader(&m_header);
  InitEXRImage(&m_image);

  if (!m_properties) m_properties = new Tiio::ExrWriterProperties();

  TEnumProperty* bitsPerPixel =
      (TEnumProperty*)(m_properties->getProperty("Bits Per Pixel"));
  m_bpp = bitsPerPixel ? std::stoi(bitsPerPixel->getValue()) : 64;
  assert(m_bpp == 48 || m_bpp == 64);

  std::wstring compressionType =
      ((TEnumProperty*)(m_properties->getProperty("Compression Type")))
          ->getValue();
  m_header.compression_type = ExrCompTypeStr.key(compressionType);

  std::wstring storageType =
      ((TEnumProperty*)(m_properties->getProperty("Storage Type")))->getValue();
  if (storageType == EXR_STORAGETYPE_TILE) {
    m_header.tiled           = 1;
    m_header.tile_size_x     = 128;
    m_header.tile_size_y     = 128;
    m_header.tile_level_mode = TINYEXR_TILE_ONE_LEVEL;
  } else
    m_header.tiled = 0;

  m_image.num_channels = (m_bpp == 64) ? 4 : 3;

  for (int c = 0; c < m_image.num_channels; c++)
    m_imageBuf[c].resize(m_info.m_lx * m_info.m_ly);

  m_image.width  = m_info.m_lx;
  m_image.height = m_info.m_ly;

  m_header.num_channels = m_image.num_channels;
  m_header.channels =
      (EXRChannelInfo*)malloc(sizeof(EXRChannelInfo) * m_header.num_channels);
  // Must be BGR(A) order, since most of EXR viewers expect this channel order.
  if (m_bpp == 64) {
    strncpy(m_header.channels[0].name, "A", 255);
    m_header.channels[0].name[strlen("A")] = '\0';
    strncpy(m_header.channels[1].name, "B", 255);
    m_header.channels[1].name[strlen("B")] = '\0';
    strncpy(m_header.channels[2].name, "G", 255);
    m_header.channels[2].name[strlen("G")] = '\0';
    strncpy(m_header.channels[3].name, "R", 255);
    m_header.channels[3].name[strlen("R")] = '\0';
  } else {
    strncpy(m_header.channels[0].name, "B", 255);
    m_header.channels[0].name[strlen("B")] = '\0';
    strncpy(m_header.channels[1].name, "G", 255);
    m_header.channels[1].name[strlen("G")] = '\0';
    strncpy(m_header.channels[2].name, "R", 255);
    m_header.channels[2].name[strlen("R")] = '\0';
  }

  m_header.pixel_types = (int*)malloc(sizeof(int) * m_header.num_channels);
  m_header.requested_pixel_types =
      (int*)malloc(sizeof(int) * m_header.num_channels);
  for (int i = 0; i < m_header.num_channels; i++) {
    m_header.pixel_types[i] =
        TINYEXR_PIXELTYPE_FLOAT;  // pixel type of input image
    m_header.requested_pixel_types[i] =
        TINYEXR_PIXELTYPE_HALF;  // pixel type of output image to be stored in
                                 // .EXR
  }
}

void ExrWriter::writeLine(char* buffer) {
  TPixel32* pix    = (TPixel32*)buffer;
  TPixel32* endPix = pix + m_info.m_lx;

  float* r_p = &m_imageBuf[0][m_row * m_info.m_lx];
  float* g_p = &m_imageBuf[1][m_row * m_info.m_lx];
  float* b_p = &m_imageBuf[2][m_row * m_info.m_lx];
  float* a_p;
  if (m_bpp == 64) a_p = &m_imageBuf[3][m_row * m_info.m_lx];
  while (pix < endPix) {
    *r_p++ = uctof(pix->r);
    *g_p++ = uctof(pix->g);
    *b_p++ = uctof(pix->b);
    if (m_bpp == 64) *a_p++ = uctof(pix->m, 1.0f);
    pix++;
  }
  m_row++;
}
void ExrWriter::writeLine(short* buffer) {
  TPixel64* pix    = (TPixel64*)buffer;
  TPixel64* endPix = pix + m_info.m_lx;

  float* r_p = &m_imageBuf[0][m_row * m_info.m_lx];
  float* g_p = &m_imageBuf[1][m_row * m_info.m_lx];
  float* b_p = &m_imageBuf[2][m_row * m_info.m_lx];
  float* a_p;
  if (m_bpp == 64) a_p = &m_imageBuf[3][m_row * m_info.m_lx];
  while (pix < endPix) {
    *r_p++ = ustof(pix->r);
    *g_p++ = ustof(pix->g);
    *b_p++ = ustof(pix->b);
    if (m_bpp == 64) *a_p++ = ustof(pix->m, 1.0f);
    pix++;
  }
  m_row++;
}

void ExrWriter::flush() {
  if (m_bpp == 64) {
    float* image_ptr[4];
    image_ptr[0]   = &(m_imageBuf[3].at(0));  // B
    image_ptr[1]   = &(m_imageBuf[2].at(0));  // G
    image_ptr[2]   = &(m_imageBuf[1].at(0));  // R
    image_ptr[3]   = &(m_imageBuf[0].at(0));  // A
    m_image.images = (unsigned char**)image_ptr;
    const char* err;
    int ret = SaveEXRImageToFileHandle(&m_image, &m_header, m_fp, &err);
    if (ret != TINYEXR_SUCCESS) {
      throw(std::string(err));
    }
  } else {
    float* image_ptr[3];
    image_ptr[0]   = &(m_imageBuf[2].at(0));  // B
    image_ptr[1]   = &(m_imageBuf[1].at(0));  // G
    image_ptr[2]   = &(m_imageBuf[0].at(0));  // R
    m_image.images = (unsigned char**)image_ptr;
    const char* err;
    int ret = SaveEXRImageToFileHandle(&m_image, &m_header, m_fp, &err);
    if (ret != TINYEXR_SUCCESS) {
      throw(std::string(err));
    }
  }
}

//============================================================

Tiio::Reader* Tiio::makeExrReader() { return new ExrReader(); }

//------------------------------------------------------------

Tiio::Writer* Tiio::makeExrWriter() { return new ExrWriter(); }