// Copyright 2008-present Contributors to the OpenImageIO project.
// SPDX-License-Identifier: BSD-3-Clause
// https://github.com/OpenImageIO/oiio/blob/master/LICENSE.md

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <ctime>

#include <OpenImageIO/dassert.h>
#include <OpenImageIO/filesystem.h>
#include <OpenImageIO/fmath.h>
#include <OpenImageIO/imageio.h>
#include <OpenImageIO/strutil.h>
#include <OpenImageIO/sysutil.h>
#include <OpenImageIO/typedesc.h>

#include "rla_pvt.h"



OIIO_PLUGIN_NAMESPACE_BEGIN

using namespace RLA_pvt;


class RLAOutput final : public ImageOutput {
public:
    RLAOutput();
    virtual ~RLAOutput();
    virtual const char* format_name(void) const override { return "rla"; }
    virtual int supports(string_view feature) const override;
    virtual bool open(const std::string& name, const ImageSpec& spec,
                      OpenMode mode = Create) override;
    virtual bool close() override;
    virtual bool write_scanline(int y, int z, TypeDesc format, const void* data,
                                stride_t xstride) override;
    virtual bool write_tile(int x, int y, int z, TypeDesc format,
                            const void* data, stride_t xstride,
                            stride_t ystride, stride_t zstride) override;

private:
    std::string m_filename;  ///< Stash the filename
    FILE* m_file;            ///< Open image handle
    std::vector<unsigned char> m_scratch;
    RLAHeader m_rla;                   ///< Wavefront RLA header
    std::vector<uint32_t> m_sot;       ///< Scanline offset table
    std::vector<unsigned char> m_rle;  ///< Run record buffer for RLE
    std::vector<unsigned char> m_tilebuffer;
    unsigned int m_dither;

    // Initialize private members to pre-opened state
    void init(void)
    {
        m_file = NULL;
        m_sot.clear();
    }

    /// Helper - sets a chromaticity from attribute
    inline void set_chromaticity(const ParamValue* p, char* dst,
                                 size_t field_size, const char* default_val);

    // Helper - handles the repetitive work of encoding and writing a
    // channel. The data is guaranteed to be in the scratch area and need
    // not be preserved.
    bool encode_channel(unsigned char* data, stride_t xstride,
                        TypeDesc chantype, int bits);

    /// Helper - write, with error detection
    bool fwrite(const void* buf, size_t itemsize, size_t nitems)
    {
        size_t n = ::fwrite(buf, itemsize, nitems, m_file);
        if (n != nitems)
            errorf("Write error: wrote %d records of %d", (int)n, (int)nitems);
        return n == nitems;
    }

    /// Helper: write buf[0..nitems-1], swap endianness if necessary
    template<typename T> bool write(const T* buf, size_t nitems = 1)
    {
        if (littleendian()
            && (is_same<T, uint16_t>::value || is_same<T, int16_t>::value
                || is_same<T, uint32_t>::value || is_same<T, int32_t>::value)) {
            T* newbuf = OIIO_ALLOCA(T, nitems);
            memcpy(newbuf, buf, nitems * sizeof(T));
            swap_endian(newbuf, nitems);
            buf = newbuf;
        }
        return fwrite(buf, sizeof(T), nitems);
    }
};



// Obligatory material to make this a recognizeable imageio plugin:
OIIO_PLUGIN_EXPORTS_BEGIN

OIIO_EXPORT ImageOutput*
rla_output_imageio_create()
{
    return new RLAOutput;
}

// OIIO_EXPORT int rla_imageio_version = OIIO_PLUGIN_VERSION;   // it's in rlainput.cpp

OIIO_EXPORT const char* rla_output_extensions[] = { "rla", nullptr };

OIIO_PLUGIN_EXPORTS_END


RLAOutput::RLAOutput() { init(); }



RLAOutput::~RLAOutput()
{
    // Close, if not already done.
    close();
}



int
RLAOutput::supports(string_view feature) const
{
    if (feature == "random_access")
        return true;
    if (feature == "displaywindow")
        return true;
    if (feature == "origin")
        return true;
    if (feature == "negativeorigin")
        return true;
    if (feature == "alpha")
        return true;
    if (feature == "nchannels")
        return true;
    if (feature == "channelformats")
        return true;
    // Support nothing else nonstandard
    return false;
}



bool
RLAOutput::open(const std::string& name, const ImageSpec& userspec,
                OpenMode mode)
{
    if (mode != Create) {
        errorf("%s does not support subimages or MIP levels", format_name());
        return false;
        // FIXME -- the RLA format supports subimages, but our writer
        // doesn't.  I'm not sure if it's worth worrying about for an
        // old format that is so rarely used.  We'll come back to it if
        // anybody actually encounters a multi-subimage RLA in the wild.
    }

    close();            // Close any already-opened file
    m_spec = userspec;  // Stash the spec
    if (m_spec.format == TypeDesc::UNKNOWN)
        m_spec.format = TypeDesc::UINT8;  // Default to uint8 if unknown

    m_file = Filesystem::fopen(name, "wb");
    if (!m_file) {
        errorf("Could not open \"%s\"", name);
        return false;
    }

    // Check for things this format doesn't support
    if (m_spec.width < 1 || m_spec.height < 1) {
        errorf("Image resolution must be at least 1x1, you asked for %d x %d",
               m_spec.width, m_spec.height);
        return false;
    }
    if (m_spec.width > 65535 || m_spec.height > 65535) {
        errorf(
            "Image resolution %d x %d too large for RLA (maxiumum 65535x65535)",
            m_spec.width, m_spec.height);
        return false;
    }

    if (m_spec.depth < 1)
        m_spec.depth = 1;
    else if (m_spec.depth > 1) {
        errorf("%s does not support volume images (depth > 1)", format_name());
        return false;
    }

    m_dither = (m_spec.format == TypeDesc::UINT8)
                   ? m_spec.get_int_attribute("oiio:dither", 0)
                   : 0;

    // prepare and write the RLA header
    memset(&m_rla, 0, sizeof(m_rla));
    // frame and window coordinates
    m_rla.WindowLeft   = m_spec.full_x;
    m_rla.WindowRight  = m_spec.full_x + m_spec.full_width - 1;
    m_rla.WindowTop    = m_spec.full_height - 1 - m_spec.full_y;
    m_rla.WindowBottom = m_rla.WindowTop - m_spec.full_height + 1;

    m_rla.ActiveLeft   = m_spec.x;
    m_rla.ActiveRight  = m_spec.x + m_spec.width - 1;
    m_rla.ActiveTop    = m_spec.height - 1 - m_spec.y;
    m_rla.ActiveBottom = m_rla.ActiveTop - m_spec.height + 1;

    m_rla.FrameNumber = m_spec.get_int_attribute("rla:FrameNumber", 0);

    // figure out what's going on with the channels
    int remaining = m_spec.nchannels;
    if (m_spec.channelformats.size()) {
        int streak;
        // accomodate first 3 channels of the same type as colour ones
        for (streak = 1; streak <= 3 && remaining > 0; ++streak, --remaining)
            if (m_spec.channelformats[streak] != m_spec.channelformats[0]
                || m_spec.alpha_channel == streak || m_spec.z_channel == streak)
                break;
        m_rla.ColorChannelType = rla_type(m_spec.channelformats[0]);
        int bits = m_spec.get_int_attribute("oiio:BitsPerSample", 0);
        m_rla.NumOfChannelBits = bits ? bits
                                      : m_spec.channelformats[0].size() * 8;
        // limit to 3 in case the loop went further
        m_rla.NumOfColorChannels = std::min(streak, 3);
        // if we have anything left and it looks like alpha, treat it as alpha
        if (remaining && m_spec.z_channel != m_rla.NumOfColorChannels) {
            for (streak = 1; remaining > 0; ++streak, --remaining)
                if (m_spec.channelformats[m_rla.NumOfColorChannels + streak]
                    != m_spec.channelformats[m_rla.NumOfColorChannels])
                    break;
            m_rla.MatteChannelType = rla_type(
                m_spec.channelformats[m_rla.NumOfColorChannels]);
            m_rla.NumOfMatteBits
                = bits ? bits
                       : m_spec.channelformats[m_rla.NumOfColorChannels].size()
                             * 8;
            m_rla.NumOfMatteChannels = streak;
        } else {
            m_rla.MatteChannelType   = CT_BYTE;
            m_rla.NumOfMatteBits     = 8;
            m_rla.NumOfMatteChannels = 0;
        }
        // and if there's something more left, put it in auxiliary
        if (remaining) {
            for (streak = 1; remaining > 0; ++streak, --remaining)
                if (m_spec.channelformats[m_rla.NumOfColorChannels
                                          + m_rla.NumOfMatteChannels + streak]
                    != m_spec.channelformats[m_rla.NumOfColorChannels
                                             + m_rla.NumOfMatteChannels])
                    break;
            m_rla.AuxChannelType = rla_type(
                m_spec.channelformats[m_rla.NumOfColorChannels
                                      + m_rla.NumOfMatteChannels]);
            m_rla.NumOfAuxBits = m_spec
                                     .channelformats[m_rla.NumOfColorChannels
                                                     + m_rla.NumOfMatteChannels]
                                     .size()
                                 * 8;
            m_rla.NumOfAuxChannels = streak;
        }
    } else {
        m_rla.ColorChannelType = m_rla.MatteChannelType = m_rla.AuxChannelType
            = rla_type(m_spec.format);
        int bits = m_spec.get_int_attribute("oiio:BitsPerSample", 0);
        if (bits) {
            m_rla.NumOfChannelBits = bits;
        } else {
            if (m_spec.channelformats.size())
                m_rla.NumOfChannelBits = m_spec.channelformats[0].size() * 8;
            else
                m_rla.NumOfChannelBits = m_spec.format.size() * 8;
        }
        m_rla.NumOfMatteBits = m_rla.NumOfAuxBits = m_rla.NumOfChannelBits;
        if (remaining >= 3) {
            // if we have at least 3 channels, treat them as colour
            m_rla.NumOfColorChannels = 3;
            remaining -= 3;
        } else {
            // otherwise let's say it's luminosity
            m_rla.NumOfColorChannels = 1;
            --remaining;
        }
        // if there's at least 1 more channel, it's alpha
        if (remaining && m_spec.z_channel != m_rla.NumOfColorChannels) {
            --remaining;
            ++m_rla.NumOfMatteChannels;
        }
        // anything left is auxiliary
        if (remaining > 0)
            m_rla.NumOfAuxChannels = remaining;
    }
    // std::cout << "color chans " << m_rla.NumOfColorChannels << " a "
    //           << m_rla.NumOfMatteChannels << " z " << m_rla.NumOfAuxChannels << "\n";
    m_rla.Revision = 0xFFFE;

    std::string colorspace = m_spec.get_string_attribute("oiio:ColorSpace",
                                                         "Unknown");
    if (Strutil::iequals(colorspace, "Linear"))
        Strutil::safe_strcpy(m_rla.Gamma, "1.0", sizeof(m_rla.Gamma));
    else if (Strutil::istarts_with(colorspace, "GammaCorrected")) {
        float g = Strutil::from_string<float>(colorspace.c_str() + 14);
        if (!(g >= 0.01f && g <= 10.0f /* sanity check */))
            g = m_spec.get_float_attribute("oiio:Gamma", 1.f);
        safe_snprintf(m_rla.Gamma, sizeof(m_rla.Gamma), "%.10f", g);
    }

    const ParamValue* p;
    // default NTSC chromaticities
    p = m_spec.find_attribute("rla:RedChroma");
    set_chromaticity(p, m_rla.RedChroma, sizeof(m_rla.RedChroma), "0.67 0.08");
    p = m_spec.find_attribute("rla:GreenChroma");
    set_chromaticity(p, m_rla.GreenChroma, sizeof(m_rla.GreenChroma),
                     "0.21 0.71");
    p = m_spec.find_attribute("rla:BlueChroma");
    set_chromaticity(p, m_rla.BlueChroma, sizeof(m_rla.BlueChroma),
                     "0.14 0.33");
    p = m_spec.find_attribute("rla:WhitePoint");
    set_chromaticity(p, m_rla.WhitePoint, sizeof(m_rla.WhitePoint),
                     "0.31 0.316");

#define STRING_FIELD(rlafield, name)                                    \
    {                                                                   \
        std::string s = m_spec.get_string_attribute(name);              \
        if (s.length()) {                                               \
            strncpy(m_rla.rlafield, s.c_str(), sizeof(m_rla.rlafield)); \
            m_rla.rlafield[sizeof(m_rla.rlafield) - 1] = 0;             \
        } else {                                                        \
            m_rla.rlafield[0] = 0;                                      \
        }                                                               \
    }

    m_rla.JobNumber = m_spec.get_int_attribute("rla:JobNumber", 0);
    STRING_FIELD(FileName, "rla:FileName");
    STRING_FIELD(Description, "ImageDescription");
    STRING_FIELD(ProgramName, "Software");
    STRING_FIELD(MachineName, "HostComputer");
    STRING_FIELD(UserName, "Artist");

    // the month number will be replaced with the 3-letter abbreviation
    time_t t = time(NULL);
    strftime(m_rla.DateCreated, sizeof(m_rla.DateCreated), "%m  %d %H:%M %Y",
             localtime(&t));
    // nice little trick - stoi() will convert the month number to integer,
    // which we then use to index this array of constants, and copy the
    // abbreviation back into the date string
    int m = clamp(Strutil::stoi(m_rla.DateCreated), 1, 12);
    static const char months[12][4] = { "JAN", "FEB", "MAR", "APR",
                                        "MAY", "JUN", "JUL", "AUG",
                                        "SEP", "OCT", "NOV", "DEC" };
    memcpy(m_rla.DateCreated, months[m - 1], 3);

    // FIXME: it appears that Wavefront have defined a set of aspect names;
    // I think it's safe not to care until someone complains
    STRING_FIELD(Aspect, "rla:Aspect");

    float aspect = m_spec.get_float_attribute("PixelAspectRatio", 1.f);
    safe_snprintf(m_rla.AspectRatio, sizeof(m_rla.AspectRatio), "%.6f", aspect);
    Strutil::safe_strcpy(m_rla.ColorChannel,
                         m_spec.get_string_attribute("rla:ColorChannel", "rgb"),
                         sizeof(m_rla.ColorChannel));
    m_rla.FieldRendered = m_spec.get_int_attribute("rla:FieldRendered", 0);

    STRING_FIELD(Time, "rla:Time");
    STRING_FIELD(Filter, "rla:Filter");
    STRING_FIELD(AuxData, "rla:AuxData");

    m_rla.rla_swap_endian();  // RLAs are big-endian
    write(&m_rla);
    m_rla.rla_swap_endian();  // flip back the endianness to native

    // write placeholder values - not all systems may expand the file with
    // zeroes upon seek
    m_sot.resize(m_spec.height, (int32_t)0);
    write(&m_sot[0], m_sot.size());

    // If user asked for tiles -- which this format doesn't support, emulate
    // it by buffering the whole image.
    if (m_spec.tile_width && m_spec.tile_height)
        m_tilebuffer.resize(m_spec.image_bytes());

    return true;
}



inline void
RLAOutput::set_chromaticity(const ParamValue* p, char* dst, size_t field_size,
                            const char* default_val)
{
    if (p && p->type().basetype == TypeDesc::FLOAT) {
        switch (p->type().aggregate) {
        case TypeDesc::VEC2:
            safe_snprintf(dst, field_size, "%.4f %.4f", ((float*)p->data())[0],
                          ((float*)p->data())[1]);
            break;
        case TypeDesc::VEC3:
            safe_snprintf(dst, field_size, "%.4f %.4f %.4f",
                          ((float*)p->data())[0], ((float*)p->data())[1],
                          ((float*)p->data())[2]);
            break;
        }
    } else
        Strutil::safe_strcpy(dst, default_val, field_size);
}



bool
RLAOutput::close()
{
    if (!m_file) {  // already closed
        init();
        return true;
    }

    bool ok = true;
    if (m_spec.tile_width) {
        // Handle tile emulation -- output the buffered pixels
        OIIO_DASSERT(m_tilebuffer.size());
        ok &= write_scanlines(m_spec.y, m_spec.y + m_spec.height, 0,
                              m_spec.format, &m_tilebuffer[0]);
        std::vector<unsigned char>().swap(m_tilebuffer);
    }

    // Now that all scanlines have been output, return to write the
    // correct scanline offset table to file and close the stream.
    fseek(m_file, sizeof(RLAHeader), SEEK_SET);
    write(&m_sot[0], m_sot.size());
    fclose(m_file);
    m_file = NULL;

    init();  // re-initialize
    return ok;
}



bool
RLAOutput::encode_channel(unsigned char* data, stride_t xstride,
                          TypeDesc chantype, int bits)
{
    if (chantype == TypeDesc::FLOAT) {
        // Special case -- float data is just dumped raw, no RLE
        uint16_t size = m_spec.width * sizeof(float);
        write(&size);
        for (int x = 0; x < m_spec.width; ++x)
            write((const float*)&data[x * xstride]);
        return true;
    }

    if (chantype == TypeDesc::UINT16 && bits != 16) {
        // Need to do bit scaling. Safe to overwrite data in place.
        for (int x = 0; x < m_spec.width; ++x) {
            unsigned short* s = (unsigned short*)(data + x * xstride);
            *s                = bit_range_convert(*s, 16, bits);
        }
    }

    m_rle.resize(2);  // reserve t bytes for the encoded size

    // multi-byte data types are sliced to MSB, nextSB, ..., LSB
    int chsize = (int)chantype.size();
    for (int byte = 0; byte < chsize; ++byte) {
        int lastval    = -1;     // last value
        int count      = 0;      // count of raw or repeats
        bool repeat    = false;  // if true, we're repeating
        int runbegin   = 0;      // where did the run begin
        int byteoffset = bigendian() ? byte : (chsize - byte - 1);
        for (int x = 0; x < m_spec.width; ++x) {
            int newval = data[x * xstride + byteoffset];
            if (count == 0) {  // beginning of a run.
                count    = 1;
                repeat   = true;  // presumptive
                runbegin = x;
            } else if (repeat) {  // We've seen one or more repeating characters
                if (newval == lastval) {
                    // another repeating value
                    ++count;
                } else {
                    // We stopped repeating.
                    if (count < 3) {
                        // If we didn't even have 3 in a row, just
                        // retroactively treat it as a raw run.
                        ++count;
                        repeat = false;
                    } else {
                        // We are ending a 3+ repetition
                        m_rle.push_back(count - 1);
                        m_rle.push_back(lastval);
                        count    = 1;
                        runbegin = x;
                    }
                }
            } else {  // Have not been repeating
                if (newval == lastval) {
                    // starting a repetition?  Output previous
                    OIIO_DASSERT(count > 1);
                    // write everything but the last char
                    --count;
                    m_rle.push_back(-count);
                    for (int i = 0; i < count; ++i)
                        m_rle.push_back(
                            data[(runbegin + i) * xstride + byteoffset]);
                    count    = 2;
                    runbegin = x - 1;
                    repeat   = true;
                } else {
                    ++count;  // another non-repeat
                }
            }

            // If the run is too long or we're at the scanline end, write
            if (count == 127 || x == m_spec.width - 1) {
                if (repeat) {
                    m_rle.push_back(count - 1);
                    m_rle.push_back(lastval);
                } else {
                    m_rle.push_back(-count);
                    for (int i = 0; i < count; ++i)
                        m_rle.push_back(
                            data[(runbegin + i) * xstride + byteoffset]);
                }
                count = 0;
            }
            lastval = newval;
        }
        OIIO_ASSERT(count == 0);
    }

    // Now that we know the size of the encoded buffer, save it at the
    // beginning
    uint16_t size = uint16_t(m_rle.size() - 2);
    m_rle[0]      = size >> 8;
    m_rle[1]      = size & 255;

    // And write the channel to the file
    return write(&m_rle[0], m_rle.size());
}



bool
RLAOutput::write_scanline(int y, int z, TypeDesc format, const void* data,
                          stride_t xstride)
{
    m_spec.auto_stride(xstride, format, spec().nchannels);
    const void* origdata = data;
    data = to_native_scanline(format, data, xstride, m_scratch, m_dither, y, z);
    OIIO_DASSERT(data != nullptr);
    if (data == origdata) {
        m_scratch.assign((unsigned char*)data,
                         (unsigned char*)data + m_spec.scanline_bytes());
        data = &m_scratch[0];
    }

    // store the offset to the scanline.  We'll swap_endian if necessary
    // when we go to actually write it.
    m_sot[m_spec.height - 1 - (y - m_spec.y)] = (uint32_t)ftell(m_file);

    size_t pixelsize = m_spec.pixel_bytes(true /*native*/);
    int offset       = 0;
    for (int c = 0; c < m_spec.nchannels; ++c) {
        TypeDesc chantype = m_spec.channelformats.size()
                                ? m_spec.channelformats[c]
                                : m_spec.format;
        int bits = (c < m_rla.NumOfColorChannels)
                       ? m_rla.NumOfChannelBits
                       : (c < (m_rla.NumOfColorChannels + m_rla.NumOfMatteBits))
                             ? m_rla.NumOfMatteBits
                             : m_rla.NumOfAuxBits;
        if (!encode_channel((unsigned char*)data + offset, pixelsize, chantype,
                            bits))
            return false;
        offset += chantype.size();
    }

    return true;
}



bool
RLAOutput::write_tile(int x, int y, int z, TypeDesc format, const void* data,
                      stride_t xstride, stride_t ystride, stride_t zstride)
{
    // Emulate tiles by buffering the whole image
    return copy_tile_to_image_buffer(x, y, z, format, data, xstride, ystride,
                                     zstride, &m_tilebuffer[0]);
}


OIIO_PLUGIN_NAMESPACE_END
