//=========================================================================
//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt 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 "smtk/session/aeva/operators/Export.h"

#include "smtk/session/aeva/Export_xml.h"
#include "smtk/session/aeva/Resource.h"
#include "smtk/session/aeva/Session.h"
#include "vtk/aeva/ext/vtkMedWriter.h"

#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/ComponentItem.h"
#include "smtk/attribute/FileItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/Resource.h"
#include "smtk/attribute/ResourceItem.h"
#include "smtk/attribute/StringItem.h"

#include "smtk/model/Face.h"
#include "smtk/model/Group.h"
#include "smtk/model/Model.h"
#include "smtk/model/Volume.h"

#include "smtk/resource/Manager.h"

#include "smtk/common/Paths.h"

#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkSmartPointer.h"
#include "vtkUnstructuredGrid.h"
#include "vtkXMLImageDataWriter.h"

#ifdef ERROR
#undef ERROR
#endif

using ResourceArray = std::vector<smtk::session::aeva::Resource::Ptr>;

namespace smtk
{
namespace session
{
namespace aeva
{

Export::Result Export::operateInternal()
{
  smtk::attribute::FileItem::Ptr filenameItem = this->parameters()->findFile("filename");

  std::string filename = filenameItem->value();

  // Infer file type from name
  std::string ext = smtk::common::Paths::extension(filename);
  // Downcase the extension
  std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
  if (ext == ".med")
  {
    return this->exportMedMesh();
  }

  // Internal storage is a VTK image, default to that:
  return this->exportVTKImage();
}

Export::Result Export::exportVTKImage()
{
  auto inputItem = this->parameters()->associations();
  auto inputValue = inputItem->valueAs<smtk::model::Entity>(0);
  smtk::session::aeva::Resource::Ptr resource;
  if (inputValue)
  {
    resource = std::static_pointer_cast<smtk::session::aeva::Resource>(inputValue->resource());
  }
  if (!resource)
  {
    smtkErrorMacro(this->log(), "No resources to export.");
    return this->createResult(smtk::operation::Operation::Outcome::FAILED);
  }

  smtk::common::UUID imageId;
  if (inputValue->isModel())
  {
    // retrieve the first cell, containing the volume image to save.
    smtk::model::Model dataset = smtk::model::Model(resource, inputValue->id());
    if (!dataset.cells().empty())
    {
      imageId = dataset.cells()[0].entity();
    }
  }
  else if (inputValue->isCellEntity())
  {
    imageId = inputValue->id();
  }

  auto const& session = resource->session();
  auto* image = vtkImageData::SafeDownCast(session->findStorage(imageId));
  if (!image)
  {
    smtkErrorMacro(this->log(), "Cannot retrieve image from model.");
    return this->createResult(smtk::operation::Operation::Outcome::FAILED);
  }

  vtkNew<vtkXMLImageDataWriter> writer;
  // force file extension to match what we're writing.
  std::string filename = smtk::common::Paths::replaceExtension(
    this->parameters()->findFile("filename")->value(), ".vti");

  writer->SetFileName(filename.c_str());
  writer->SetInputData(image);
  writer->Write();
  auto result = this->createResult(Export::Outcome::SUCCEEDED);
  return result;
}

Export::Result Export::exportMedMesh()
{
  auto inputItem = this->parameters()->associations();
  auto inputValue = inputItem->valueAs<smtk::model::Entity>(0);
  smtk::session::aeva::Resource::Ptr resource = nullptr;
  smtk::session::aeva::Session::Ptr session = nullptr;
  m_result = this->createResult(Export::Outcome::FAILED);

  this->prepareResourceAndSession(m_result, resource, session, false);
  if (!resource)
  {
    smtkErrorMacro(this->log(), "No resources to export.");
    return m_result;
  }
  // organize data from the resource in a multiblock as the writer expects.
  vtkSmartPointer<vtkMultiBlockDataSet> readerInput = vtkMultiBlockDataSet::New();
  vtkSmartPointer<vtkMultiBlockDataSet> geomInput = vtkMultiBlockDataSet::New();
  readerInput->SetBlock(0, geomInput);
  vtkSmartPointer<vtkMultiBlockDataSet> groupInput = vtkMultiBlockDataSet::New();
  readerInput->SetBlock(1, groupInput);

  // TODO: Currently available med files start with a single surface,
  // and an optional encoded tetra volume. This may not be true in the future.
  // We always create a volume on import, even if it's empty.
  // Look for the volume and surface, add to the geom MB.
  // Side-sets, or groups, are marked with a property flag and added to the group MB.
  smtk::common::UUID volId;
  if (!inputValue->isModel())
  {
    smtkErrorMacro(this->log(), "Expected model as input to export.");
    return m_result;
  }
  smtk::model::Model dataset = smtk::model::Model(resource, inputValue->id());
  vtkSmartPointer<vtkUnstructuredGrid> volume;
  std::string volumeName;
  vtkSmartPointer<vtkUnstructuredGrid> surface;
  std::string surfaceName;
  for (auto const& cellRef : dataset.cells())
  {
    auto const& cellId = cellRef.entity();
    auto* unstructuredGrid = vtkUnstructuredGrid::SafeDownCast(session->findStorage(cellId));
    if (unstructuredGrid && Session::isSideSet(*cellRef.component()))
    {
      // add to the group multi-block
      auto i = groupInput->GetNumberOfBlocks();
      groupInput->SetBlock(i, unstructuredGrid);
      groupInput->GetMetaData(i)->Set(vtkCompositeDataSet::NAME(), cellRef.name());
    }
    else if (!volume && cellRef.dimension() == 3 &&
      (!unstructuredGrid || unstructuredGrid->GetCellType(0) == VTK_TETRA))
    {
      // unstructuredGrid can be empty, if the surface encloses a rigid volume with no tetra mesh.
      volume = unstructuredGrid;
      volumeName = cellRef.name();
      // retrieve the surface from the volume relations
      for (auto const& relation : cellRef.relations())
      {
        if (relation.dimension() == 2)
        {
          unstructuredGrid =
            vtkUnstructuredGrid::SafeDownCast(session->findStorage(relation.entity()));
          if (unstructuredGrid && unstructuredGrid->GetCellType(0) == VTK_TRIANGLE)
          {
            surface = unstructuredGrid;
            surfaceName = relation.name();
          }
        }
      }
    }
    else if (unstructuredGrid && !surface && cellRef.dimension() == 2 &&
      unstructuredGrid->GetCellType(0) == VTK_TRIANGLE)
    {
      surface = unstructuredGrid;
      surfaceName = cellRef.name();
    }
  }

  if (!surface)
  {
    smtkErrorMacro(this->log(), "Cannot retrieve surface from model.");
    return m_result;
  }
  unsigned int block = 0;
  if (volume)
  {

    geomInput->SetBlock(block, volume);
    geomInput->GetMetaData(block)->Set(vtkCompositeDataSet::NAME(), volumeName);
    ++block;
  }
  geomInput->SetBlock(block, surface);
  geomInput->GetMetaData(block)->Set(vtkCompositeDataSet::NAME(), surfaceName);

  vtkNew<vtkMedWriter> writer;
  // force file extension to match what we're writing.
  std::string filename = smtk::common::Paths::replaceExtension(
    this->parameters()->findFile("filename")->value(), ".med");

  writer->SetFileName(filename);
  writer->SetInputData(readerInput);
  writer->Update();

  m_result->findInt("outcome")->setValue(0, static_cast<int>(Export::Outcome::SUCCEEDED));
  return m_result;
}

const char* Export::xmlDescription() const
{
  return Export_xml;
}

bool exportResource(const smtk::resource::ResourcePtr& resource)
{
  Export::Ptr exportResource = Export::create();
  exportResource->parameters()->findResource("resource")->setValue(resource);
  Export::Result result = exportResource->operate();
  return (result->findInt("outcome")->value() == static_cast<int>(Export::Outcome::SUCCEEDED));
}
}
}
}
