//=========================================================================
//=========================================================================
//  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 "vtkAppendFilter.h"
#include "vtkCell3D.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkPolyData.h"
#include "vtkSmartPointer.h"
#include "vtkUnstructuredGrid.h"
#include "vtkXMLImageDataWriter.h"

SMTK_THIRDPARTY_PRE_INCLUDE
#include "boost/filesystem.hpp"
#include "boost/system/error_code.hpp"
#include <boost/dll.hpp>
SMTK_THIRDPARTY_POST_INCLUDE

#ifdef ERROR
#undef ERROR
#endif

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

namespace smtk
{
namespace session
{
namespace aeva
{

Export::Result Export::operateInternal()
{
  m_result = this->createResult(smtk::operation::Operation::Outcome::FAILED);

  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();
  m_result = this->createResult(Export::Outcome::FAILED);
  smtk::model::Models models;
  smtk::session::aeva::Session::Ptr session = nullptr;
  smtk::session::aeva::Resource::Ptr resource = nullptr;
  // we might get multiple resources or models as input. Extract all models.
  for (std::size_t k = 0; k < inputItem->numberOfValues(); ++k)
  {
    resource = inputItem->valueAs<smtk::session::aeva::Resource>(k);
    auto inputValue = inputItem->valueAs<smtk::model::Entity>(k);
    if (resource)
    {
      auto resModels =
        resource->entitiesMatchingFlagsAs<smtk::model::Models>(smtk::model::MODEL_ENTITY, false);
      models.insert(models.end(), resModels.begin(), resModels.end());
      if (!session)
      {
        session = resource->session();
      }
    }
    else if (inputValue)
    {
      if (!inputValue->isModel())
      {
        smtkErrorMacro(this->log(), "Expected model as input to export.");
        return m_result;
      }
      this->prepareResourceAndSession(m_result, resource, session, false);
      models.push_back(smtk::model::Model(resource, inputValue->id()));
    }
  }
  // For the list of models, find which ones have a vtkImageData as cell[0]
  std::vector<vtkImageData*> images;
  for (auto& dataset : models)
  {
    smtk::common::UUID imageId;
    // retrieve the first cell, containing the volume image to save.
    if (!dataset.cells().empty())
    {
      imageId = dataset.cells()[0].entity();
    }
    // For some reason, session can be different for different resources. Re-fetch the session.
    session =
      std::dynamic_pointer_cast<smtk::session::aeva::Resource>(dataset.entityRecord()->resource())
        ->session();
    auto* image = session ? vtkImageData::SafeDownCast(session->findStorage(imageId)) : nullptr;
    if (!image)
    {
      smtkWarningMacro(this->log(), "Cannot retrieve image from model: " << dataset.name());
      continue;
    }
    images.push_back(image);
  }
  // For the found images, write one or more files.
  // Append a number to the filename for multiple files.
  for (int count = 0; count < images.size(); ++count)
  {
    auto* image = images[count];

    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");
    if (images.size() > 1)
    {
      // add a number to multi-file output. Use Paths::replaceFilename() when smtk includes it.
      std::ostringstream newfn;
      newfn << smtk::common::Paths::stem(filename) << std::to_string(count) << ".vti";
      filename = boost::filesystem::path(filename).parent_path().append(newfn.str()).string();
    }

    writer->SetFileName(filename.c_str());
    writer->SetInputData(image);
    if (!writer->Write())
    {
      smtkErrorMacro(this->log(), "vtkXMLImageDataWriter returned failure.");
      return m_result;
    }
  }
  m_result->findInt("outcome")->setValue(0, static_cast<int>(Export::Outcome::SUCCEEDED));
  return m_result;
}

Export::Result Export::exportMedMesh()
{
  auto inputItem = this->parameters()->associations();

  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);
  m_result = this->createResult(Export::Outcome::FAILED);

  for (int k = 0; k < static_cast<int>(inputItem->numberOfValues()); ++k)
  {
    smtk::session::aeva::Session::Ptr session = nullptr;
    smtk::session::aeva::Resource::Ptr resource =
      inputItem->valueAs<smtk::session::aeva::Resource>(k);
    auto inputValue = inputItem->valueAs<smtk::model::Entity>(k);
    smtk::model::Models models;
    if (resource)
    {
      models =
        resource->entitiesMatchingFlagsAs<smtk::model::Models>(smtk::model::MODEL_ENTITY, false);
      session = resource->session();
    }
    else if (inputValue)
    {
      if (!inputValue->isModel())
      {
        smtkErrorMacro(this->log(), "Expected model as input to export.");
        return m_result;
      }
      this->prepareResourceAndSession(m_result, resource, session, false);
      models.push_back(smtk::model::Model(resource, inputValue->id()));
    }

    if (!resource)
    {
      smtkErrorMacro(this->log(), "No resources to export.");
      return m_result;
    }

    for (auto& dataset : models)
    {
      // organize data from the resource in a multiblock as the writer expects.
      vtkSmartPointer<vtkMultiBlockDataSet> writerInput = vtkMultiBlockDataSet::New();
      vtkSmartPointer<vtkMultiBlockDataSet> geomInput = vtkMultiBlockDataSet::New();
      writerInput->SetBlock(0, geomInput);
      vtkSmartPointer<vtkMultiBlockDataSet> groupInput = vtkMultiBlockDataSet::New();
      writerInput->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;
      vtkSmartPointer<vtkUnstructuredGrid> volume;
      std::string volumeName;
      vtkSmartPointer<vtkUnstructuredGrid> surface;
      vtkNew<vtkUnstructuredGrid> newUG;
      std::string surfaceName;
      for (auto const& cellRef : dataset.cells())
      {
        auto const& cellId = cellRef.entity();
        auto* unstructuredGrid = vtkUnstructuredGrid::SafeDownCast(session->findStorage(cellId));
        auto* polyData = vtkPolyData::SafeDownCast(session->findStorage(cellId));
        if (polyData && !unstructuredGrid)
        {
          // allow polydata surfaces like .stl files to be exported as .med - convert to UG
          vtkNew<vtkAppendFilter> appendFilter;
          appendFilter->AddInputData(polyData);
          appendFilter->Update();

          newUG->ShallowCopy(appendFilter->GetOutput());
          unstructuredGrid = newUG;
        }
        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->GetNumberOfCells() > 0 &&
              !!vtkCell3D::SafeDownCast(unstructuredGrid->GetCell(0)))))
        {
          // 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->GetNumberOfCells() > 0))
        {
          // 2D surfaces like .stl and 3D cells from other unstructured grids get caught here.
          surface = unstructuredGrid;
          surfaceName = cellRef.name();
        }
      }

      if (!surface)
      {
        smtkWarningMacro(this->log(), "Cannot retrieve surface from model: " << dataset.name());
        continue;
      }
      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);
      writerInput->GetInformation()->Set(vtkMultiBlockDataSet::NAME(), dataset.name());
      writer->AddInputDataObject(0, writerInput);
    }
  }
  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));
}
}
}
}
