XR_DEFINE_HANDLE(XrSpatialImageTrackingDatabaseEXT)

XrResult xrCreateSpatialImageTrackingDatabaseAsyncEXT(
    XrSession                                   session,
    const XrSpatialImageTrackingDatabaseCreateInfoEXT* createInfo,
    XrFutureEXT*                                future);

typedef struct XrSpatialImageTrackingDatabaseCreateInfoEXT {
    XrStructureType                      type;
    const void*                          next;
    uint32_t                             spatialReferenceImageCount;
    const XrSpatialReferenceImageEXT*    spatialReferenceImages;
} XrSpatialImageTrackingDatabaseCreateInfoEXT;

XrResult xrCreateSpatialImageTrackingDatabaseCompleteEXT(
    XrSession                                   session,
    XrFutureEXT                                 future,
    XrCreateSpatialImageTrackingDatabaseCompletionEXT* completion);

typedef struct XrCreateSpatialImageTrackingDatabaseCompletionEXT {
    XrStructureType                      type;
    void*                                next;
    XrResult                             futureResult;
    XrSpatialImageTrackingDatabaseEXT    database;
} XrCreateSpatialImageTrackingDatabaseCompletionEXT;

XrResult xrDestroySpatialImageTrackingDatabaseEXT(
    XrSpatialImageTrackingDatabaseEXT           database);

XrResult xrEnumerateSpatialReferenceImageFormatsEXT(
    XrInstance                                  instance,
    XrSystemId                                  systemId,
    XrSpatialCapabilityEXT                      capability,
    uint32_t                                    formatCapacityInput,
    uint32_t*                                   formatCountOutput,
    XrSpatialReferenceImageFormatEXT*           formats);

typedef enum XrSpatialReferenceImageFormatEXT {
    XR_SPATIAL_REFERENCE_IMAGE_FORMAT_RGBA_8888_EXT = 1,
    XR_SPATIAL_REFERENCE_IMAGE_FORMAT_RGB_888_EXT = 2,
    XR_SPATIAL_REFERENCE_IMAGE_FORMAT_YUV_420_888_EXT = 3,
    XR_SPATIAL_REFERENCE_IMAGE_FORMAT_MAX_ENUM_EXT = 0x7FFFFFFF
} XrSpatialReferenceImageFormatEXT;

typedef struct XrSpatialImageStaticOptimizationEXT {
    XrStructureType    type;
    const void*        next;
    XrBool32           optimizeForStaticImage;
} XrSpatialImageStaticOptimizationEXT;

typedef struct XrSpatialImageSizeEXT {
    XrStructureType    type;
    const void*        next;
    float              physicalWidth;
} XrSpatialImageSizeEXT;

typedef struct XrSpatialCapabilityConfigurationImageTrackingEXT {
    XrStructureType                             type;
    const void*                                 next;
    XrSpatialCapabilityEXT                      capability;
    uint32_t                                    enabledComponentCount;
    const XrSpatialComponentTypeEXT*            enabledComponents;
    uint32_t                                    imageTrackingDatabaseCount;
    const XrSpatialImageTrackingDatabaseEXT*    imageTrackingDatabases;
} XrSpatialCapabilityConfigurationImageTrackingEXT;

typedef struct XrSpatialReferenceImagePlaneEXT {
    uint32_t    bufferSize;
    uint8_t*    buffer;
    uint32_t    rowStride;
    uint32_t    pixelStride;
} XrSpatialReferenceImagePlaneEXT;

typedef struct XrSpatialReferenceImageEXT {
    XrStructureType                           type;
    const void*                               next;
    uint32_t                                  width;
    uint32_t                                  height;
    XrSpatialReferenceImageFormatEXT          format;
    uint32_t                                  planeCount;
    const XrSpatialReferenceImagePlaneEXT*    planes;
} XrSpatialReferenceImageEXT;

typedef struct XrSpatialComponentImage2DListEXT {
    XrStructureType             type;
    const void*                 next;
    uint32_t                    imageCount;
    XrSpatialImage2DDataEXT*    images;
} XrSpatialComponentImage2DListEXT;

typedef struct XrSpatialImage2DDataEXT {
    XrSpatialImageTrackingDatabaseEXT    imageTrackingDatabase;
    uint32_t                             referenceImageIndex;
} XrSpatialImage2DDataEXT;

XrInstance instance; // Previously initialized
XrSystemId systemId; // Previously initialized
XrSession session; // Previously initialized

uint32_t capabilityCountOutput = 0;
CHK_XR(xrEnumerateSpatialCapabilitiesEXT(instance, systemId, 0, &capabilityCountOutput, nullptr));
std::vector<XrSpatialCapabilityEXT> capabilities(capabilityCountOutput);
CHK_XR(xrEnumerateSpatialCapabilitiesEXT(instance, systemId, capabilities.size(), &capabilityCountOutput, capabilities.data()));

if (std::find(capabilities.begin(), capabilities.end(), XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT) == capabilities.end()) {
  // Image tracking not supported
}

XrSpatialCapabilityComponentTypesEXT imageTrackingComponents{XR_TYPE_SPATIAL_CAPABILITY_COMPONENT_TYPES_EXT};
CHK_XR(xrEnumerateSpatialCapabilityComponentTypesEXT(instance, systemId, XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT, &imageTrackingComponents));
std::vector<XrSpatialComponentTypeEXT> imageTrackingCapabilityComponents(imageTrackingComponents.componentTypeCountOutput);
imageTrackingComponents.componentTypes = imageTrackingCapabilityComponents.data();
CHK_XR(xrEnumerateSpatialCapabilityComponentTypesEXT(instance, systemId, XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT, &imageTrackingComponents));

if (std::find(imageTrackingCapabilityComponents.begin(), imageTrackingCapabilityComponents.end(), XR_SPATIAL_COMPONENT_TYPE_IMAGE_2D_EXT) == imageTrackingCapabilityComponents.end()) {
  // The image2d component is not supported
}

uint32_t capabilityFeatureCountOutput = 0;
CHK_XR(xrEnumerateSpatialCapabilityFeaturesEXT(instance, systemId, XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT, 0, &capabilityFeatureCountOutput, nullptr));
std::vector<XrSpatialCapabilityFeatureEXT> capabilityFeatures(capabilityFeatureCountOutput);
CHK_XR(xrEnumerateSpatialCapabilityFeaturesEXT(instance, systemId, XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT, capabilityFeatures.size(), &capabilityFeatureCountOutput, capabilityFeatures.data()));

// Helper function to query optional feature support
const auto supportsFeature = [&capabilityFeatures](XrSpatialCapabilityFeatureEXT feature) {
  return std::find(capabilityFeatures.begin(), capabilityFeatures.end(), feature) != capabilityFeatures.end();
};

// supported state for optional features
const bool supportsAutomaticScaleEstimation = supportsFeature(XR_SPATIAL_CAPABILITY_FEATURE_IMAGE_TRACKING_AUTOMATIC_SIZE_IMAGES_EXT);
const bool supportsStaticImages = supportsFeature(XR_SPATIAL_CAPABILITY_FEATURE_IMAGE_TRACKING_STATIC_IMAGES_EXT);
const bool supportsFixedSizeImages = supportsFeature(XR_SPATIAL_CAPABILITY_FEATURE_IMAGE_TRACKING_FIXED_SIZE_IMAGES_EXT);

uint32_t formatCountOutput = 0;
CHK_XR(xrEnumerateSpatialReferenceImageFormatsEXT(instance, systemId, XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT, 0, &formatCountOutput, nullptr));
std::vector<XrSpatialReferenceImageFormatEXT> formats(formatCountOutput);
CHK_XR(xrEnumerateSpatialReferenceImageFormatsEXT(instance, systemId, XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT, formats.size(), &formatCountOutput, formats.data()));

if (std::find(formats.begin(), formats.end(), XR_SPATIAL_REFERENCE_IMAGE_FORMAT_RGBA_8888_EXT) == formats.end()) {
  // The desired format is not supported
}

// Helper function to wait for XrFutureEXT completion
auto waitUntilReady = [&](XrFutureEXT future) {
  XrFuturePollInfoEXT pollInfo{XR_TYPE_FUTURE_POLL_INFO_EXT};
  XrFuturePollResultEXT pollResult{XR_TYPE_FUTURE_POLL_RESULT_EXT};
  pollInfo.future = future;
  do {
    // sleep(1);
    CHK_XR(xrPollFutureEXT(instance, &pollInfo, &pollResult));
  } while (pollResult.state != XR_FUTURE_STATE_READY_EXT);
};

// Data/metadata about reference images (identifiers, 3D models, trigger functions, etc.)
std::vector<std::string> referenceImageIdentifiers;

XrSpatialImageTrackingDatabaseEXT imageTrackingDatabase{};
{
  uint32_t imgSize; // preloaded interleaved RGBA buffer size
  uint8_t* img; // preloaded interleaved RGBA buffer data

  // interleaved RGBA buffer memory layout and buffer positions

  // first row of pixels
  // img[0] | img[1] | img[2] | img[3] | img[0]+pStride | img[1]+pStride | ...
  //  R0    |   G0   |   B0   |   A0   |       R1       |       G1       | ...

  // second row of pixels
  // img[0]+rStride | img[1]+rStride | img[2]+rStride | img[3]+rStride | img[0]+rStride+pStride | ...
  //      R640      |      G640      |      B640      |      A640      |          R641          | ...

  XrSpatialReferenceImagePlaneEXT spatialReferenceImagePlanes[4];

  for (int i = 0; i < 4; ++i) { // for each plane (0=R, 1=G, 2=B, 3=A)
    spatialReferenceImagePlanes[i].buffer = img + i; // set each plane buffer pointer to the first pixel of the plane
    spatialReferenceImagePlanes[i].bufferSize = imgSize - i; // adjust remeinaing buffer size according to offset applied
    spatialReferenceImagePlanes[i].rowStride = 640 * 4; // set byte offset to start of next row for this plane
    spatialReferenceImagePlanes[i].pixelStride = 4; // set byte offset to next pixel for this plane
  }

  XrSpatialReferenceImageEXT spatialReferenceImage{XR_TYPE_SPATIAL_REFERENCE_IMAGE_EXT};
  spatialReferenceImage.width = 640;
  spatialReferenceImage.height = 480;
  spatialReferenceImage.format = XR_SPATIAL_REFERENCE_IMAGE_FORMAT_RGBA_8888_EXT;
  spatialReferenceImage.planeCount = 4;
  spatialReferenceImage.planes = spatialReferenceImagePlanes;

  referenceImageIdentifiers.emplace_back("MyReferenceImage");

  // Helper function to append to the next chain
  XrBaseOutStructure* appendTo = reinterpret_cast<XrBaseOutStructure*>(&spatialReferenceImage);
  const auto appendToNextChain = [&appendTo](XrBaseOutStructure* toAppend) {
    appendTo->next = toAppend;
    appendTo = toAppend;
  };

  XrSpatialImageSizeEXT imageSize{XR_TYPE_SPATIAL_IMAGE_SIZE_EXT};
  imageSize.physicalWidth = 0.1f;

  if (!supportsAutomaticScaleEstimation) {
    if (supportsFixedSizeImages) {
      appendToNextChain(reinterpret_cast<XrBaseOutStructure*>(&imageSize));
    }
  }

  XrSpatialImageStaticOptimizationEXT imageStaticOptimization{XR_TYPE_SPATIAL_IMAGE_STATIC_OPTIMIZATION_EXT};
  imageStaticOptimization.optimizeForStaticImage = XR_TRUE;

  if (supportsStaticImages) {
    appendToNextChain(reinterpret_cast<XrBaseOutStructure*>(&imageStaticOptimization));
  }

  XrSpatialImageTrackingDatabaseCreateInfoEXT spatialImageTrackingDatabaseCreateInfo{XR_TYPE_SPATIAL_IMAGE_TRACKING_DATABASE_CREATE_INFO_EXT};
  spatialImageTrackingDatabaseCreateInfo.spatialReferenceImages = &spatialReferenceImage;
  spatialImageTrackingDatabaseCreateInfo.spatialReferenceImageCount = 1;

  XrFutureEXT createImageTrackingDatabaseFuture = XR_NULL_FUTURE_EXT;
  CHK_XR(xrCreateSpatialImageTrackingDatabaseAsyncEXT(session, &spatialImageTrackingDatabaseCreateInfo, &createImageTrackingDatabaseFuture));

  waitUntilReady(createImageTrackingDatabaseFuture);

  XrCreateSpatialImageTrackingDatabaseCompletionEXT imageTrackingCompletion{XR_TYPE_CREATE_SPATIAL_IMAGE_TRACKING_DATABASE_COMPLETION_EXT};
  CHK_XR(xrCreateSpatialImageTrackingDatabaseCompleteEXT(session, createImageTrackingDatabaseFuture, &imageTrackingCompletion));
  if (imageTrackingCompletion.futureResult != XR_SUCCESS) {
    // Error, image tracking database creation failed
  }

  imageTrackingDatabase = imageTrackingCompletion.database;
}

XrSpatialContextEXT spatialContext{};
{
  const std::array<XrSpatialComponentTypeEXT, 2> enabledComponents = {
    XR_SPATIAL_COMPONENT_TYPE_BOUNDED_2D_EXT,
    XR_SPATIAL_COMPONENT_TYPE_IMAGE_2D_EXT
  };

  // Configure the image tracking capability for the spatial context
  XrSpatialCapabilityConfigurationImageTrackingEXT imageTrackingCapabilityConfig{XR_TYPE_SPATIAL_CAPABILITY_CONFIGURATION_IMAGE_TRACKING_EXT};
  imageTrackingCapabilityConfig.capability = XR_SPATIAL_CAPABILITY_IMAGE_TRACKING_EXT;
  imageTrackingCapabilityConfig.enabledComponentCount = enabledComponents.size();
  imageTrackingCapabilityConfig.enabledComponents = enabledComponents.data();
  imageTrackingCapabilityConfig.imageTrackingDatabaseCount = 1;
  imageTrackingCapabilityConfig.imageTrackingDatabases = &imageTrackingDatabase;

  std::vector<XrSpatialCapabilityConfigurationBaseHeaderEXT*> capabilityConfigs;
  capabilityConfigs.push_back(reinterpret_cast<XrSpatialCapabilityConfigurationBaseHeaderEXT*>(&imageTrackingCapabilityConfig));

  XrSpatialContextCreateInfoEXT spatialContextCreateInfo{XR_TYPE_SPATIAL_CONTEXT_CREATE_INFO_EXT};
  spatialContextCreateInfo.capabilityConfigCount = capabilityConfigs.size();
  spatialContextCreateInfo.capabilityConfigs = capabilityConfigs.data();

  XrFutureEXT createContextFuture = XR_NULL_FUTURE_EXT;
  CHK_XR(xrCreateSpatialContextAsyncEXT(session, &spatialContextCreateInfo, &createContextFuture));

  waitUntilReady(createContextFuture);

  XrCreateSpatialContextCompletionEXT completion{XR_TYPE_CREATE_SPATIAL_CONTEXT_COMPLETION_EXT};
  CHK_XR(xrCreateSpatialContextCompleteEXT(session, createContextFuture, &completion));
  if (completion.futureResult != XR_SUCCESS) {
    // Error, spatial context creation failed
  }

  spatialContext = completion.spatialContext;
}

auto discoverSpatialEntities = [&](XrSpatialContextEXT spatialContext, XrTime time) {
  // We want to look for entities that have the following components
  std::array<XrSpatialComponentTypeEXT, 2> snapshotComponents {XR_SPATIAL_COMPONENT_TYPE_BOUNDED_2D_EXT, XR_SPATIAL_COMPONENT_TYPE_IMAGE_2D_EXT};

  XrSpatialDiscoverySnapshotCreateInfoEXT snapshotCreateInfo{XR_TYPE_SPATIAL_DISCOVERY_SNAPSHOT_CREATE_INFO_EXT};
  snapshotCreateInfo.componentTypeCount = snapshotComponents.size();
  snapshotCreateInfo.componentTypes = snapshotComponents.data();
  XrFutureEXT discoveryFuture = XR_NULL_FUTURE_EXT;
  CHK_XR(xrCreateSpatialDiscoverySnapshotAsyncEXT(spatialContext, &snapshotCreateInfo, &discoveryFuture));

  waitUntilReady(discoveryFuture);

  XrCreateSpatialDiscoverySnapshotCompletionInfoEXT completionInfo{XR_TYPE_CREATE_SPATIAL_DISCOVERY_SNAPSHOT_COMPLETION_INFO_EXT};
  completionInfo.baseSpace = localSpace;
  completionInfo.time = time;
  completionInfo.future = discoveryFuture;

  XrCreateSpatialDiscoverySnapshotCompletionEXT completion{XR_TYPE_CREATE_SPATIAL_DISCOVERY_SNAPSHOT_COMPLETION_EXT};
  CHK_XR(xrCreateSpatialDiscoverySnapshotCompleteEXT(spatialContext, &completionInfo, &completion));
  if (completion.futureResult == XR_SUCCESS) {
    // Query for the bounded2d component data
    XrSpatialComponentTypeEXT componentsToQuery[] = {XR_SPATIAL_COMPONENT_TYPE_BOUNDED_2D_EXT, XR_SPATIAL_COMPONENT_TYPE_IMAGE_2D_EXT};
    XrSpatialComponentDataQueryConditionEXT queryCond{XR_TYPE_SPATIAL_COMPONENT_DATA_QUERY_CONDITION_EXT};
    queryCond.componentTypes = componentsToQuery;

    XrSpatialComponentDataQueryResultEXT queryResult{XR_TYPE_SPATIAL_COMPONENT_DATA_QUERY_RESULT_EXT};
    CHK_XR(xrQuerySpatialComponentDataEXT(completion.snapshot, &queryCond, &queryResult));

    std::vector<XrSpatialEntityIdEXT> entityIds(queryResult.entityIdCountOutput);
    std::vector<XrSpatialEntityTrackingStateEXT> entityStates(queryResult.entityStateCountOutput);
    queryResult.entityIdCapacityInput = entityIds.size();
    queryResult.entityIds = entityIds.data();
    queryResult.entityStateCapacityInput = entityStates.size();
    queryResult.entityStates = entityStates.data();

    std::vector<XrSpatialBounded2DDataEXT> bounded2d(queryResult.entityIdCountOutput);
    XrSpatialComponentBounded2DListEXT boundsList{XR_TYPE_SPATIAL_COMPONENT_BOUNDED_2D_LIST_EXT};
    boundsList.boundCount = bounded2d.size();
    boundsList.bounds = bounded2d.data();
    queryResult.next = &boundsList;

    std::vector<XrSpatialImage2DDataEXT> image2d(queryResult.entityIdCountOutput);
    XrSpatialComponentImage2DListEXT imagesList{XR_TYPE_SPATIAL_COMPONENT_IMAGE_2D_LIST_EXT};
    imagesList.imageCount = image2d.size();
    imagesList.images = image2d.data();
    queryResult.next = &imagesList;

    CHK_XR(xrQuerySpatialComponentDataEXT(completion.snapshot, &queryCond, &queryResult));

    for (int32_t i = 0; i < queryResult.entityIdCountOutput; ++i) {
      if (entityStates[i] == XR_SPATIAL_ENTITY_TRACKING_STATE_TRACKING_EXT) {
        const XrSpatialImage2DDataEXT& currentImage2d = image2d[i];
        const std::string& currentReferenceImageIdentifier = referenceImageIdentifiers[currentImage2d.referenceImageIndex];
        std::cout << currentReferenceImageIdentifier << std::endl; // Outputs the "MyReferenceImage" string the app defined previously; could be any other assocatied data or action.

        // Access bounded2d[i] for image bounds
      }
    }

    CHK_XR(xrDestroySpatialSnapshotEXT(completion.snapshot));
  }
};

while (1) {
  // For every frame in frame loop

  XrFrameState frameState;  // Previously returned from xrWaitFrame
  const XrTime time = frameState.predictedDisplayTime;

  // Poll for the XR_TYPE_EVENT_DATA_SPATIAL_DISCOVERY_RECOMMENDED_EXT event
  XrEventDataBuffer event = {XR_TYPE_EVENT_DATA_BUFFER};
  XrResult result = xrPollEvent(instance, &event);
  if (result == XR_SUCCESS) {
      switch (event.type) {
          case XR_TYPE_EVENT_DATA_SPATIAL_DISCOVERY_RECOMMENDED_EXT: {
              const XrEventDataSpatialDiscoveryRecommendedEXT& eventdata = *reinterpret_cast<XrEventDataSpatialDiscoveryRecommendedEXT*>(&event);
              // Discover spatial entities for the context that we received the "discovery
              // recommended" event for
              discoverSpatialEntities(eventdata.spatialContext, time);
              break;
          }
      }
  }

  // Finish frame loop
}

CHK_XR(xrDestroySpatialContextEXT(spatialContext));

// Unload spatial image tracking resources internally
CHK_XR(xrDestroySpatialImageTrackingDatabaseEXT(imageTrackingDatabase));