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

#include "Node.hpp"
#include <hipdnn_data_sdk/utilities/ShapeUtilities.hpp>
#include <hipdnn_frontend/Error.hpp>
#include <hipdnn_frontend/attributes/{{ op.frontend.effective_attributes_filename_full }}.hpp>
#include <hipdnn_frontend/attributes/GraphAttributes.hpp>
#include <hipdnn_frontend/detail/{{ op.packer_filename }}>
{% if op.frontend.unpacker_function %}
#include <hipdnn_frontend/detail/{{ op.unpacker_filename }}>
{% endif %}
#include <hipdnn_frontend/detail/ScopedHipdnnBackendDescriptor.hpp>

namespace hipdnn_frontend::graph
{
class {{ op.frontend.node_class }} : public BaseNode<{{ op.frontend.node_class }}, {{ op.frontend.node_type_enum }}>
{
public:
    {{ op.frontend.attributes_class }} attributes;

    {{ op.frontend.node_class }}({{ op.frontend.attributes_class }}&& attrs, const GraphAttributes& graphAttrs)
        : BaseNode(graphAttrs)
        , attributes(std::move(attrs))
    {
    }
{% if op.frontend.unpacker_function %}

    Error unpack_from_descriptor(
        hipdnnBackendDescriptor_t opDesc,
        std::unordered_map<int64_t, std::shared_ptr<TensorAttributes>>& tensorMap) override
    {
        {{ op.frontend.attributes_class }} attrs;
        HIPDNN_CHECK_ERROR(detail::{{ op.frontend.unpacker_function }}(opDesc, tensorMap, attrs));
        attributes = std::move(attrs);
        return {};
    }
{% endif %}

    Error pre_validate_node() const override
    {
{% if op.validation %}
{% if op.validation.required_input_tensors %}
        // Validate required input tensors
{% for tensor_name in op.validation.required_input_tensors %}
{% for tensor in op.frontend.inputs %}
{% if tensor.name == tensor_name %}
        HIPDNN_RETURN_IF_FALSE(attributes.{{ tensor.effective_getter_name }}(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "{{ op.frontend.node_class }} missing {{ tensor.name }} (input) for pre-validation");
{% endif %}
{% endfor %}
{% for tensor in op.frontend.outputs %}
{% if tensor.name == tensor_name %}
        HIPDNN_RETURN_IF_FALSE(attributes.{{ tensor.effective_getter_name }}(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "{{ op.frontend.node_class }} missing {{ tensor.name }} (output) for pre-validation");
{% endif %}
{% endfor %}

{% endfor %}
{% endif %}
{% if op.validation.required_input_dims %}
        // Validate required tensor dimensions
{% for tensor_name in op.validation.required_input_dims %}
{% for tensor in op.frontend.all_tensors %}
{% if tensor.name == tensor_name %}
        HIPDNN_RETURN_IF_TRUE(attributes.{{ tensor.effective_getter_name }}()->get_dim().empty(),
                              ErrorCode::ATTRIBUTE_NOT_SET,
                              "{{ op.frontend.node_class }} missing {{ tensor.name }} dimensions for pre-validation");
{% endif %}
{% endfor %}

{% endfor %}
{% endif %}
{% if op.validation.custom_checks %}
        // Custom validation checks
{% for check in op.validation.custom_checks %}
        // TODO: Implement custom check: {{ check }}
{% endfor %}

{% endif %}
{% else %}
        // Validate required tensor pointers
{% for tensor in op.frontend.required_inputs %}
        HIPDNN_RETURN_IF_FALSE(attributes.{{ tensor.effective_getter_name }}(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "{{ op.frontend.node_class }} missing {{ tensor.name }} (input) for pre-validation");

{% endfor %}
{% for tensor in op.frontend.required_outputs %}
        HIPDNN_RETURN_IF_FALSE(attributes.{{ tensor.effective_getter_name }}(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "{{ op.frontend.node_class }} missing {{ tensor.name }} (output) for pre-validation");

{% endfor %}
{% endif %}
        return {};
    }

    Error infer_properties_node() override
    {
        // Validate required tensor pointers
{% for tensor in op.frontend.required_inputs %}
        HIPDNN_RETURN_IF_FALSE(attributes.{{ tensor.effective_getter_name }}(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "{{ op.frontend.node_class }} missing {{ tensor.name }} for setting properties");

{% endfor %}
{% for tensor in op.frontend.required_outputs %}
        HIPDNN_RETURN_IF_FALSE(attributes.{{ tensor.effective_getter_name }}(),
                               ErrorCode::ATTRIBUTE_NOT_SET,
                               "{{ op.frontend.node_class }} missing {{ tensor.name }} for setting properties");

{% endfor %}
        HIPDNN_CHECK_ERROR(attributes.fill_from_context(graph_attributes));

{% if op.infer_properties and op.infer_properties.strategy == "match_input" %}
{% set ref_input = op.infer_properties.reference_input %}
{% set first_output = op.frontend.outputs[0] %}
        // Infer output dimensions from {{ ref_input }}
        auto {{ ref_input }}Tensor = attributes.get_{{ ref_input }}();
        auto {{ first_output.name }}Tensor = attributes.{{ first_output.effective_getter_name }}();

        if({{ first_output.name }}Tensor->get_dim().empty())
        {
            {{ first_output.name }}Tensor->set_dim({{ ref_input }}Tensor->get_dim());
        }

        // Infer output strides if not set
        if({{ first_output.name }}Tensor->get_stride().empty())
        {
            auto& {{ ref_input }}Strides = {{ ref_input }}Tensor->get_stride();
            auto& {{ first_output.name }}Dims = {{ first_output.name }}Tensor->get_dim();

            HIPDNN_RETURN_IF_TRUE(
                {{ ref_input }}Strides.empty(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "{{ op.frontend.node_class }}: Cannot infer output strides - missing {{ ref_input }} strides");

            HIPDNN_RETURN_IF_TRUE(
                {{ first_output.name }}Dims.empty(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "{{ op.frontend.node_class }}: Cannot infer output strides - missing output dimensions");

            auto strideOrder = hipdnn_data_sdk::utilities::extractStrideOrder({{ ref_input }}Strides);
            auto {{ first_output.name }}Strides = hipdnn_data_sdk::utilities::generateStrides({{ first_output.name }}Dims, strideOrder);
            {{ first_output.name }}Tensor->set_stride({{ first_output.name }}Strides);
        }
{% else %}
{% if op.frontend.outputs %}
{% set first_output = op.frontend.outputs[0] %}
{% if op.frontend.inputs %}
{% set first_input = op.frontend.inputs[0] %}
        // TODO: Implement output dimension inference
{% if op.infer_properties and op.infer_properties.dimension_formula %}
        // Formula hint: {{ op.infer_properties.dimension_formula }}
{% endif %}

        auto {{ first_output.name }}Tensor = attributes.{{ first_output.effective_getter_name }}();
        auto {{ first_input.name }}Tensor = attributes.{{ first_input.effective_getter_name }}();

        // Infer output strides if not set
        if({{ first_output.name }}Tensor->get_stride().empty())
        {
            auto& {{ first_input.name }}Strides = {{ first_input.name }}Tensor->get_stride();
            auto& {{ first_output.name }}Dims = {{ first_output.name }}Tensor->get_dim();

            HIPDNN_RETURN_IF_TRUE(
                {{ first_input.name }}Strides.empty(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "{{ op.frontend.node_class }}: Cannot infer output strides - missing input strides");

            HIPDNN_RETURN_IF_TRUE(
                {{ first_output.name }}Dims.empty(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "{{ op.frontend.node_class }}: Cannot infer output strides - missing output dimensions");

            HIPDNN_RETURN_IF_NE(
                {{ first_input.name }}Strides.size(),
                {{ first_output.name }}Dims.size(),
                ErrorCode::ATTRIBUTE_NOT_SET,
                "{{ op.frontend.node_class }}: Stride dimension mismatch between input and output tensors");

            auto strideOrder = hipdnn_data_sdk::utilities::extractStrideOrder({{ first_input.name }}Strides);
            auto {{ first_output.name }}Strides = hipdnn_data_sdk::utilities::generateStrides({{ first_output.name }}Dims, strideOrder);
            {{ first_output.name }}Tensor->set_stride({{ first_output.name }}Strides);
        }
{% endif %}
{% endif %}
{% endif %}

        return {};
    }

    Error create_operation(
        std::unordered_map<int64_t, detail::ScopedHipdnnBackendDescriptor>& tensorDescs,
        std::vector<detail::ScopedHipdnnBackendDescriptor>& operations) const override
    {
        return detail::{{ op.frontend.packer_function }}(attributes, tensorDescs, operations);
    }
};
}