// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier:  MIT

#pragma once

#include <memory>
#include <mutex>

#include <hipdnn_flatbuffers_sdk/data_objects/convolution_bwd_attributes_generated.h>
#include <hipdnn_flatbuffers_sdk/data_objects/tensor_attributes_generated.h>
#include <miopen/miopen.h>

#include <hipdnn_plugin_sdk/interfaces/IPlan.hpp>

#include "HipdnnMiopenHandle.hpp"
#include "HipdnnMiopenSettings.hpp"
#include "MiopenConvDescriptor.hpp"
#include "MiopenTensor.hpp"

namespace miopen_plugin
{

class ConvBwdParams
{
public:
    ConvBwdParams(
        const hipdnn_flatbuffers_sdk::data_objects::ConvolutionBwdAttributes& attributes,
        const std::unordered_map<int64_t,
                                 const hipdnn_flatbuffers_sdk::data_objects::TensorAttributes*>&
            tensorMap,
        bool deterministicEnabled = false);

    ConvBwdParams(const ConvBwdParams&) = delete;
    ConvBwdParams& operator=(const ConvBwdParams&) = delete;

    ConvBwdParams(ConvBwdParams&&) = default;
    ConvBwdParams& operator=(ConvBwdParams&&) = default;

    const MiopenTensor& dx() const;
    const MiopenTensor& w() const;
    const MiopenTensor& dy() const;
    const MiopenConvDescriptor& conv() const;

    bool validTensors() const;

private:
    size_t _spatialDimCount;
    MiopenTensor _dx;
    MiopenTensor _w;
    MiopenTensor _dy;
    MiopenConvDescriptor _conv;
    bool _tensorsValid;
};

class ConvBwdPlan : public hipdnn_plugin_sdk::IPlan<HipdnnMiopenHandle>
{
public:
    ConvBwdPlan(const HipdnnMiopenHandle& handle,
                ConvBwdParams&& params,
                const HipdnnMiopenSettings& executionSettings);
    ~ConvBwdPlan() override = default;

    ConvBwdPlan(const ConvBwdPlan&) = delete;
    ConvBwdPlan& operator=(const ConvBwdPlan&) = delete;

    ConvBwdPlan(ConvBwdPlan&& other) = delete;
    ConvBwdPlan& operator=(ConvBwdPlan&& other) = delete;

    size_t getWorkspaceSize(const HipdnnMiopenHandle& handle) const override;

    void execute(const HipdnnMiopenHandle& handle,
                 const hipdnnPluginDeviceBuffer_t* deviceBuffers,
                 uint32_t numDeviceBuffers,
                 void* workspace = nullptr) const override;

private:
    ConvBwdParams _params;
    mutable std::mutex _algorithmMutex;
    mutable std::optional<miopenConvBwdDataAlgorithm_t> _algorithm;
    mutable size_t _workspaceSize = 0;
    HipdnnMiopenSettings _executionSettings;
};

} // namespace miopen_plugin