// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkJSONDataSetWriter.h"

#include "vtkArchiver.h"
#include "vtkArrayDispatch.h"
#include "vtkCellData.h"
#include "vtkDataArray.h"
#include "vtkDataSet.h"
#include "vtkDataSetAttributes.h"
#include "vtkIdTypeArray.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkMatrix3x3.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkSmartPointer.h"

#include "vtksys/FStream.hxx"
#include "vtksys/MD5.h"
#include "vtksys/SystemTools.hxx"

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

//------------------------------------------------------------------------------
VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkJSONDataSetWriter);
vtkCxxSetObjectMacro(vtkJSONDataSetWriter, Archiver, vtkArchiver);

//------------------------------------------------------------------------------
vtkJSONDataSetWriter::vtkJSONDataSetWriter()
{
  this->Archiver = vtkArchiver::New();
  this->ValidStringCount = 1;
  this->GetPointArraySelection()->SetUnknownArraySetting(1);
  this->GetCellArraySelection()->SetUnknownArraySetting(1);
}

//------------------------------------------------------------------------------
vtkJSONDataSetWriter::~vtkJSONDataSetWriter()
{
  this->SetArchiver(nullptr);
}

//------------------------------------------------------------------------------

vtkDataSet* vtkJSONDataSetWriter::GetInput()
{
  return vtkDataSet::SafeDownCast(this->Superclass::GetInput());
}

//------------------------------------------------------------------------------

vtkDataSet* vtkJSONDataSetWriter::GetInput(int port)
{
  return vtkDataSet::SafeDownCast(this->Superclass::GetInput(port));
}

//------------------------------------------------------------------------------

std::string vtkJSONDataSetWriter::WriteDataSetAttributes(
  vtkDataSetAttributes* fields, const char* className)
{
  int nbArrayWritten = 0;
  vtkIdType activeTCoords = -1;
  vtkIdType activeScalars = -1;
  vtkIdType activeNormals = -1;
  vtkIdType activeGlobalIds = -1;
  vtkIdType activeTensors = -1;
  vtkIdType activePedigreeIds = -1;
  vtkIdType activeVectors = -1;

  vtkIdType nbFields = fields->GetNumberOfArrays();

  if (nbFields == 0)
  {
    return "";
  }

  std::stringstream jsonSnippet;
  jsonSnippet << "  \"" << className << "\": {"
              << "\n    \"vtkClass\": \"vtkDataSetAttributes\","
              << "\n    \"arrays\": [\n";
  for (vtkIdType idx = 0; idx < nbFields; idx++)
  {
    vtkDataArray* field = fields->GetArray(idx);
    if (field == nullptr)
    {
      continue;
    }

    if (std::string(className) == "pointData")
    {
      if (!this->GetPointArraySelection()->ArrayIsEnabled(field->GetName()))
      {
        vtkDebugMacro("Skipping writing point array " << field->GetName());
        continue;
      }
    }
    else if (std::string(className) == "cellData")
    {
      if (!this->GetCellArraySelection()->ArrayIsEnabled(field->GetName()))
      {
        vtkDebugMacro("Skipping cell array " << field->GetName());
        continue;
      }
    }

    if (nbArrayWritten)
    {
      jsonSnippet << ",\n";
    }

    jsonSnippet << "      { \"data\": " << this->WriteArray(field, "vtkDataArray") << "}";

    // Update active field if any
    activeTCoords = field == fields->GetTCoords() ? nbArrayWritten : activeTCoords;
    activeScalars = field == fields->GetScalars() ? nbArrayWritten : activeScalars;
    activeNormals = field == fields->GetNormals() ? nbArrayWritten : activeNormals;
    activeGlobalIds = field == fields->GetGlobalIds() ? nbArrayWritten : activeGlobalIds;
    activeTensors = field == fields->GetTensors() ? nbArrayWritten : activeTensors;
    activePedigreeIds = field == fields->GetPedigreeIds() ? nbArrayWritten : activePedigreeIds;
    activeVectors = field == fields->GetVectors() ? nbArrayWritten : activeVectors;

    // Increment the number of array currently in the list
    nbArrayWritten++;
  }
  jsonSnippet << "\n    ],\n"
              << "    \"activeTCoords\": " << activeTCoords << ",\n"
              << "    \"activeScalars\": " << activeScalars << ",\n"
              << "    \"activeNormals\": " << activeNormals << ",\n"
              << "    \"activeGlobalIds\": " << activeGlobalIds << ",\n"
              << "    \"activeTensors\": " << activeTensors << ",\n"
              << "    \"activePedigreeIds\": " << activePedigreeIds << ",\n"
              << "    \"activeVectors\": " << activeVectors << "\n"
              << "  }";

  return jsonSnippet.str();
}

//------------------------------------------------------------------------------

std::string vtkJSONDataSetWriter::WriteArray(
  vtkDataArray* array, const char* className, const char* arrayName)
{
  bool needConvert;
  std::string id = vtkJSONDataSetWriter::GetUID(array, needConvert);
  std::stringstream arrayPath;
  arrayPath << "data/" << id;
  bool success = vtkJSONDataSetWriter::WriteArrayContents(array, arrayPath.str().c_str());

  if (!success)
  {
    return "{}";
  }

  const char* INDENT = "    ";
  std::stringstream ss;
  ss << "{\n"
     << INDENT << "  \"vtkClass\": \"" << className << "\",\n"
     << INDENT << "  \"name\": \""
     << this->GetValidString(arrayName == nullptr ? array->GetName() : arrayName) << "\",\n"
     << INDENT << "  \"numberOfComponents\": " << array->GetNumberOfComponents() << ",\n"
     << INDENT << "  \"dataType\": \"" << vtkJSONDataSetWriter::GetShortType(array, needConvert)
     << "Array\",\n"
     << INDENT << "  \"ref\": {\n"
     << INDENT << "     \"encode\": \"LittleEndian\",\n"
     << INDENT << "     \"basepath\": \"data\",\n"
     << INDENT << "     \"id\": \"" << id << "\"\n"
     << INDENT << "  },\n"
     << INDENT << "  \"size\": " << array->GetNumberOfValues() << "\n"
     << INDENT << "}";

  return ss.str();
}

//------------------------------------------------------------------------------
void vtkJSONDataSetWriter::Write(vtkDataSet* dataset)
{
  vtkImageData* imageData = vtkImageData::SafeDownCast(dataset);
  vtkPolyData* polyData = vtkPolyData::SafeDownCast(dataset);
  this->ValidDataSet = false;

  // Get input and check data
  if (dataset == nullptr)
  {
    vtkErrorMacro(<< "No data to write!");
    return;
  }

  // Capture vtkDataSet definition
  std::stringstream metaJsonFile;
  metaJsonFile << "{\n";
  metaJsonFile << "  \"vtkClass\": \"" << dataset->GetClassName() << "\"";

  // ImageData
  if (imageData)
  {
    this->ValidDataSet = true;

    // Spacing
    metaJsonFile << ",\n  \"spacing\": [" << imageData->GetSpacing()[0] << ", "
                 << imageData->GetSpacing()[1] << ", " << imageData->GetSpacing()[2] << "]";

    // Origin
    metaJsonFile << ",\n  \"origin\": [" << imageData->GetOrigin()[0] << ", "
                 << imageData->GetOrigin()[1] << ", " << imageData->GetOrigin()[2] << "]";

    // Extent
    metaJsonFile << ",\n  \"extent\": [" << imageData->GetExtent()[0] << ", "
                 << imageData->GetExtent()[1] << ", " << imageData->GetExtent()[2] << ", "
                 << imageData->GetExtent()[3] << ", " << imageData->GetExtent()[4] << ", "
                 << imageData->GetExtent()[5] << "]";

    // Direction
    // Write the matrix using vtk.js convention for direction matrices (transpose the matrix)
    auto direction = imageData->GetDirectionMatrix()->GetData();
    metaJsonFile << ",\n  \"direction\": [" << direction[0] << ", " << direction[3] << ", "
                 << direction[6] << ", " << direction[1] << ", " << direction[4] << ", "
                 << direction[7] << ", " << direction[2] << ", " << direction[5] << ", "
                 << direction[8] << "]";
  }

  // PolyData
  if (polyData && polyData->GetPoints())
  {
    this->ValidDataSet = true;

    vtkPoints* points = polyData->GetPoints();
    metaJsonFile << ",\n  \"points\": "
                 << this->WriteArray(points->GetData(), "vtkPoints", "points");

    // Verts
    vtkNew<vtkIdTypeArray> cells;
    polyData->GetVerts()->ExportLegacyFormat(cells);
    if (cells->GetNumberOfValues())
    {
      metaJsonFile << ",\n  \"verts\": " << this->WriteArray(cells, "vtkCellArray", "verts");
    }

    // Lines
    polyData->GetLines()->ExportLegacyFormat(cells);
    if (cells->GetNumberOfValues())
    {
      metaJsonFile << ",\n  \"lines\": " << this->WriteArray(cells, "vtkCellArray", "lines");
    }

    // Strips
    polyData->GetStrips()->ExportLegacyFormat(cells);
    if (cells->GetNumberOfValues())
    {
      metaJsonFile << ",\n  \"strips\": " << this->WriteArray(cells, "vtkCellArray", "strips");
    }

    // Polys
    polyData->GetPolys()->ExportLegacyFormat(cells);
    if (cells->GetNumberOfValues())
    {
      metaJsonFile << ",\n  \"polys\": " << this->WriteArray(cells, "vtkCellArray", "polys");
    }
  }

  // PointData
  bool isEmpty = true;
  std::string fieldJSON = this->WriteDataSetAttributes(dataset->GetPointData(), "pointData");
  if (!fieldJSON.empty())
  {
    isEmpty = false;
    metaJsonFile << ",\n" << fieldJSON;
  }

  // CellData
  fieldJSON = this->WriteDataSetAttributes(dataset->GetCellData(), "cellData");
  if (!fieldJSON.empty())
  {
    isEmpty = false;
    metaJsonFile << ",\n" << fieldJSON;
  }

  metaJsonFile << "}\n";

  // Create archive only if there's actually something to write
  if (this->ValidDataSet || !isEmpty)
  {
    this->GetArchiver()->OpenArchive();
    std::string metaJsonFileStr = metaJsonFile.str();
    this->GetArchiver()->InsertIntoArchive(
      "index.json", metaJsonFileStr.c_str(), metaJsonFileStr.size());
    this->GetArchiver()->CloseArchive();
  }
}

//------------------------------------------------------------------------------
void vtkJSONDataSetWriter::WriteData()
{
  vtkDataSet* dataset = this->GetInput();
  this->Write(dataset);
}

//------------------------------------------------------------------------------
struct vtkJSONDataSetWriterFunctor
{
  template <class TArray>
  void operator()(TArray* input, vtkJSONDataSetWriter* self, const char* filePath)
  {
    using ValueType = typename TArray::ValueType;
    // Check if we need to convert the (u)int64 to (u)int32
    if (std::is_integral_v<ValueType> && sizeof(ValueType) > 4)
    {
      if (std::is_signed_v<ValueType>)
      {
        vtkNew<vtkAOSDataArrayTemplate<unsigned int>> uint32;
        uint32->DeepCopy(input);
        const char* content = reinterpret_cast<const char*>(uint32->GetPointer(0));
        size_t size = uint32->GetNumberOfValues() * uint32->GetDataTypeSize();
        self->GetArchiver()->InsertIntoArchive(filePath, content, size);
      }
      else
      {
        vtkNew<vtkAOSDataArrayTemplate<int>> int32;
        int32->DeepCopy(input);
        const char* content = reinterpret_cast<const char*>(int32->GetPointer(0));
        size_t size = int32->GetNumberOfValues() * int32->GetDataTypeSize();
        self->GetArchiver()->InsertIntoArchive(filePath, content, size);
      }
    }
    else
    {
      const char* content = reinterpret_cast<const char*>(input->GetPointer(0));
      size_t size = input->GetNumberOfValues() * input->GetDataTypeSize();
      self->GetArchiver()->InsertIntoArchive(filePath, content, size);
    }
  }
};

//------------------------------------------------------------------------------
bool vtkJSONDataSetWriter::WriteArrayContents(vtkDataArray* input, const char* filePath)
{
  auto aos = input->ToAOSDataArray();
  vtkJSONDataSetWriterFunctor functor;
  if (!vtkArrayDispatch::DispatchByArray<vtkArrayDispatch::AOSArrays>::Execute(
        aos, functor, this, filePath))
  {
    return false;
  }
  return true;
}

//------------------------------------------------------------------------------
namespace
{
class vtkSingleFileArchiver : public vtkArchiver
{
public:
  static vtkSingleFileArchiver* New();
  vtkTypeMacro(vtkSingleFileArchiver, vtkArchiver);

  void OpenArchive() override {}
  void CloseArchive() override {}
  void InsertIntoArchive(const std::string& filePath, const char* data, std::size_t size) override
  {
    vtksys::ofstream file;
    file.open(filePath.c_str(), ios::out | ios::binary);
    file.write(data, size);
    file.close();
  }

private:
  vtkSingleFileArchiver() = default;
  ~vtkSingleFileArchiver() override = default;
};
vtkStandardNewMacro(vtkSingleFileArchiver);
}

//------------------------------------------------------------------------------
bool vtkJSONDataSetWriter::WriteArrayAsRAW(vtkDataArray* array, const char* filePath)
{
  vtkNew<vtkJSONDataSetWriter> writer;
  vtkNew<vtkSingleFileArchiver> archiver;
  writer->SetArchiver(archiver);
  return writer->WriteArrayContents(array, filePath);
}

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

  os << indent << "Archiver:" << endl;
  this->Archiver->PrintSelf(os, indent.GetNextIndent());

  os << indent << "PointArraySelection:" << endl;
  this->PointArraySelection->PrintSelf(os, indent.GetNextIndent());

  os << indent << "CelltArraySelection:" << endl;
  this->CellArraySelection->PrintSelf(os, indent.GetNextIndent());
}

//------------------------------------------------------------------------------
int vtkJSONDataSetWriter::FillInputPortInformation(int, vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkDataSet");
  return 1;
}

//------------------------------------------------------------------------------
// Static helper functions
//------------------------------------------------------------------------------

void vtkJSONDataSetWriter::ComputeMD5(const unsigned char* content, int size, std::string& hash)
{
  unsigned char digest[16];
  char md5Hash[33];
  md5Hash[32] = '\0';

  vtksysMD5* md5 = vtksysMD5_New();
  vtksysMD5_Initialize(md5);
  vtksysMD5_Append(md5, content, size);
  vtksysMD5_Finalize(md5, digest);
  vtksysMD5_DigestToHex(digest, md5Hash);
  vtksysMD5_Delete(md5);

  hash = md5Hash;
}

//------------------------------------------------------------------------------

std::string vtkJSONDataSetWriter::GetShortType(vtkDataArray* input, bool& needConversion)
{
  needConversion = false;
  std::stringstream ss;
  switch (input->GetDataType())
  {
    case VTK_UNSIGNED_CHAR:
    case VTK_UNSIGNED_SHORT:
    case VTK_UNSIGNED_INT:
    case VTK_UNSIGNED_LONG:
    case VTK_UNSIGNED_LONG_LONG:
      ss << "Uint";
      if (input->GetDataTypeSize() <= 4)
      {
        ss << (input->GetDataTypeSize() * 8);
      }
      else
      {
        needConversion = true;
        ss << "32";
      }

      break;

    case VTK_CHAR:
    case VTK_SIGNED_CHAR:
    case VTK_SHORT:
    case VTK_INT:
    case VTK_LONG:
    case VTK_LONG_LONG:
    case VTK_ID_TYPE:
      ss << "Int";
      if (input->GetDataTypeSize() <= 4)
      {
        ss << (input->GetDataTypeSize() * 8);
      }
      else
      {
        needConversion = true;
        ss << "32";
      }
      break;

    case VTK_FLOAT:
    case VTK_DOUBLE:
      ss << "Float";
      ss << (input->GetDataTypeSize() * 8);
      break;

    case VTK_BIT:
    case VTK_STRING:
    case VTK_VARIANT:
    default:
      ss << "xxx";
      break;
  }

  return ss.str();
}

//------------------------------------------------------------------------------

std::string vtkJSONDataSetWriter::GetUID(vtkDataArray* input, bool& needConversion)
{
  auto aosInput = input->ToAOSDataArray();
  // NOLINTNEXTLINE(bugprone-unsafe-functions)
  const unsigned char* content = (const unsigned char*)aosInput->GetVoidPointer(0);
  int size = input->GetNumberOfValues() * input->GetDataTypeSize();
  std::string hash;
  vtkJSONDataSetWriter::ComputeMD5(content, size, hash);

  std::stringstream ss;
  ss << vtkJSONDataSetWriter::GetShortType(input, needConversion) << "_"
     << input->GetNumberOfValues() << "-" << hash;

  return ss.str();
}

//------------------------------------------------------------------------------

std::string vtkJSONDataSetWriter::GetValidString(const char* name)
{
  if (name != nullptr && strlen(name))
  {
    return name;
  }
  std::stringstream ss;
  ss << "invalid_" << this->ValidStringCount++;

  return ss.str();
}
VTK_ABI_NAMESPACE_END