#include <cstdio>

#include <cstdio>
#include <cstdlib>
#include <vector>

#include "tinyexr.h"

#define TINY_DNG_WRITER_IMPLEMENTATION
#include "tiny_dng_writer.h"

static bool Create32bitFpTiff(
const float *data, // [width x height x in_channels]
const size_t width,
const size_t height,
const size_t in_channels,
const size_t channels,
tinydngwriter::DNGImage *dng_image) {

if (in_channels < 1) return false;

unsigned int image_width = uint32_t(width);
unsigned int image_height = uint32_t(height);

//dng_image->SetSubfileType(false, false, false);
dng_image->SetImageWidth(image_width);
dng_image->SetImageLength(image_height);
dng_image->SetRowsPerStrip(image_height);
dng_image->SetSamplesPerPixel(uint16_t(channels));
std::vector<uint16_t> bps(channels);
for (size_t i = 0; i < bps.size(); i++) {
bps[i] = 32;
}
dng_image->SetBitsPerSample(static_cast<unsigned int>(channels), bps.data());
dng_image->SetPlanarConfig(tinydngwriter::PLANARCONFIG_CONTIG);
dng_image->SetCompression(tinydngwriter::COMPRESSION_NONE);
if (channels == 1) {
dng_image->SetPhotometric(
tinydngwriter::PHOTOMETRIC_BLACK_IS_ZERO); // grayscale
} else {
dng_image->SetPhotometric(
tinydngwriter::PHOTOMETRIC_RGB);
}
dng_image->SetXResolution(1.0);
dng_image->SetYResolution(1.0);
dng_image->SetResolutionUnit(tinydngwriter::RESUNIT_NONE);

std::vector<uint16_t> formats(channels);
for (size_t i = 0; i < formats.size(); i++) {
formats[i] = tinydngwriter::SAMPLEFORMAT_IEEEFP;
}
dng_image->SetSampleFormat(static_cast<unsigned int>(channels), formats.data());

std::vector<float> buf;
buf.resize(size_t(channels) * image_width * image_height);

for (size_t i = 0; i < image_width * image_height; i++) {
size_t in_c = 0;
for (size_t c = 0; c < channels; c++) {
buf[channels * i + c] = data[in_channels * i + in_c];
in_c++;
in_c = std::min(in_c, in_channels - 1);
}
}

//size_t max_dump_pixels = 4096;
//for (size_t i = 0; i < std::min(max_dump_pixels, buf.size()); i++) {
// std::cout << "val[" << i << "] = " << buf[i] << "\n";
//}
//std::cout << "last = " << buf.at(image_width * image_height * channels - 1) << "\n";

// We must retain pointer address of `buf` until calling DNGWriter::WriteToFile
dng_image->SetImageData(reinterpret_cast<unsigned char *>(buf.data()),
buf.size() * sizeof(float));

if (!dng_image->Error().empty()) {
std::cout << "Err: " << dng_image->Error() << "\n";

return false;
}

return true;
}

int main(int argc, char** argv)
{
if (argc < 3) {
printf("Usage: exr2fptiff input.exr output.tiff\n");
exit(-1);
}

std::string input_filename = argv[1];
std::string output_filename = argv[2];

// Get # of layers
size_t num_layers{0};
{

EXRVersion exr_version;
{
int ret = ParseEXRVersionFromFile(&exr_version, input_filename.c_str());
if (ret != 0) {
std::cerr << "Invalid EXR file: " << input_filename << "\n";
return EXIT_FAILURE;
}

if (exr_version.multipart) {
std::cerr << "Multipart EXR file is not supported in this example.\n";
return EXIT_FAILURE;
}
}

EXRHeader exr_header;
InitEXRHeader(&exr_header);

const char* err = nullptr;
int ret = ParseEXRHeaderFromFile(&exr_header, &exr_version, argv[1], &err);
if (ret != TINYEXR_SUCCESS) {
if (err) {
std::cerr << "Parse EXR error: " << err << "\n";
FreeEXRErrorMessage(err); // free's buffer for an error message
} else {
std::cerr << "Parse EXR error.\n";
}
return EXIT_FAILURE;
}

num_layers = size_t(exr_header.num_channels);
if (num_layers == 0) {
std::cerr << "no layers found\n";
return EXIT_FAILURE;
}

if (num_layers > 4) {
std::cerr << "This program supports up to 4(e.g. RGBA) layers.\n";
return EXIT_FAILURE;
}

FreeEXRHeader(&exr_header);
}

std::cout << "# of channels = " << num_layers << "\n";

// Use legacy but easy-to-use API to read image.
float *rgba{nullptr};
int width;
int height;
{
const char *err;
int ret = LoadEXR(&rgba, &width, &height, input_filename.c_str(), &err);

if (ret != TINYEXR_SUCCESS) {
if (err) {
std::cerr << "Load EXR error: " << err << "\n";
FreeEXRErrorMessage(err); // free's buffer for an error message
} else {
std::cerr << "Load EXR error.\n";
}
return EXIT_FAILURE;
}

}

bool big_endian = false;

tinydngwriter::DNGImage tiff;
tiff.SetBigEndian(big_endian);

bool ret = Create32bitFpTiff(rgba, size_t(width), size_t(height), /* in_channels */4, size_t(num_layers), &tiff);
if (!ret) {
std::cerr << "Failed to create floating point tiff data\n";
return EXIT_FAILURE;
}

// 4. Free image data
free(rgba);

tinydngwriter::DNGWriter dng_writer(big_endian);
ret = dng_writer.AddImage(&tiff);
if (!ret) {
std::cerr << "Failed to add TIFF image to TIFF writer.\n";
return EXIT_FAILURE;
}

// 5. write tiff
std::string err;
ret = dng_writer.WriteToFile(output_filename.c_str(), &err);

if (!err.empty()) {
std::cerr << err;
}

if (!ret) {
return EXIT_FAILURE;
}

return EXIT_SUCCESS;
}