/**
 * OpenColorIO ColorSpace Iop.
 */

#include "OCIOColorSpace.h"

namespace OCIO = OCIO_NAMESPACE;

#include <string>
#include <sstream>
#include <fstream>
#include <stdexcept>

#include <DDImage/Channel.h>
#include "DDImage/Colorspace_KnobI.h"
#include <DDImage/PixelIop.h>
#include <DDImage/NukeWrapper.h>
#include <DDImage/Row.h>
#include <DDImage/Knobs.h>
#include <DDImage/Enumeration_KnobI.h>
#include <DDImage/NodeI.h>

OCIOColorSpace::OCIOColorSpace(Node *n)
  : DD::Image::OCIOPluginBase(n)
{
    m_inputColorSpaceIndex = 0;
    m_outputColorSpaceIndex = 0;
    m_reverseTransform = false;

    m_processor = nullptr;
    m_reverseProcessor = nullptr;

    setExtendedOpProperties(std::make_shared<DD::Image::ExtendedOpProperties>(kNumOCIOContextVariables));
}

OCIOColorSpace::~OCIOColorSpace()
{

}

DD::Image::Hash OCIOColorSpace::gpuEngine_shader_hash_at_impl(double time)
{
  DD::Image::Hash hash;

  // Using append() takes into account the current config.
  // TODO: No doubt there's some way to animate the config selection but this doesn't
  // support that. In practice, I can't believe anyone would do this.
  append(hash);

  // We also need to hash the knobs that determine which processor gets used,
  // evaluated at the specified time.

  if (DD::Image::Knob* inKnob = knob("in_colorspace")) {
    hash.append(inKnob->get_value_at(time));
  }

  if (DD::Image::Knob* outKnob = knob("out_colorspace")) {
    hash.append(outKnob->get_value_at(time));
  }

  if (DD::Image::Knob* reverseKnob = knob("reverse_transform")) {
    hash.append(reverseKnob->get_value_at(time));
  }

  return hash;
}

void OCIOColorSpace::knobs(DD::Image::Knob_Callback f)
{
    updateConfigBeforeCreatingKnobs(f);

    DD::Image::Colorspace_knob(f, &m_inputColorSpaceIndex, OCIO::ROLE_SCENE_LINEAR, "in_colorspace", "input");
    DD::Image::Tooltip(f, "\nInput data is taken to be in this color transform. \n\n"
                          "The first menu level is the color space family name. \n\n"
                          "The sub menu items are the names of the color spaces themselves."
                          "\n\n"
                          "<b>Warning</b>"
                          "\n"
                          "The label for this knob has been changed - however the knob name is the same as in 13.0"
                          " - in a future major release the knob name will also be changed.");


    DD::Image::Colorspace_knob(f, &m_outputColorSpaceIndex, OCIO::ROLE_SCENE_LINEAR, "out_colorspace", "output");
    DD::Image::Tooltip(f, "\nImage data is converted to this color transform for output. \n\n"
                          "The first menu level is the color space family name. \n\n"
                          "The sub menu items are the names of the color spaces themselves."
                          "\n\n"
                          "<b>Warning</b>"
                          "\n"
                          "The label for this knob has been changed - however the knob name is the same as in 13.0"
                          " - in a future major release the knob name will also be changed.");

    // Internal hidden knob, used to do the reverse transform (out_colorspace -> in_colorspace).
    // This is to facilitate the Write Node's Read from Write functionality
    DD::Image::Bool_knob(f, &m_reverseTransform, "reverse_transform", "reverse transform");
    DD::Image::SetFlags(f, DD::Image::Knob::HIDDEN | DD::Image::Knob::DO_NOT_WRITE);

    DD::Image::Button(f, "swap", "swap input/output");
    DD::Image::SetFlags(f, DD::Image::Knob::STARTLINE);
}

int OCIOColorSpace::knob_changed(DD::Image::Knob* k)
{
  if (k->is("swap")) {
    const bool inColorspaceError = (knob("in_colorspace")->enumerationKnob()->getError() != nullptr);
    const bool outColorspaceError = (knob("out_colorspace")->enumerationKnob()->getError() != nullptr);

    if (!inColorspaceError && !outColorspaceError) {
      const int inColorspaceValue = knob("in_colorspace")->get_value();
      const int outColorspaceValue = knob("out_colorspace")->get_value();
      
      knob("in_colorspace")->set_value(outColorspaceValue);
      knob("out_colorspace")->set_value(inColorspaceValue);
    }
  }

  return 1;
}

void OCIOColorSpace::append(DD::Image::Hash& localhash)
{
    // TODO: Hang onto the context, what if getting it
    // (and querying getCacheID) is expensive?
    try {
        OCIO::ConstConfigRcPtr config = getConfigForProject();
        OCIO::ConstContextRcPtr context = getLocalContext();    // The first thing this does is call getConfigForProject!
        std::string configCacheID = config->getCacheID(context);
        localhash.append(configCacheID);
    }
    catch(OCIO::Exception &e) {
        error(e.what());
        return;
    }
}

void OCIOColorSpace::updateOCIOProcessors()
{
  const std::string inputName = knob("in_colorspace")->enumerationKnob()->getSelectedItemString();
  const std::string outputName = knob("out_colorspace")->enumerationKnob()->getSelectedItemString();

  // First error on invalid color space settings. This may happen with deprecated colour spaces.
  // Note that it is possible to send invalid colour spaces to OCIO but the error messages returned
  // are not user friendly. See TP 387269 for details.
  if (knob("in_colorspace")->enumerationKnob()->getError()) {
      std::stringstream errorMsg;
      errorMsg << "Invalid input LUT selected: " << inputName << std::endl;
      error(errorMsg.str().c_str());
      throw OCIO::Exception(errorMsg.str().c_str());
  }
  if (knob("out_colorspace")->enumerationKnob()->getError()) {
      std::stringstream errorMsg;
      errorMsg << "Invalid output LUT selected: " << outputName << std::endl;
      error(errorMsg.str().c_str());
      throw OCIO::Exception(errorMsg.str().c_str());
  }

  OCIO::ConstConfigRcPtr config = getConfigForProject();
  OCIO::ConstContextRcPtr context = getLocalContext();

  m_processor = config->getProcessor(context, inputName.c_str(), outputName.c_str());
  m_cpuProcessor = m_processor->getDefaultCPUProcessor();

  if (m_reverseTransform) {
    m_reverseProcessor = config->getProcessor(context, outputName.c_str(), inputName.c_str());
    m_cpuReverseProcessor = m_reverseProcessor->getDefaultCPUProcessor();
  }
}

void OCIOColorSpace::_validate(bool for_real)
{
    try {
        updateOCIOProcessors();
        setupGPUShaders(m_processor);

        if (m_reverseTransform) {
          setupGPUShaders(m_reverseProcessor);
        }
    }
    catch (OCIO::Exception &e) {
        error(e.what());
        return;
    }

    if (m_processor->isNoOp()) {
        set_out_channels(DD::Image::Mask_None); // prevents engine() from being called
    }
    else {
        // Tell Nuke that we're going to possibly operate on all channels.
        // The caveat here is that this Op is contained within a NukeWrapper
        // object which allows the user to specify the channels they want this
        // object to operate on. So, here we tell Nuke we will operate on
        // anything, but in pixel_engine() we pass through any non RBG
        // channels,
        set_out_channels(DD::Image::Mask_All);
    }

    // Call through to the base implementation
    DD::Image::OCIOPluginBase::_validate(for_real);
}

// Note that this is copied by others (OCIODisplay)
void OCIOColorSpace::in_channels(int /* n unused */, DD::Image::ChannelSet& mask) const
{
    DD::Image::ChannelSet done;
    foreach (c, mask) {
      const bool isRGB = DD::Image::colourIndex(c) < 3;
      if (isRGB) {
        const bool have = done & c;
        if (!have) {
          // because the colour conversion needs R, G *and* B from any
          // layer-channel-set, expand the channels list to include the other
          // corresponding R, G or B channels for that.
          // i.e. if 'mask' has a layer 'CustomLayer' and only contain has R &
          // G, after this call it will also contain B
          done.addBrothers(c, 3);
        }
      }
    }
    mask += done;
}

// See Saturation::pixel_engine for a well-commented example.
// Note that this is copied by others (OCIODisplay)
void OCIOColorSpace::pixel_engine(
    const DD::Image::Row& in,
    int /* rowY unused */, int rowX, int rowXBound,
    DD::Image::ChannelMask outputChannels,
    DD::Image::Row& out)
{
    int rowWidth = rowXBound - rowX;

    in.debug( rowX, rowXBound, outputChannels );

    DD::Image::ChannelSet done;
    foreach (requestedChannel, outputChannels) {
        // Skip channels which had their trios processed already,
        if (done & requestedChannel) {
            continue;
        }

        // Pass through channels which are not selected for processing
        // and non-rgb channels.
        // This operates on Layer RGB channels e.g. layer.r, layer.g etc...
        if (requestedChannel == DD::Image::Chan_Mask
            || colourIndex(requestedChannel) >= 3) {
            out.copy(in, requestedChannel, rowX, rowXBound);
            continue;
        }

        DD::Image::Channel rChannel = DD::Image::brother(requestedChannel, 0);
        DD::Image::Channel gChannel = DD::Image::brother(requestedChannel, 1);
        DD::Image::Channel bChannel = DD::Image::brother(requestedChannel, 2);

        done += rChannel;
        done += gChannel;
        done += bChannel;

        // IMPORTANT: must be before we aquire the out-pointers as the Rows may
        //            be co-locational i.e. &in === &out, and the data on the
        //            in-row is a borrowed, read-only reference and may not be
        //            in the out Row's writable channels
        DD::Image::Row::ReadablePtr rIn = in[rChannel] + rowX;
        DD::Image::Row::ReadablePtr gIn = in[gChannel] + rowX;
        DD::Image::Row::ReadablePtr bIn = in[bChannel] + rowX;

        DD::Image::Row::WritablePtr rOut = out.writable(rChannel) + rowX;
        DD::Image::Row::WritablePtr gOut = out.writable(gChannel) + rowX;
        DD::Image::Row::WritablePtr bOut = out.writable(bChannel) + rowX;

        mFnAssertMsg(rOut != gOut && rOut != bOut && gOut != bOut,
            "channel outputs, where we will do the in-place modification must "
            "be unique as we get channel cross-talk in OCIO" );

#if FN_DEBUG
        if (    (rOut && rOut<(float*)0x400)
             || (gOut && gOut<(float*)0x400)
             || (bOut && bOut<(float*)0x400) ) {
          error( "Outputs returned invalid memory, check that the correct get/request has been called" );
          mFnAssert( false );
          return;
        }

        if (    (rIn && rIn<(float*)0x400)
             || (gIn && gIn<(float*)0x400)
             || (bIn && bIn<(float*)0x400) ) {
          error( "Outputs returned invalid memory, check that the correct get/request has been called" );
          mFnAssert( false );
          return;
        }
#endif // FN_DEBUG

        out.debug( rowX, rowXBound, outputChannels );

        // OCIO modifies in-place
        // Note: xOut can equal xIn in some circumstances, such as when the
        // 'Black' (throwaway) scanline is used. We thus must guard memcpy,
        // which does not allow for overlapping regions.
        // We must also guard against r/g/b In being NULL, which can happen
        // when an OCIOColorspace is used internally with a Read node, in which
        // case the Node has no input. In this case, we're passing the same row
        // reference for in and out, so the memcpy isn't required anyway.
        if (rIn && rOut != rIn) memcpy(rOut, rIn, sizeof(float)*rowWidth);
        if (gIn && gOut != gIn) memcpy(gOut, gIn, sizeof(float)*rowWidth);
        if (bIn && bOut != bIn) memcpy(bOut, bIn, sizeof(float)*rowWidth);

        out.debug( rowX, rowXBound, outputChannels );

        try {
            OCIO::PlanarImageDesc img(rOut, gOut, bOut, nullptr, rowWidth, /*height*/ 1);

            if (m_reverseTransform) {
              applyImageCPU(m_cpuReverseProcessor, img);
            }
            else {
              applyImageCPU(m_cpuProcessor, img);
            }
        }
        catch(OCIO::Exception &e) {
            error(e.what());
        }

        out.debug( rowX, rowXBound, outputChannels );
    }
}

const DD::Image::Op::Description OCIOColorSpace::description("OCIOColorSpace", build);

const char* OCIOColorSpace::Class() const
{
    return description.name;
}

const char* OCIOColorSpace::displayName() const
{
    return description.name;
}

const char* OCIOColorSpace::node_help() const
{
    // TODO more detailed help text
    return "Use OpenColorIO to convert from one color space to another.";
}

// This class is necessary in order to call knobsAtTheEnd(). Otherwise, the NukeWrapper knobs
// will be added to the Context tab instead of the primary tab.
class OCIOColorSpaceNukeWrapper : public DD::Image::NukeWrapper
{
public:
    OCIOColorSpaceNukeWrapper(DD::Image::PixelIop* op) : DD::Image::NukeWrapper(op)
    {
    }

    void attach() override
    {
        wrapped_iop()->attach();
    }

    void detach() override
    {
        wrapped_iop()->detach();
    }

    void knobs(DD::Image::Knob_Callback f) override
    {
        OCIOColorSpace* csIop = dynamic_cast<OCIOColorSpace*>(wrapped_iop());
        if(!csIop) {
          return;
        }

        DD::Image::NukeWrapper::knobs(f);

        csIop->getExtendedOpProperties()->extendedPropertiesKnobs(f, "Context", nullptr, DD::Image::OCIOPluginBase::kContextTooltip);
    }

    std::string getLibraryName() const override { return wrapped_iop()->getLibraryName(); }
    Op::VersionInfo getLibraryVersion() const override { return wrapped_iop()->getLibraryVersion(); }
};

static DD::Image::Op* build(Node *node)
{
    DD::Image::NukeWrapper *op = (new OCIOColorSpaceNukeWrapper(new OCIOColorSpace(node)));
    op->channels(DD::Image::Mask_RGB);
    return op;
}