#include "EXRFiles.h"

EXRFile::EXRFile(std::string const& input)
{
    char const* err = nullptr;
    if (LoadEXR(&flt, &width, &height, input.c_str(), &err) != TINYEXR_SUCCESS)
    {
        if (err)
        {
            std::cerr << "TinyEXR ERROR: " << err << std::endl;
            FreeEXRErrorMessage(err);
            exit(1);
        }
    }
}

OEXRFile::OEXRFile(std::vector<float3> const& dat, int width, int height, int compressionType) :
    width(std::move(width)), height(std::move(height)),
    compressionType(compressionType), infos(3)
{
    for (float3 const& col : dat)
    {
        r.push_back(col.x);
        g.push_back(col.y);
        b.push_back(col.z);
        //a.push_back(col.w);
    }

    for (int i = 0; i < 3; ++i)
    {
        pixelType.push_back(TINYEXR_PIXELTYPE_FLOAT);
        reqPixelType.push_back(TINYEXR_PIXELTYPE_FLOAT);
    }

    initChannelInfo();
    initHeader();
}

OEXRFile::OEXRFile(std::vector<float2> const& dat, int width, int height, int compressionType) :
    width(std::move(width)), height(std::move(height)),
    compressionType(compressionType), infos(3)
{
    for (float2 const& col : dat)
    {
        r.push_back(col.x);
        g.push_back(col.y);
        b.push_back(0);
        //a.push_back(col.w);
    }

    for (int i = 0; i < 3; ++i)
    {
        pixelType.push_back(TINYEXR_PIXELTYPE_FLOAT);
        reqPixelType.push_back(TINYEXR_PIXELTYPE_FLOAT);
    }

    initChannelInfo();
    initHeader();
}

void OEXRFile::initChannelInfo()
{
    infos.resize(4);
    // strcpy(infos[0].name, "A");
    strcpy(infos[0].name, "B");
    strcpy(infos[1].name, "G");
    strcpy(infos[2].name, "R");

    for (auto& info : infos)
    {
        info.name[1] = '\0';
    }
}

void OEXRFile::initHeader()
{
    InitEXRHeader(&hd);
    hd.num_channels = 3;
    hd.channels = infos.data();
    hd.pixel_types = pixelType.data();
    hd.requested_pixel_types = reqPixelType.data();
    hd.compression_type = compressionType;
}

void OEXRFile::write(std::string const& filename)
{
    EXRImage im;
    InitEXRImage(&im);
    im.num_channels = 3;
    im.width = width;
    im.height = height;

    std::array<float*, 3> arr{ b.data(), g.data(), r.data() };

    im.images = reinterpret_cast<unsigned char**>(arr.data());

    char const* err = nullptr;
    if (SaveEXRImageToFile(&im, &hd, filename.c_str(), &err) != TINYEXR_SUCCESS)
    {
        std::cerr << "TinyEXR ERROR: " << err << std::endl;
        FreeEXRErrorMessage(err);
        exit(1);
    }
}