//=========================================================================
//  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/SelectionResponder.h"
#include "smtk/session/aeva/Predicates.h"
#include "smtk/session/aeva/SelectionResponder_xml.h"
#include "vtk/aeva/ext/vtkGlobalIdBooleans.h"

#include "smtk/extension/vtk/geometry/Backend.h"
#include "smtk/extension/vtk/geometry/Geometry.h"
#include "smtk/extension/vtk/source/vtkResourceMultiBlockSource.h"

#include "smtk/view/Selection.h"

#include "smtk/session/aeva/CellSelection.h"
#include "smtk/session/aeva/Resource.h"

#include "smtk/model/Entity.h"
#include "smtk/model/Face.h"
#include "smtk/model/Model.h"

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

#include "smtk/operation/Manager.h"
#include "smtk/operation/MarkGeometry.h"

#include "smtk/io/Logger.h"

// ParaView client
#include "pqView.h"

// ParaView server
#include "vtkPVRenderView.h"

// VTK
#include "vtkAppendDataSets.h"
#include "vtkCellData.h"
#include "vtkCompositeDataIterator.h"
#include "vtkExtractSelection.h"
#include "vtkIdTypeArray.h"
#include "vtkInformation.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkPolyData.h"
#include "vtkSelection.h"
#include "vtkSelectionNode.h"
#include "vtkUnsignedIntArray.h"

#ifdef ERROR
#undef ERROR
#endif

namespace smtk
{
namespace session
{
namespace aeva
{

namespace // anonymous
{

// Deal with old SMTK versions that did not provide access to the interaction mode.
template<typename ResponderClass>
struct InteractionMode
{
  template<typename Test>
  static void test(Test* /* responder */, int& mode, ...)
  {
    mode = vtkPVRenderView::INTERACTION_MODE_POLYGON;
  }

  template<typename Test>
  static void test(Test* responder, int& mode, decltype(&Test::interactionMode) /* unused */)
  {
    mode = responder->interactionMode();
  }

  InteractionMode(ResponderClass* responder, int& mode)
  {
    InteractionMode::test<ResponderClass>(responder, mode, nullptr);
  }
};

} // anonmyous namespace

SelectionResponder::SelectionResponder() = default;

SelectionResponder::~SelectionResponder() = default;

bool SelectionResponder::fetchResourceAndSession(std::shared_ptr<Resource>& resource,
  std::shared_ptr<Session>& session) const
{
  auto assoc = this->parameters()->associations();
  smtk::resource::ResourcePtr rsrc = assoc->valueAs<smtk::resource::Resource>();
  resource = std::dynamic_pointer_cast<smtk::session::aeva::Resource>(rsrc);
  if (!resource)
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "No associated resource for aeva selection.");
    return false;
  }
  session = resource->session();
  return session && resource;
}

std::shared_ptr<smtk::model::Entity> SelectionResponder::selectedSideSet() const
{
  std::shared_ptr<smtk::model::Entity> result;
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return result;
  }
  auto seln = this->smtkSelection();
  auto bits = seln->selectionValueFromLabel("selected");
  if (!bits)
  {
    bits = 1;
  }

  for (const auto& selnIt : seln->currentSelection())
  {
    int val = selnIt.second;
    if ((bits & val) == bits && selnIt.first)
    {
      auto comp = std::dynamic_pointer_cast<smtk::model::Entity>(selnIt.first);
      if (comp && session->isSideSet(*comp))
      {
        if (result)
        {
          result = smtk::model::Entity::Ptr();
          break;
        }
        result = comp;
      }
    }
  }
  return result;
}

vtkSmartPointer<vtkDataObject> SelectionResponder::appendCellPrimitives(
  bool includeExistingPrimitiveSelection)
{
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  vtkSmartPointer<vtkDataObject> wholeSelection;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return wholeSelection;
  }

  auto cellSelection = CellSelection::instance();

  smtk::extension::vtk::geometry::Backend vtk;
  const auto& baseGeometry = resource->geometry(vtk);
  if (!baseGeometry)
  {
    smtkErrorMacro(this->log(), "No VTK geometry.");
    return wholeSelection;
  }
  const auto& geometry =
    dynamic_cast<const smtk::extension::vtk::geometry::Geometry&>(*baseGeometry);

  vtkNew<vtkAppendDataSets> append;
  append->MergePointsOn();
  auto seln = this->smtkSelection();
  auto bits = seln->selectionValueFromLabel("selected");
  if (!bits)
  {
    bits = 1;
  }

  bool haveInput = false;
  for (const auto& selnIt : seln->currentSelection())
  {
    int val = selnIt.second;
    if ((bits & val) == bits && selnIt.first)
    {
      auto comp = std::dynamic_pointer_cast<smtk::model::Entity>(selnIt.first);
      if (!includeExistingPrimitiveSelection && comp == cellSelection)
      {
        continue;
      }
      auto data = geometry.data(comp);
      if (data)
      {
        append->AddInputData(data);
        haveInput = true;
      }
    }
  }
  if (haveInput)
  {
    append->Update();
    wholeSelection = append->GetOutputDataObject(0);
  }

  return wholeSelection;
}

vtkSmartPointer<vtkDataObject> SelectionResponder::appendSelectedPrimitives(
  bool includeExistingPrimitiveSelection)
{
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  vtkSmartPointer<vtkDataObject> wholeSelection;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return wholeSelection;
  }

  auto* geometry = this->vtkData();
  ::vtkSelection* selnBlock = this->vtkSelection();
  vtkNew<vtkExtractSelection> extractor;
  extractor->SetInputDataObject(0, geometry);
  extractor->SetInputDataObject(1, selnBlock);
  // extractor->PreserveTopologyOn(); // <-- defines a selection array as point-data or cell-data
  extractor->Update();
  auto* selectedStuff = vtkMultiBlockDataSet::SafeDownCast(extractor->GetOutputDataObject(0));

  // Since both selectedStuff and geometry have the same structure,
  // we can apply the iterator to both.
  vtkCompositeDataIterator* iter = selectedStuff->NewIterator();
  // Users may have selected stuff from a side set. But aeva (at least for now)
  // only allows cell-selections from element blocks. So, instead of creating
  // a selection per iterator entry, we append all selected cells on a per-model
  // basis and create a "selection" per element-block entry.
  std::set<vtkSmartPointer<vtkDataObject> > selectedPieces;
  // Note that if a CellSelection instance exists, it must be in the selection.
  auto cellSelection = CellSelection::instance();
  bool addedCellSeln = false;

  // I. Collect pieces of the VTK selection from different model entities
  for (iter->InitTraversal(); !iter->IsDoneWithTraversal(); iter->GoToNextItem())
  {
    auto* seln = iter->GetCurrentDataObject();
    auto uuid = vtkResourceMultiBlockSource::GetDataObjectUUID(geometry->GetMetaData(iter));
    auto srcComp = resource->findEntity(uuid);
    bool isPrimitiveSelection = (cellSelection && uuid == cellSelection->id());
    if (!srcComp && !isPrimitiveSelection)
    {
      continue;
    }
    if (isPrimitiveSelection && !includeExistingPrimitiveSelection)
    {
      // Skip primitive selection when told to do so.
      continue;
    }
    selectedPieces.insert(seln);
  }
  iter->Delete();

  if (!selectedPieces.empty())
  {
    if (selectedPieces.size() > 1)
    {
      vtkNew<vtkAppendDataSets> append;
      append->MergePointsOn();
      for (auto const& input : selectedPieces)
      {
        append->AddInputData(input);
      }
      append->Update();
      wholeSelection = append->GetOutputDataObject(0);
    }
    else
    {
      wholeSelection = *selectedPieces.begin();
    }
  }
  return wholeSelection;
}

bool SelectionResponder::replaceWithPrimitiveSelection(Result& result)
{
  bool didWork = false;
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return didWork;
  }

  vtkSmartPointer<vtkDataObject> primitiveData = this->appendSelectedPrimitives(false);
  auto cellSelection = CellSelection::instance();
  bool replacedPrimitives;
  if (!cellSelection)
  {
    cellSelection = CellSelection::create(resource, primitiveData, this->manager());
    auto created = result->findComponent("created");
    created->appendValue(cellSelection);
    replacedPrimitives = false;
  }
  else
  {
    cellSelection->replaceData(primitiveData);
    auto modified = result->findComponent("modified");
    modified->appendValue(cellSelection);
    replacedPrimitives = true;
  }
  // TODO: Set dimensionality bit of CellSelection to match
  //       what is selected (nodes, edges, faces, volumes).
  //       CellSelection class must allow this?
  std::set<smtk::resource::ComponentPtr> selection;
  selection.insert(cellSelection);
  auto smtkSelection = this->smtkSelection();
  bool didModify = smtkSelection->modifySelection(selection,
    m_smtkSelectionSource,
    m_smtkSelectionValue,
    view::SelectionAction::FILTERED_REPLACE,
    /* bitwise */ true,
    /* notify */ true);

  // Either we updated the CellSelection with new primitives or we
  // modified the SMTK selection:
  didWork = replacedPrimitives | didModify;

  return didWork;
}

bool SelectionResponder::addPrimitiveSelection(Result& result)
{
  bool didWork = false;
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return didWork;
  }

  vtkSmartPointer<vtkDataObject> primitiveData =
    this->appendSelectedPrimitives(/* include pre-existing primitives */ false);
  auto cellSelection = CellSelection::instance();

  // If a single SMTK side-set cell was selected and all
  // the global IDs match its primary geometry(ies), modify that
  // cell's geometry. Otherwise, create a CellSelection
  // holding all primitives in the current VTK and pre-existing
  // SMTK selection.
  smtk::model::Entity::Ptr sideSet = this->selectedSideSet();
  if (sideSet)
  {
    std::set<smtk::model::Entity*> sideSetPrimaries;
    bool validSideSet = session->primaryGeometries(sideSet.get(), sideSetPrimaries);
    // smtk::model::Entity* sideSetPrimary = session->primaryGeometry(sideSet);

    std::set<smtk::model::Entity*> primitivePrimaries;
    bool validPrimary = session->primaryGeometries(
      primitiveData, vtkDataObject::CELL, primitivePrimaries, resource.get());

    bool primaryIsSubset = std::includes(sideSetPrimaries.begin(),
      sideSetPrimaries.end(),
      primitivePrimaries.begin(),
      primitivePrimaries.end());
    if (validSideSet && validPrimary && primaryIsSubset)
    {
      vtkSmartPointer<vtkDataObject> sideSetGeom = session->findStorage(sideSet->id());
      vtkNew<vtkGlobalIdBooleans> unite;
      unite->SetOperationToUnion();
      unite->SetWorkpiece(sideSetGeom);
      unite->SetTool(primitiveData);
      unite->Update();
      session->addStorage(sideSet->id(), unite->GetOutputDataObject(0));
      smtk::operation::MarkGeometry(resource).markModified(sideSet);
      auto modified = result->findComponent("modified");
      modified->appendValue(sideSet);
      return true;
    }
  }

  bool replacedPrimitives;
  if (!cellSelection)
  {
    cellSelection = CellSelection::create(resource, primitiveData, this->manager());
    auto created = result->findComponent("created");
    created->appendValue(cellSelection);
    replacedPrimitives = false;
  }
  else
  {
    vtkNew<vtkGlobalIdBooleans> unite;
    unite->SetOperationToUnion();
    unite->SetWorkpiece(primitiveData);
    unite->SetTool(cellSelection->data());
    unite->Update();
    cellSelection->replaceData(unite->GetOutputDataObject(0));
    auto modified = result->findComponent("modified");
    modified->appendValue(cellSelection);
    replacedPrimitives = true;
  }
  // TODO: Set dimensionality bit of CellSelection to match
  //       what is selected (nodes, edges, faces, volumes).
  //       CellSelection class must allow this?
  std::set<smtk::resource::ComponentPtr> selection;
  selection.insert(cellSelection);
  auto smtkSelection = this->smtkSelection();
  bool didModify = smtkSelection->modifySelection(selection,
    m_smtkSelectionSource,
    m_smtkSelectionValue,
    view::SelectionAction::FILTERED_REPLACE,
    /* bitwise */ true,
    /* notify */ true);

  // Either we updated the CellSelection with new primitives or we
  // modified the SMTK selection:
  didWork = replacedPrimitives | didModify;

  return didWork;
}

bool SelectionResponder::subtractPrimitiveSelection(Result& result)
{
  bool didWork = false;
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return didWork;
  }

  auto cellSelection = CellSelection::instance();

  // If a single SMTK side-set cell was selected and all
  // the global IDs match its primary geometry(ies), modify that
  // cell's geometry. Otherwise, create a CellSelection
  // holding all primitives in the current VTK and pre-existing
  // SMTK selection.
  vtkSmartPointer<vtkDataObject> primitiveData =
    this->appendSelectedPrimitives(/* include pre-existing primitives */ false);
  smtk::model::Entity::Ptr sideSet = this->selectedSideSet();
  if (sideSet && sideSet != cellSelection)
  {
    std::set<smtk::model::Entity*> sideSetPrimaries;
    bool validSideSet = session->primaryGeometries(sideSet.get(), sideSetPrimaries);

    std::set<smtk::model::Entity*> primitivePrimaries;
    bool validPrimary = session->primaryGeometries(
      primitiveData, vtkDataObject::CELL, primitivePrimaries, resource.get());

    bool primaryIsSubset = std::includes(sideSetPrimaries.begin(),
      sideSetPrimaries.end(),
      primitivePrimaries.begin(),
      primitivePrimaries.end());
    if (validSideSet && validPrimary && primaryIsSubset)
    {
      vtkSmartPointer<vtkDataObject> sideSetGeom = session->findStorage(sideSet->id());
      vtkNew<vtkGlobalIdBooleans> diff;
      diff->SetOperationToDifference();
      diff->SetWorkpiece(sideSetGeom);
      diff->SetTool(primitiveData);
      diff->Update();
      session->addStorage(sideSet->id(), diff->GetOutputDataObject(0));
      smtk::operation::MarkGeometry(resource).markModified(sideSet);
      auto modified = result->findComponent("modified");
      modified->appendValue(sideSet);
      return true;
    }
  }

  vtkSmartPointer<vtkDataObject> cellPrimitives =
    this->appendCellPrimitives(/* include pre-existing primitives */ false);
  vtkSmartPointer<vtkDataObject> primitiveDataIncl =
    this->appendSelectedPrimitives(/* include pre-existing primitives */ true);
  bool replacedPrimitives = false;
  if (primitiveDataIncl && NumberOfCells(primitiveDataIncl) > 0 && cellPrimitives &&
    NumberOfCells(cellPrimitives) > 0)
  {
    vtkNew<vtkGlobalIdBooleans> diff;
    diff->SetOperationToDifference();
    diff->SetWorkpiece(cellPrimitives);
    diff->SetTool(primitiveDataIncl);
    diff->Update();
    if (!cellSelection)
    {
      cellSelection =
        CellSelection::create(resource, diff->GetOutputDataObject(0), this->manager());
      auto created = result->findComponent("created");
      created->appendValue(cellSelection);
    }
    else
    {
      cellSelection->replaceData(diff->GetOutputDataObject(0));
      auto modified = result->findComponent("modified");
      modified->appendValue(cellSelection);
    }
    replacedPrimitives = true;
  }
  else if (cellSelection && primitiveDataIncl && NumberOfCells(primitiveDataIncl) > 0)
  {
    vtkNew<vtkGlobalIdBooleans> diff;
    diff->SetOperationToDifference();
    diff->SetWorkpiece(cellSelection->data());
    diff->SetTool(primitiveDataIncl);
    diff->Update();
    cellSelection->replaceData(diff->GetOutputDataObject(0));
    auto modified = result->findComponent("modified");
    modified->appendValue(cellSelection);
    replacedPrimitives = true;
  }

  // TODO: Set dimensionality bit of CellSelection to match
  //       what is selected (nodes, edges, faces, volumes).
  //       CellSelection class must allow this?
  std::set<smtk::resource::ComponentPtr> selection;
  if (cellSelection)
  {
    selection.insert(cellSelection);
  }
  auto smtkSelection = this->smtkSelection();
  bool didModify = smtkSelection->modifySelection(selection,
    m_smtkSelectionSource,
    m_smtkSelectionValue,
    view::SelectionAction::FILTERED_REPLACE,
    /* bitwise */ true,
    /* notify */ true);

  // Either we updated the CellSelection with new primitives or we
  // modified the SMTK selection:
  didWork = replacedPrimitives | didModify;

  return didWork;
}

bool SelectionResponder::xorPrimitiveSelection(Result& /* result */)
{
  bool didWork = false;
  std::shared_ptr<Resource> resource;
  std::shared_ptr<Session> session;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return didWork;
  }

  smtkErrorMacro(this->log(), "Toggle (xor) selection is not implemented yet.");

  return didWork;
}

SelectionResponder::Result SelectionResponder::operateInternal()
{
  auto result = this->createResult(smtk::operation::Operation::Outcome::SUCCEEDED);
  bool worked = false;
  // We only attempt to deal with point/primitve selections, not block selections
  // or picks during camera manipulation:
  int uiMode;
  InteractionMode<SelectionResponder>{ this, uiMode };
  if ((uiMode == vtkPVRenderView::INTERACTION_MODE_SELECTION ||
        uiMode == vtkPVRenderView::INTERACTION_MODE_POLYGON) &&
    !this->selectingBlocks())
  {
    switch (this->modifier())
    {
      default:                           // fall through
      case pqView::PV_SELECTION_DEFAULT: // replace selection with primitive selection
        worked = this->replaceWithPrimitiveSelection(result);
        break;
      case pqView::PV_SELECTION_ADDITION: // add to existing selection
        worked = this->addPrimitiveSelection(result);
        break;
      case pqView::PV_SELECTION_SUBTRACTION: // subtract from existing selection
        worked = this->subtractPrimitiveSelection(result);
        break;
      case pqView::PV_SELECTION_TOGGLE: // xor with existing selection
        worked = this->xorPrimitiveSelection(result);
        break;
    }
  }

  if (!worked)
  {
    result->findInt("outcome")->setValue(
      static_cast<int>(smtk::operation::Operation::Outcome::FAILED));
  }
  return result;
}

const char* SelectionResponder::xmlDescription() const
{
  return SelectionResponder_xml;
}

} // namespace aeva
} // namespace session
} // namespace smtk
