// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause

#include "vtkOpenXRSceneObserver.h"

#include "vtkCommand.h"
#include "vtkObjectFactory.h"
#include "vtkOpenXR.h" // For OpenXR types
#include "vtkOpenXRManager.h"
#include "vtkOpenXRSceneComponent.h"
#include "vtkOpenXRUtilities.h"

#include <algorithm>
#include <unordered_map>

VTK_ABI_NAMESPACE_BEGIN

vtkStandardNewMacro(vtkOpenXRSceneObserver)

  // Components are mapped using the IDS generated by the runtime
  // Note the string is used as a hashable raw array (unprintable character may be stored)
  using ComponentMap = std::unordered_map<std::string, vtkSmartPointer<vtkOpenXRSceneComponent>>;

struct vtkOpenXRSceneObserver::vtkInternals
{
  vtkOpenXRManager* Manager{};
  xr::ExtensionDispatchTable Extensions{};
  XrSceneObserverMSFT SceneObserver{};
  XrTime NextObserverQueryTime{};
  XrTime ObserverStartTime{};
  std::vector<XrSceneComputeFeatureMSFT> SupportedFeatures;
  std::vector<XrSceneComputeFeatureMSFT> EnabledFeatures;
  bool SupportsMarker{};
  ComponentMap Components;

  bool GetMarkersInfo(XrSceneMSFT scene, vtkOpenXRSceneObserver* parent)
  {
    auto& manager = *this->Manager;
    auto& ext = this->Extensions;

    XrSceneMarkerTypeMSFT type{ XR_SCENE_MARKER_TYPE_QR_CODE_MSFT };
    XrSceneMarkerTypeFilterMSFT filter{ XR_TYPE_SCENE_MARKER_TYPE_FILTER_MSFT };
    filter.markerTypeCount = 1;
    filter.markerTypes = &type;

    XrSceneMarkersMSFT markers{ XR_TYPE_SCENE_MARKERS_MSFT };

    XrSceneComponentsGetInfoMSFT getInfo{ XR_TYPE_SCENE_COMPONENTS_GET_INFO_MSFT };
    getInfo.componentType = XR_SCENE_COMPONENT_TYPE_MARKER_MSFT;
    getInfo.next = &filter;

    XrSceneMarkerQRCodesMSFT qrcodes{ XR_TYPE_SCENE_MARKER_QR_CODES_MSFT };
    qrcodes.next = &markers;

    XrSceneComponentsMSFT components{ XR_TYPE_SCENE_COMPONENTS_MSFT };
    components.next = &qrcodes;
    if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
          ext.xrGetSceneComponentsMSFT(scene, &getInfo, &components),
          "Failed to get scene components"))
    {
      return false;
    }

    std::vector<XrSceneComponentMSFT> componentsBuffer;
    componentsBuffer.resize(components.componentCountOutput);
    components.componentCapacityInput = components.componentCountOutput;
    components.components = componentsBuffer.data();

    std::vector<XrSceneMarkerMSFT> markersBuffer;
    markersBuffer.resize(components.componentCountOutput);
    markers.sceneMarkers = markersBuffer.data();
    markers.sceneMarkerCapacityInput = static_cast<uint32_t>(markersBuffer.size());

    std::vector<XrSceneMarkerQRCodeMSFT> qrcodesBuffer;
    qrcodesBuffer.resize(components.componentCountOutput);
    qrcodes.qrCodes = qrcodesBuffer.data();
    qrcodes.qrCodeCapacityInput = static_cast<uint32_t>(qrcodesBuffer.size());

    if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
          ext.xrGetSceneComponentsMSFT(scene, &getInfo, &components),
          "Failed to get scene objects"))
    {
      return false;
    }

    for (uint32_t i = 0; i < markers.sceneMarkerCapacityInput; ++i)
    {
      const auto& marker = markersBuffer[i];
      if (marker.lastSeenTime < this->ObserverStartTime)
      {
        continue; // ignore any component older that our observer
      }

      const auto& compID = componentsBuffer[i].id;
      const std::string compIDStr{ reinterpret_cast<const char*>(compID.bytes), 16ull };

      auto compIt = this->Components.find(compIDStr);
      bool isNew = false;
      if (compIt == this->Components.end())
      {
        isNew = true;
        auto comp = vtkSmartPointer<vtkOpenXRSceneComponent>::New();
        comp->Initialize(vtkOpenXRSceneComponent::Marker);
        compIt = this->Components.emplace(compIDStr, std::move(comp)).first;
      }
      else if (marker.lastSeenTime <= compIt->second->GetLastModifiedTime())
      {
        continue; // did not changed.
      }

      XrSceneComponentsLocateInfoMSFT locateInfo{ XR_TYPE_SCENE_COMPONENTS_LOCATE_INFO_MSFT };
      locateInfo.baseSpace = manager.GetReferenceSpace();
      locateInfo.time = manager.GetPredictedDisplayTime();
      locateInfo.componentIdCount = 1;
      locateInfo.componentIds = &compID;

      XrSceneComponentLocationMSFT location;
      XrSceneComponentLocationsMSFT locations{ XR_TYPE_SCENE_COMPONENT_LOCATIONS_MSFT };
      locations.locationCount = 1;
      locations.locations = &location;

      if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
            ext.xrLocateSceneComponentsMSFT(scene, &locateInfo, &locations),
            "Failed to get scene objects"))
      {
        return false;
      }

      uint32_t count = 0;
      if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
            ext.xrGetSceneMarkerDecodedStringMSFT(scene, &compID, 0, &count, nullptr),
            "failed to query qrcode string"))
      {
        return false;
      }

      std::string text;
      text.resize(count);
      if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
            ext.xrGetSceneMarkerDecodedStringMSFT(
              scene, &compID, static_cast<uint32_t>(text.size()), &count, text.data()),
            "failed to query qrcode text"))
      {
        return false;
      }

      vtkNew<vtkMatrix4x4> matrix;
      vtkOpenXRUtilities::SetMatrixFromXrPose(matrix, location.pose);

      vtkOpenXRSceneComponent* sceneComp = compIt->second.Get();
      sceneComp->UpdateMarkerRepresentation(marker.lastSeenTime, matrix,
        static_cast<double>(marker.size.width), static_cast<double>(marker.size.height),
        std::move(text));

      if (isNew)
      {
        parent->InvokeEvent(vtkCommand::UpdateDataEvent, sceneComp);
      }
    }

    return true;
  }
};

//------------------------------------------------------------------------------
vtkOpenXRSceneObserver::vtkOpenXRSceneObserver()
  : Impl{ new vtkInternals{} }
{
}

//------------------------------------------------------------------------------
vtkOpenXRSceneObserver::~vtkOpenXRSceneObserver()
{
  if (this->Impl->SceneObserver)
  {
    this->Impl->Extensions.xrDestroySceneObserverMSFT(this->Impl->SceneObserver);
  }
}

//------------------------------------------------------------------------------
bool vtkOpenXRSceneObserver::Initialize()
{
  this->Impl->Manager = &vtkOpenXRManager::GetInstance();
  this->Impl->Extensions.PopulateDispatchTable(this->Impl->Manager->GetXrRuntimeInstance());
  return this->CreateMSFTSceneObserver();
}

//------------------------------------------------------------------------------
bool vtkOpenXRSceneObserver::UpdateSceneData()
{
  auto& manager = *this->Impl->Manager;
  const auto currentTime = manager.GetPredictedDisplayTime();

  if (this->Impl->ObserverStartTime == 0) // initialize start time at first query
  {
    this->Impl->ObserverStartTime = currentTime;
  }

  if (currentTime < this->Impl->NextObserverQueryTime)
  {
    return true; // successfully did nothing!
  }

  const double minIntervalNanoSeconds = this->MinimumInterval * 1.0e9;
  this->Impl->NextObserverQueryTime = currentTime + static_cast<XrTime>(minIntervalNanoSeconds);

  auto& ext = this->Impl->Extensions;

  XrSceneComputeStateMSFT state;
  if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
        ext.xrGetSceneComputeStateMSFT(this->Impl->SceneObserver, &state),
        "can not query observer state"))
  {
    return false;
  }

  // get result
  if (state == XR_SCENE_COMPUTE_STATE_COMPLETED_MSFT)
  {
    XrSceneCreateInfoMSFT info{ XR_TYPE_SCENE_CREATE_INFO_MSFT };
    XrSceneMSFT scene;
    if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
          ext.xrCreateSceneMSFT(this->Impl->SceneObserver, &info, &scene),
          "Failed to create scene"))
    {
      return false;
    }

    this->Impl->GetMarkersInfo(scene, this);

    ext.xrDestroySceneMSFT(scene);
  }

  if (state == XR_SCENE_COMPUTE_STATE_NONE_MSFT || state == XR_SCENE_COMPUTE_STATE_COMPLETED_MSFT)
  {
    // Clipping sphere
    XrSceneSphereBoundMSFT sphere{};
    sphere.center = manager.GetViewPose(0)->position;
    sphere.radius = static_cast<float>(this->ClippingRadius);

    XrNewSceneComputeInfoMSFT info{ XR_TYPE_NEW_SCENE_COMPUTE_INFO_MSFT };
    info.requestedFeatureCount = static_cast<uint32_t>(this->Impl->EnabledFeatures.size());
    info.requestedFeatures = this->Impl->EnabledFeatures.data();
    info.consistency = static_cast<XrSceneComputeConsistencyMSFT>(this->ComputeConsistency);
    info.bounds.space = manager.GetReferenceSpace();
    info.bounds.time = manager.GetPredictedDisplayTime();
    info.bounds.spheres = &sphere;
    info.bounds.sphereCount = 1;

    if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
          ext.xrComputeNewSceneMSFT(this->Impl->SceneObserver, &info),
          "Failed to start scene observer"))
    {
      return false;
    }

    return true;
  }

  vtkWarningWithObjectMacro(nullptr, << "Scene observer failed");
  return false;
}

//------------------------------------------------------------------------------
bool vtkOpenXRSceneObserver::EnableComputeFeature(SceneFeature feature)
{
  if (!this->IsComputeFeatureSupported(feature))
  {
    return false;
  }

  if (this->IsComputeFeatureEnabled(feature))
  {
    return true; // already enabled
  }

  this->Impl->EnabledFeatures.emplace_back(static_cast<XrSceneComputeFeatureMSFT>(feature));
  return true;
}

//------------------------------------------------------------------------------
void vtkOpenXRSceneObserver::DisableComputeFeature(SceneFeature feature)
{
  auto& enabled = this->Impl->EnabledFeatures;
  const auto it =
    std::find(enabled.begin(), enabled.end(), static_cast<XrSceneComputeFeatureMSFT>(feature));
  if (it != enabled.end())
  {
    enabled.erase(it);
  }
}

//------------------------------------------------------------------------------
bool vtkOpenXRSceneObserver::IsComputeFeatureEnabled(SceneFeature feature) const
{
  const auto& enabled = this->Impl->EnabledFeatures;
  const auto it =
    std::find(enabled.begin(), enabled.end(), static_cast<XrSceneComputeFeatureMSFT>(feature));
  return it != enabled.end();
}

//------------------------------------------------------------------------------
bool vtkOpenXRSceneObserver::IsComputeFeatureSupported(SceneFeature feature) const
{
  const auto& supported = this->Impl->SupportedFeatures;
  const auto it =
    std::find(supported.begin(), supported.end(), static_cast<XrSceneComputeFeatureMSFT>(feature));
  return it != supported.end();
}

//------------------------------------------------------------------------------
bool vtkOpenXRSceneObserver::CreateMSFTSceneObserver()
{
  XrSceneObserverCreateInfoMSFT info{ XR_TYPE_SCENE_OBSERVER_CREATE_INFO_MSFT };

  auto& manager = *this->Impl->Manager;
  auto& ext = this->Impl->Extensions;
  auto session = manager.GetSession();

  if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
        ext.xrCreateSceneObserverMSFT(session, &info, &this->Impl->SceneObserver),
        "Failed to create MSFT scene observer"))
  {
    return false;
  }

  auto instance = manager.GetXrRuntimeInstance();
  const auto systemId = manager.GetSystemID();

  // Check support for markers
  uint32_t count = 0;
  if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
        ext.xrEnumerateSceneComputeFeaturesMSFT(instance, systemId, 0, &count, nullptr),
        "Failed to enumerate scene compute features"))
  {
    return false;
  }

  std::vector<XrSceneComputeFeatureMSFT> features;
  features.resize(count);
  if (!manager.XrCheckOutput(vtkOpenXRManager::WarningOutput,
        ext.xrEnumerateSceneComputeFeaturesMSFT(instance, systemId, count, &count, features.data()),
        "Failed to enumerate scene compute features"))
  {
    return false;
  }

  this->Impl->SupportedFeatures = std::move(features);

  // Enable all implemented features
  this->EnableComputeFeature(SceneFeature::Markers);

  return true;
}

VTK_ABI_NAMESPACE_END