// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

//
// Implements a custom operator kernel which counts the number of calls to Compute(), but otherwise is a no-op.
//

#pragma once

#include "test.h"

template <typename T>
struct NullShapeInferrer : winrt::implements<NullShapeInferrer<T>, IMLOperatorShapeInferrer>
{
    STDMETHOD(InferOutputShapes)(IMLOperatorShapeInferenceContext* context) noexcept
    {
        WINML_EXPECT_NO_THROW(OperatorHelper::ShapeInferenceFunction<T>(context));
        return S_OK;
    }
};

struct NullOperator : winrt::implements<NullOperator, IMLOperatorKernel>
{
    NullOperator(std::atomic<uint32_t>* callCount) : m_callCount(callCount) {}

    STDMETHOD(Compute)(IMLOperatorKernelContext* context)
    {
        winrt::com_ptr<IMLOperatorTensor> outputTensor;
        WINML_EXPECT_HRESULT_SUCCEEDED(context->GetOutputTensor(0, outputTensor.put()));

        ++(*m_callCount);
        return S_OK;
    }

private:
    std::atomic<uint32_t>* m_callCount;
};

struct NullOperatorFactory : winrt::implements<NullOperatorFactory, IMLOperatorKernelFactory>
{
    NullOperatorFactory(std::atomic<uint32_t>* callCount) : m_callCount(callCount) {}

    STDMETHOD(CreateKernel)(
        IMLOperatorKernelCreationContext* context,
        IMLOperatorKernel** kernel)
    {
        ORT_UNUSED_PARAMETER(context);
        auto op = winrt::make<NullOperator>(m_callCount);
        op.copy_to(kernel);
        return S_OK;
    }

    static MLOperatorEdgeDescription CreateEdgeDescriptor(MLOperatorEdgeType type, MLOperatorTensorDataType dataType)
    {
        ORT_UNUSED_PARAMETER(type);
        MLOperatorEdgeDescription desc;
        desc.edgeType = MLOperatorEdgeType::Tensor;
        desc.tensorDataType = dataType;
        return desc;
    }

    static void RegisterKernel(
        const char* name,
        const char* domain,
        int versionSince,
        winrt::com_ptr<IMLOperatorRegistry> registry,
        winrt::com_ptr<IMLOperatorShapeInferrer> shapeInferrer,
        std::atomic<uint32_t>* callCount)
    {
        MLOperatorKernelDescription kernelDescription;
        kernelDescription.domain = domain;
        kernelDescription.name = name;
        kernelDescription.minimumOperatorSetVersion = versionSince;
        kernelDescription.executionType = MLOperatorExecutionType::D3D12;

        MLOperatorEdgeTypeConstrant typeConstraint;
        typeConstraint.typeLabel = "T";
        std::vector<MLOperatorEdgeDescription> allowedEdges
        {
            CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Double),
            CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float),
            CreateEdgeDescriptor(MLOperatorEdgeType::Tensor, MLOperatorTensorDataType::Float16)
        };
        typeConstraint.allowedTypes = allowedEdges.data();
        typeConstraint.allowedTypeCount = static_cast<uint32_t>(allowedEdges.size());

        std::vector<MLOperatorEdgeTypeConstrant> typeConstraints{ typeConstraint };
        kernelDescription.typeConstraints = typeConstraints.data();
        kernelDescription.typeConstraintCount = static_cast<uint32_t>(typeConstraints.size());

        kernelDescription.defaultAttributes = nullptr;
        kernelDescription.defaultAttributeCount = 0;
        kernelDescription.options = MLOperatorKernelOptions::None;
        kernelDescription.executionOptions = 0;

        auto factory = winrt::make<NullOperatorFactory>(callCount);

        WINML_EXPECT_HRESULT_SUCCEEDED(registry->RegisterOperatorKernel(
            &kernelDescription,
            factory.get(),
            shapeInferrer.get()
            ));
    }

private:
    std::atomic<uint32_t>* m_callCount;
};
