/*
 * Copyright (c) 2023 Arm Limited.
 *
 * SPDX-License-Identifier: MIT
 *
 * 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.
 */

#ifndef CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H
#define CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H

#include "ckw/Kernel.h"
#include "ckw/TensorInfo.h"
#include "ckw/TensorOperand.h"
#include "ckw/TileInfo.h"
#include "ckw/TileOperand.h"

#include <memory>

namespace ckw
{

namespace prototype
{
struct GpuKernelWriterAttribute;
class IGpuKernelWriter;
} // namespace prototype

/** Kernel writer. */
class KernelWriter
{
public:
    // =============================================================================================
    // Constructors and destructor
    // =============================================================================================

    /** Initialize a new instance of kernel writer.
     *
     * @param[in] kernel The kernel to be written to.
     */
    explicit KernelWriter(Kernel &kernel);

    /** Destructor */
    ~KernelWriter();

    /** No copy constructor. */
    KernelWriter(const KernelWriter &) = delete;

    /** No copy assignment. */
    KernelWriter &operator=(const KernelWriter &) = delete;

    // =============================================================================================
    // Scope management
    // =============================================================================================

    /** Get the current ID space. */
    int32_t id_space() const;

    /** Set the current ID space. */
    KernelWriter &id_space(int32_t id_space);

    /** Switch to and return a new ID space. */
    int32_t next_id_space();

    // =============================================================================================
    // Tensor and tile declaration
    // =============================================================================================

    /** Define a tensor argument.
     *
     * @param[in] name The name of the tensor.
     * @param[in] info The tensor info.
     *
     * @return The @ref TensorOperand object.
     */
    TensorOperand &create_tensor_argument(const char *name, const TensorInfo &info);

    /** Define a compile-time constant scalar argument.
     *
     * @param[in] name  The name of the tile.
     * @param[in] value The value of the tile.
     *
     * @return The @ref TileOperand object.
     */
    TileOperand &create_tile_argument(const char *name, int32_t value);

    /** Declare a new tile.
     *
     * The name of the tile must be unique in the current ID space.
     *
     * @param[in] name The name of the tile.
     * @param[in] ...  The necessary arguments to create a new @ref TileOperand.
     *
     * @return The @ref TileOperand object.
     */
    template <typename... TArgs>
    TileOperand &declare_tile(const char *name, TArgs &&...args)
    {
        const auto var_name = generate_variable_name(name);
        auto       operand  = new TileOperand(var_name, ::std::forward<TArgs>(args)...);
        register_operand(operand, true);

        return *operand;
    }

    // =============================================================================================
    // Load and store
    // =============================================================================================

    /** Load the data from the tensor memory to the tile using the sampling information.
     *
     * @param[out] tile    The tile to be loaded.
     * @param[in]  tensor  The tensor to be read.
     * @param[in]  sampler The tensor sampling information.
     */
    void op_load(TileOperand &tile, TensorOperand &tensor, const TensorTileSampler &sampler);

    /** Store the tile to the tensor using the specified sampling information.
     *
     * @param[out] dst     The tensor that the tile is written to.
     * @param[in]  src     The tile to be stored.
     * @param[in]  sampler The tensor sampling information.
     */
    void op_store(TensorOperand &tensor, const TileOperand &tile, const TensorTileSampler &sampler);

    // =============================================================================================
    // Data processing
    // =============================================================================================

    /** Write assignment: `<dst> = <src>`.
     *
     * @param[in] dst The destination tile.
     * @param[in] src The source tile.
     */
    void op_assign(TileOperand &dst, const TileOperand &src);

    /** Write binary expression: `<dst> = <lhs> <op> <rhs>`.
     *
     * @param[in] dst The destination tile.
     * @param[in] lhs The LHS operand.
     * @param[in] rhs The RHS operand.
     * @param[in] op  The binary operator.
     */
    void op_binary_expression(TileOperand &dst, const TileOperand &lhs, const TileOperand &rhs, BinaryOp op);

    /** Write function applied to scalar value: `<dst> = <func>(<src>)`.
     *
     * @param[in] dst  The destination tile.
     * @param[in] src  The source tile.
     * @param[in] func The function to be applied to the source tile.
     */
    void op_scalar_function(TileOperand &dst, const TileOperand &src, ScalarUnaryFunction func);

    // =============================================================================================
    // Misc
    // =============================================================================================

    /** Set `dst` the global ID of dimension `dim`.
     *
     * @param[in] dst The tile to be written to.
     * @param[in] dim The global ID dimension.
     */
    void op_get_global_id(TileOperand &dst, int32_t dim);

    // =============================================================================================
    // Code generation
    // =============================================================================================

    /** Generate the source code of the kernel. */
    ::std::string generate_code();

private:
    /** Generate the full variable name based on the original name and the ID space.
     *
     * @param[in] name The name of the variable.
     *
     * @return The full variable name.
     */
    ::std::string generate_variable_name(const char *name) const;

    /** Register the operand to the kernel.
     *
     * The operand is uniquely owned by the kernel afterward.
     *
     * @param[in] operand   The operand to be registered.
     * @param[in] declaring Whether the tile declaration is generated.
     */
    void register_operand(OperandBase *operand, bool declaring);

private:
    Kernel                                                *_kernel;
    ::std::unique_ptr<prototype::GpuKernelWriterAttribute> _impl_attr;
    ::std::unique_ptr<prototype::IGpuKernelWriter>         _impl;

    int32_t _id_space{ 0 };
    int32_t _max_id_space{ 0 };
};

} // namespace ckw

#endif // CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H