/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (C) 2023-2025 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 *******************************************************************************/

#pragma once

#include <functional>
#include <memory>
#include <tuple>
#include <vector>

// CK includes
#include <combined_element_wise_operation.hpp>
#include <device_elementwise_dynamic_vector_dims_impl.hpp>

#include "performance.hpp"
#include "elementwise_meta_traits.hpp"
#include "util.hpp"

namespace hiptensor
{
    class ElementwiseSolution
    {
    public:
        // Due to unique_ptr ownership of members,
        // ElementwiseSolutions should also be considered unique.
        // This means disabling default and copy ctor
        ElementwiseSolution()                           = delete;
        ElementwiseSolution(ElementwiseSolution const&) = delete;
        virtual ~ElementwiseSolution()                  = default;
        ElementwiseSolution& operator=(ElementwiseSolution const&) = delete;

        // This class is intended to receive DeviceOp kernel pointers from
        // the CK generator and take ownership.
        explicit ElementwiseSolution(
            std::unique_ptr<ck::tensor_operation::device::BaseOperator>&& deviceOp);

        // Must specialize incoming arg handling
        virtual bool initArgs(std::vector<float> const&                    scalarValues,
                              std::vector<std::vector<std::size_t>> const& inLengthsArray,
                              std::vector<std::vector<std::size_t>> const& inStridesArray,
                              std::vector<std::vector<int32_t>> const&     inModesArray,
                              std::vector<std::vector<std::size_t>> const& outLengthsArray,
                              std::vector<std::vector<std::size_t>> const& outStridesArray,
                              std::vector<std::vector<int32_t>> const&     outModesArray,
                              std::vector<hiptensorOperator_t> const&      operators,
                              std::vector<const void*> const&              inBuffers,
                              std::vector<void*> const&                    outBuffers)
            = 0;

        float operator()(StreamConfig const& streamConfig = StreamConfig{});

        /// Accessors

        // Unique ID for the kernel
        size_t uid() const;

        // Problem size
        ck::index_t problemSize() const;

        // Kernel's name encoding
        std::string kernelName() const;

        // Reset all arguments
        void resetArgs();

    protected:
        // Derived runtime arguments
        ck::index_t mDim;
        ck::index_t mSize;
        bool        mValid;
        uint32_t    mThreadDim;

        // Kernel Params
        std::unique_ptr<ck::tensor_operation::device::BaseOperator> mDeviceOp;
        std::unique_ptr<ck::tensor_operation::device::BaseArgument> mInvokerArgPtr;
        std::unique_ptr<ck::tensor_operation::device::BaseInvoker>  mInvokerPtr;
    };

} // namespace hiptensor

#include "elementwise_solution_impl.hpp"

