/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkDataEncoder.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkDataEncoder.h"

#include "vtkBase64Utilities.h"
#include "vtkCommand.h"
#include "vtkImageData.h"
#include "vtkJPEGWriter.h"
#include "vtkMultiThreader.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPNGWriter.h"
#include "vtkUnsignedCharArray.h"

#include <cassert>
#include <cmath>
#include <condition_variable>
#include <map>
#include <mutex>
#include <vector>

#include <vtksys/SystemTools.hxx>

#define MAX_NUMBER_OF_THREADS_IN_POOL 32
//****************************************************************************
namespace
{
class vtkSharedData
{
public:
  struct OutputValueType
  {
  public:
    vtkTypeUInt32 TimeStamp;
    vtkSmartPointer<vtkUnsignedCharArray> Data;
    OutputValueType()
      : TimeStamp(0)
      , Data(nullptr)
    {
    }
  };

  enum
  {
    ENCODING_NONE = 0,
    ENCODING_BASE64 = 1
  };
  struct InputValueType
  {
  public:
    vtkTypeUInt32 OutputStamp;
    vtkSmartPointer<vtkImageData> Image;
    int Quality;
    int Encoding;
    InputValueType()
      : OutputStamp(0)
      , Image(nullptr)
      , Quality(100)
      , Encoding(ENCODING_BASE64)
    {
    }
  };

  typedef std::map<vtkTypeUInt32, InputValueType> InputMapType;
  typedef std::map<vtkTypeUInt32, OutputValueType> OutputMapType;

private:
  bool Done;
  std::mutex DoneLock;
  std::mutex OutputsLock;
  std::condition_variable_any OutputsAvailable;

  std::mutex ThreadDoneLock;
  std::condition_variable_any ThreadDone;
  int ActiveThreadCount;

  //------------------------------------------------------------------------
  // OutputsLock must be held before accessing any of the following members.
  OutputMapType Outputs;

  //------------------------------------------------------------------------
  // Constructs used to synchronization.
  std::mutex InputsLock;
  std::condition_variable_any InputsAvailable;

  //------------------------------------------------------------------------
  // InputsLock must be held before accessing any of the following members.
  InputMapType Inputs;

public:
  //------------------------------------------------------------------------
  vtkSharedData()
    : Done(false)
    , ActiveThreadCount(0)
  {
  }

  //------------------------------------------------------------------------
  // Each thread should call this method when it starts. It helps us clean up
  // threads when they are done.
  void BeginWorker()
  {
    this->ThreadDoneLock.lock();
    this->ActiveThreadCount++;
    this->ThreadDoneLock.unlock();
  }

  //------------------------------------------------------------------------
  // Each thread should call this method when it ends.
  void EndWorker()
  {
    this->ThreadDoneLock.lock();
    this->ActiveThreadCount--;
    bool last_thread = (this->ActiveThreadCount == 0);
    this->ThreadDoneLock.unlock();
    if (last_thread)
    {
      this->ThreadDone.notify_one();
    }
  }

  //------------------------------------------------------------------------
  void RequestAndWaitForWorkersToEnd()
  {
    // Get the done lock so we other threads don't end up testing the Done
    // flag and quitting before this thread starts to wait for them to quit.
    this->DoneLock.lock();
    this->Done = true;

    // Grab the ThreadDoneLock. so even if any thread ends up check this->Done
    // as soon as we release the lock, it won't get a chance to terminate.
    this->ThreadDoneLock.lock();

    // release the done lock. Let threads test for this->Done flag.
    this->DoneLock.unlock();

    // Tell all workers that inputs are available, so they will try to check
    // the input as well as the done flag.
    this->InputsAvailable.notify_all();

    // Now wait for thread to terminate releasing this->ThreadDoneLock as soon
    // as we start waiting. Thus, no threads have got a chance to call
    // EndWorker() till the main thread starts waiting for them.
    this->ThreadDone.wait(this->ThreadDoneLock);

    this->ThreadDoneLock.unlock();

    // reset Done flag since all threads have died.
    this->Done = false;
  }

  //------------------------------------------------------------------------
  bool IsDone()
  {
    this->DoneLock.lock();
    bool val = this->Done;
    this->DoneLock.unlock();
    return val;
  }

  //------------------------------------------------------------------------
  void PushAndTakeReference(
    vtkTypeUInt32 key, vtkImageData*& data, vtkTypeUInt64 stamp, int quality, int encoding)
  {
    this->InputsLock.lock();
    {
      vtkSharedData::InputValueType& value = this->Inputs[key];
      value.Image.TakeReference(data);
      value.OutputStamp = stamp;
      value.Quality = quality;
      value.Encoding = encoding;
      data = nullptr;
    }
    this->InputsLock.unlock();
    this->InputsAvailable.notify_one();
  }

  //------------------------------------------------------------------------
  vtkTypeUInt64 GetExpectedOutputStamp(vtkTypeUInt32 key)
  {
    vtkTypeUInt64 stamp = 0;
    this->InputsLock.lock();
    vtkSharedData::InputMapType::iterator iter = this->Inputs.find(key);
    if (iter != this->Inputs.end())
    {
      stamp = iter->second.OutputStamp;
    }
    this->InputsLock.unlock();
    return stamp;
  }

  //------------------------------------------------------------------------
  // NOTE: This method may suspend the calling thread until inputs become
  // available.
  vtkTypeUInt64 GetNextInputToProcess(
    vtkTypeUInt32& key, vtkSmartPointer<vtkImageData>& image, int& quality, int& encoding)
  {
    vtkTypeUInt32 stamp = 0;

    this->InputsLock.lock();
    do
    {
      // Check if we have an input available, if so, return it.
      InputMapType::iterator iter;
      for (iter = this->Inputs.begin(); iter != this->Inputs.end(); ++iter)
      {
        if (iter->second.Image != nullptr)
        {
          key = iter->first;
          image = iter->second.Image;
          iter->second.Image = nullptr;
          stamp = iter->second.OutputStamp;
          quality = iter->second.Quality;
          encoding = iter->second.Encoding;
          break;
        }
      }
      if (image == nullptr && !this->IsDone())
      {
        // No data is available, let's wait till it becomes available.
        this->InputsAvailable.wait(this->InputsLock);
      }

    } while (image == nullptr && !this->IsDone());

    this->InputsLock.unlock();
    return stamp;
  }

  //------------------------------------------------------------------------
  void SetOutputReference(
    const vtkTypeUInt32& key, vtkTypeUInt64 timestamp, vtkUnsignedCharArray*& dataRef)
  {
    this->OutputsLock.lock();
    assert(dataRef->GetReferenceCount() == 1);
    OutputMapType::iterator iter = this->Outputs.find(key);
    if (iter == this->Outputs.end() || iter->second.Data == nullptr ||
      iter->second.TimeStamp < timestamp)
    {
      // cout << "Done: " <<
      //  vtkMultiThreader::GetCurrentThreadID() << " "
      //  << key << ", " << timestamp << endl;
      this->Outputs[key].TimeStamp = timestamp;
      this->Outputs[key].Data.TakeReference(dataRef);
      dataRef = nullptr;
    }
    else
    {
      dataRef->Delete();
      dataRef = nullptr;
    }
    this->OutputsLock.unlock();
    this->OutputsAvailable.notify_all();
  }

  //------------------------------------------------------------------------
  bool CopyLatestOutputIfDifferent(vtkTypeUInt32 key, vtkUnsignedCharArray* data)
  {
    vtkTypeUInt64 dataTimeStamp = 0;
    this->OutputsLock.lock();
    {
      const vtkSharedData::OutputValueType& output = this->Outputs[key];
      if (output.Data != nullptr &&
        (output.Data->GetMTime() > data->GetMTime() ||
          output.Data->GetNumberOfTuples() != data->GetNumberOfTuples()))
      {
        data->DeepCopy(output.Data);
        data->Modified();
      }
      dataTimeStamp = output.TimeStamp;
    }
    this->OutputsLock.unlock();

    vtkTypeUInt64 outputTS = 0;

    this->InputsLock.lock();
    vtkSharedData::InputMapType::iterator iter = this->Inputs.find(key);
    if (iter != this->Inputs.end())
    {
      outputTS = iter->second.OutputStamp;
    }
    this->InputsLock.unlock();

    return (dataTimeStamp >= outputTS);
  }

  //------------------------------------------------------------------------
  void Flush(vtkTypeUInt32 key, vtkTypeUInt64 timestamp)
  {
    this->OutputsLock.lock();
    while (this->Outputs[key].TimeStamp < timestamp)
    {
      // output is not yet ready, we have to wait.
      this->OutputsAvailable.wait(this->OutputsLock);
    }

    this->OutputsLock.unlock();
  }
};

VTK_THREAD_RETURN_TYPE Worker(void* calldata)
{
  // cout << "Start Thread: " << vtkMultiThreader::GetCurrentThreadID() << endl;
  vtkMultiThreader::ThreadInfo* info = reinterpret_cast<vtkMultiThreader::ThreadInfo*>(calldata);
  vtkSharedData* sharedData = reinterpret_cast<vtkSharedData*>(info->UserData);

  sharedData->BeginWorker();

  while (true)
  {
    vtkTypeUInt32 key = 0;
    vtkSmartPointer<vtkImageData> image;
    vtkTypeUInt64 timestamp = 0;
    // these are defaults, reset by GetNextInputToProcess
    int quality = 100;
    int encoding = 1;

    timestamp = sharedData->GetNextInputToProcess(key, image, quality, encoding);

    if (timestamp == 0 || image == nullptr)
    {
      // end thread.
      break;
    }

    // cout << "Working Thread: " << vtkMultiThreader::GetCurrentThreadID() << endl;

    // Do the encoding.
    vtkNew<vtkJPEGWriter> writer;
    writer->WriteToMemoryOn();
    writer->SetInputData(image);
    writer->SetQuality(quality);
    writer->Write();
    vtkUnsignedCharArray* data = writer->GetResult();

    vtkUnsignedCharArray* result = vtkUnsignedCharArray::New();
    if (encoding)
    {
      result->SetNumberOfComponents(1);
      result->SetNumberOfTuples(std::ceil(1.5 * data->GetNumberOfTuples()));
      unsigned long size = vtkBase64Utilities::Encode(
        data->GetPointer(0), data->GetNumberOfTuples(), result->GetPointer(0), /*mark_end=*/0);
      result->SetNumberOfTuples(static_cast<vtkIdType>(size) + 1);
      result->SetValue(size, 0);
    }
    else
    {
      result->ShallowCopy(data);
    }
    // Pass over the "result" reference.
    sharedData->SetOutputReference(key, timestamp, result);
    assert(result == nullptr);
  }

  // cout << "Closing Thread: " << vtkMultiThreader::GetCurrentThreadID() << endl;
  sharedData->EndWorker();
  return VTK_THREAD_RETURN_VALUE;
}
}

//****************************************************************************
class vtkDataEncoder::vtkInternals
{
private:
  std::map<vtkTypeUInt32, vtkSmartPointer<vtkUnsignedCharArray>> ClonedOutputs;
  std::vector<int> RunningThreadIds;

public:
  vtkNew<vtkMultiThreader> Threader;
  vtkSharedData SharedData;
  vtkTypeUInt64 Counter;

  vtkSmartPointer<vtkUnsignedCharArray> lastBase64Image;

  vtkInternals()
    : Counter(0)
  {
    lastBase64Image = vtkSmartPointer<vtkUnsignedCharArray>::New();
  }

  void TerminateAllWorkers()
  {
    // request and wait for all threads to close.
    if (!this->RunningThreadIds.empty())
    {
      this->SharedData.RequestAndWaitForWorkersToEnd();
    }

    // Stop threads
    while (!this->RunningThreadIds.empty())
    {
      this->Threader->TerminateThread(this->RunningThreadIds.back());
      this->RunningThreadIds.pop_back();
    }
  }

  void SpawnWorkers(vtkTypeUInt32 numberOfThreads)
  {
    for (vtkTypeUInt32 cc = 0; cc < numberOfThreads; cc++)
    {
      this->RunningThreadIds.push_back(this->Threader->SpawnThread(&Worker, &this->SharedData));
    }
  }

  // Since changes to vtkObjectBase::ReferenceCount are not thread safe, we have
  // this level of indirection between the outputs stored in SharedData and
  // passed back to the user/main thread.
  bool GetLatestOutput(vtkTypeUInt32 key, vtkSmartPointer<vtkUnsignedCharArray>& data)
  {
    vtkSmartPointer<vtkUnsignedCharArray>& output = this->ClonedOutputs[key];
    if (!output)
    {
      output = vtkSmartPointer<vtkUnsignedCharArray>::New();
    }
    data = output;
    return this->SharedData.CopyLatestOutputIfDifferent(key, data);
  }

  // Once an imagedata has been written to memory as a jpg or png, this
  // convenience function can encode that image as a Base64 string.
  const char* GetBase64EncodedImage(vtkUnsignedCharArray* encodedInputImage)
  {
    this->lastBase64Image->SetNumberOfComponents(1);
    this->lastBase64Image->SetNumberOfTuples(
      std::ceil(1.5 * encodedInputImage->GetNumberOfTuples()));
    unsigned long size = vtkBase64Utilities::Encode(encodedInputImage->GetPointer(0),
      encodedInputImage->GetNumberOfTuples(), this->lastBase64Image->GetPointer(0), /*mark_end=*/0);

    this->lastBase64Image->SetNumberOfTuples(static_cast<vtkIdType>(size) + 1);
    this->lastBase64Image->SetValue(size, 0);

    return reinterpret_cast<char*>(this->lastBase64Image->GetPointer(0));
  }
};

vtkStandardNewMacro(vtkDataEncoder);
//------------------------------------------------------------------------------
vtkDataEncoder::vtkDataEncoder()
  : Internals(new vtkInternals())
{
  this->MaxThreads = 3;
  this->Initialize();
}

//------------------------------------------------------------------------------
vtkDataEncoder::~vtkDataEncoder()
{
  this->Internals->TerminateAllWorkers();
  delete this->Internals;
  this->Internals = nullptr;
}

//------------------------------------------------------------------------------
void vtkDataEncoder::SetMaxThreads(vtkTypeUInt32 maxThreads)
{
  if (maxThreads < MAX_NUMBER_OF_THREADS_IN_POOL && maxThreads > 0)
  {
    this->MaxThreads = maxThreads;
  }
}

//------------------------------------------------------------------------------
void vtkDataEncoder::Initialize()
{
  this->Internals->TerminateAllWorkers();
  this->Internals->SpawnWorkers(this->MaxThreads);
}

//------------------------------------------------------------------------------
void vtkDataEncoder::PushAndTakeReference(
  vtkTypeUInt32 key, vtkImageData*& data, int quality, int encoding)
{
  // if data->ReferenceCount != 1, it means the caller thread is keep an extra
  // reference and that's bad.
  assert(data->GetReferenceCount() == 1);

  this->Internals->SharedData.PushAndTakeReference(
    key, data, ++this->Internals->Counter, quality, encoding);
  assert(data == nullptr);
}

//------------------------------------------------------------------------------
bool vtkDataEncoder::GetLatestOutput(vtkTypeUInt32 key, vtkSmartPointer<vtkUnsignedCharArray>& data)
{
  return this->Internals->GetLatestOutput(key, data);
}

//------------------------------------------------------------------------------
const char* vtkDataEncoder::EncodeAsBase64Png(vtkImageData* img, int compressionLevel)
{
  // Perform in-memory write of image as png
  vtkNew<vtkPNGWriter> writer;
  writer->WriteToMemoryOn();
  writer->SetInputData(img);
  writer->SetCompressionLevel(compressionLevel);
  writer->Write();

  // Return Base64-encoded string
  return this->Internals->GetBase64EncodedImage(writer->GetResult());
}

//------------------------------------------------------------------------------
const char* vtkDataEncoder::EncodeAsBase64Jpg(vtkImageData* img, int quality)
{
  // Perform in-memory write of image as jpg
  vtkNew<vtkJPEGWriter> writer;
  writer->WriteToMemoryOn();
  writer->SetInputData(img);
  writer->SetQuality(quality);
  writer->Write();

  // Return Base64-encoded string
  return this->Internals->GetBase64EncodedImage(writer->GetResult());
}

//------------------------------------------------------------------------------
void vtkDataEncoder::Flush(vtkTypeUInt32 key)
{
  vtkTypeUInt64 outputTS = this->Internals->SharedData.GetExpectedOutputStamp(key);
  if (outputTS != 0)
  {
    // Now wait till we see the outputTS in the output for key.
    // cout << "Wait till : " << outputTS << endl;
    this->Internals->SharedData.Flush(key, outputTS);
  }
}

//------------------------------------------------------------------------------
void vtkDataEncoder::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
}

//------------------------------------------------------------------------------
void vtkDataEncoder::Finalize()
{
  this->Internals->TerminateAllWorkers();
}
