// Copyright (C) 2018-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include <oneapi/dnnl/dnnl.hpp>
#include <sstream>
#include <oneapi/dnnl/dnnl_debug.h>

#include "intel_gpu/runtime/utils.hpp"
#include "test_utils.h"

#include "intel_gpu/runtime/format.hpp"
#include "graph/impls/onednn/utils.hpp"

using namespace cldnn;
using namespace cldnn::onednn;

struct dnnl_desc_params {
    // Descriptor info to test
    dnnl::memory::dims dims;
    dnnl::memory::data_type data_type;
    dnnl::memory::dims strides;  // In case of plain (non-blocked) formats the strides between dimensions
};

struct desc_stride_test_params {
    dnnl_desc_params test_desc;
    std::vector<std::vector<size_t>> expected_orders;
};

class format_test_stride : public testing::TestWithParam<desc_stride_test_params> {
public:
    static std::string PrintToString(testing::TestParamInfo<desc_stride_test_params> param_info) {
        auto strides = param_info.param.test_desc.strides;
        std::string res = " stride : {";
        for (auto stride : strides) {
            res += std::to_string(stride) + " ";
        }
        res += "}";

        return res;
    }
};

TEST_P(format_test_stride, test_candidates_using_stride) {
    auto param = GetParam();

    dnnl::memory::desc test_desc(param.test_desc.dims, param.test_desc.data_type, param.test_desc.strides);
    auto candidates= onednn::get_candidate_orders(test_desc);

    ASSERT_EQ(candidates.size(), param.expected_orders.size());
    bool is_same = true;
    for (size_t idx = 0; idx < param.expected_orders.size(); idx++) {
        auto expected = param.expected_orders.at(idx);
        bool found_match = false;
        for (size_t idx = 0 ; idx < candidates.size() ; idx++) {
            if (std::equal(candidates[idx].begin(), candidates[idx].end(), expected.begin()))
                found_match = true;
        }

        if (!found_match)
            is_same = false;
    }

    ASSERT_TRUE(is_same);
}

INSTANTIATE_TEST_SUITE_P(smoke, format_test_stride,
    testing::ValuesIn(std::vector<desc_stride_test_params>{
        {{{1, 3, 8, 8}, dnnl::memory::data_type::f16, {768, 256, 16, 1}}, {{0, 1, 2, 3}}},
        {{{6, 1, 1, 8}, dnnl::memory::data_type::f16, {96, 16, 16, 1}}, {{0, 1, 2, 3}, {0, 2, 1, 3}}},
        {{{1, 1, 1, 16}, dnnl::memory::data_type::f16, {32, 32, 32, 1}}, {{0, 1, 2, 3}, {0, 2, 1, 3}, {1, 0, 2, 3}, {2, 1, 0, 3}, {1, 2, 0, 3}, {2, 0, 1, 3}}}
    }),
    format_test_stride::PrintToString);


struct format_matching_test_params {
    dnnl::memory::dims dims;
    dnnl::memory::data_type data_type;
    dnnl::memory::format_tag dnnl_format;
    cldnn::format cldnn_format;
};

class data_format_test_match_dnnl : public testing::TestWithParam<format_matching_test_params> {
public:
    static std::string PrintToString(testing::TestParamInfo<format_matching_test_params> param_info) {
        auto dnnl_format = param_info.param.dnnl_format;
        std::string res = " data format (dnnl::memory::format_tag) : " + std::string(dnnl_fmt_tag2str((dnnl_format_tag_t)dnnl_format));
        res += " > " + format::traits(param_info.param.cldnn_format).str;

        return res;
    }
};

TEST_P(data_format_test_match_dnnl, test_match_data_format) {
    auto param = GetParam();

    dnnl::memory::desc test_desc(param.dims, param.data_type, param.dnnl_format);
    auto result = onednn::find_data_format(test_desc);
    ASSERT_TRUE(result == param.cldnn_format);
}

INSTANTIATE_TEST_SUITE_P(smoke, data_format_test_match_dnnl,
    testing::ValuesIn(std::vector<format_matching_test_params>{
        {{1, 3, 8, 8}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::nchw, cldnn::format::bfyx},
        {{1, 3, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::aBcd4b, cldnn::format::b_fs_yx_fsv4},
        {{1, 12, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::nChw16c, cldnn::format::b_fs_yx_fsv16},
        {{1, 12, 16, 16}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::aBcd32b, cldnn::format::b_fs_yx_fsv32},
        {{1, 3, 16, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::aBcde4b, cldnn::format::b_fs_zyx_fsv4},
        {{12, 12, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::NChw16n16c, cldnn::format::bs_fs_yx_bsv16_fsv16},
        {{32, 32, 16, 16}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::NChw32n32c, cldnn::format::bs_fs_yx_bsv32_fsv32},
    }),
    data_format_test_match_dnnl::PrintToString);

class data_format_test_not_match_dnnl : public testing::TestWithParam<format_matching_test_params> {
public:
    static std::string PrintToString(testing::TestParamInfo<format_matching_test_params> param_info) {
        auto dnnl_format = param_info.param.dnnl_format;
        std::string res = " Failed case for " + std::string(dnnl_fmt_tag2str((dnnl_format_tag_t)dnnl_format));
        return res;
    }
};

TEST_P(data_format_test_not_match_dnnl, test_not_match_data_format) {
    auto param = GetParam();

    dnnl::memory::desc test_desc(param.dims, param.data_type, param.dnnl_format);
    auto result = onednn::find_data_format(test_desc);
    ASSERT_FALSE(result == param.cldnn_format);
}

INSTANTIATE_TEST_SUITE_P(smoke, data_format_test_not_match_dnnl,
    testing::ValuesIn(std::vector<format_matching_test_params>{
        {{1, 3, 8, 8}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::nchw, cldnn::format::byxf},
        {{1, 3, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::aBcd4b, cldnn::format::b_fs_yx_fsv2},
        {{1, 12, 16, 16}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::nChw16c, cldnn::format::b_fs_yx_fsv32},
        {{1, 12, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::aBcd32b, cldnn::format::b_fs_yx_fsv16},
        {{1, 3, 16, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::aBcde4b, cldnn::format::b_fs_zyx_fsv16},
        {{32, 32, 16, 16}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::NChw16n16c, cldnn::format::bs_fs_yx_bsv32_fsv32},
        {{12, 12, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::NChw32n32c, cldnn::format::bs_fs_yx_bsv16_fsv16},
    }),
    data_format_test_not_match_dnnl::PrintToString);

class weight_format_test_match_dnnl : public testing::TestWithParam<format_matching_test_params> {
public:
    static std::string PrintToString(testing::TestParamInfo<format_matching_test_params> param_info) {
        auto dnnl_format = param_info.param.dnnl_format;
        std::string res = " weight format (dnnl::memory::format_tag) : " + std::string(dnnl_fmt_tag2str((dnnl_format_tag_t)dnnl_format));
        res += " > " + format::traits(param_info.param.cldnn_format).str;

        return res;
    }
};

TEST_P(weight_format_test_match_dnnl, test_match_data_format) {
    auto param = GetParam();

    dnnl::memory::desc test_desc(param.dims, param.data_type, param.dnnl_format);
    auto result = onednn::find_format(test_desc, false);
    ASSERT_TRUE(result == param.cldnn_format);
}

INSTANTIATE_TEST_SUITE_P(smoke, weight_format_test_match_dnnl,
    testing::ValuesIn(std::vector<format_matching_test_params>{
        {{1, 3, 8, 8}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::abcd, cldnn::format::oiyx},
        {{16, 16, 8, 8}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::ABcd16b16a, cldnn::format::os_is_yx_isv16_osv16},
    }),
    weight_format_test_match_dnnl::PrintToString);

struct memory_desc_to_fmt_conversion_params {
    dnnl::memory::dims dims;
    dnnl::memory::data_type data_type;
    dnnl::memory::format_tag dnnl_format;
    bool grouped;
    cldnn::format_traits expected;
};

class memory_desc_to_fmt_conversion_test : public testing::TestWithParam<memory_desc_to_fmt_conversion_params> {
public:
    static std::string PrintToString(testing::TestParamInfo<memory_desc_to_fmt_conversion_params> param_info) {
        auto dnnl_format = param_info.param.dnnl_format;
        std::stringstream s;
        ov::Shape shape(param_info.param.dims.begin(), param_info.param.dims.end());
        s << shape;
        return std::string(dnnl_fmt_tag2str((dnnl_format_tag_t)dnnl_format)) + "_" + s.str();
    }
};

TEST_P(memory_desc_to_fmt_conversion_test, test_match_data_format) {
    auto param = GetParam();

    dnnl::memory::desc test_desc(param.dims, param.data_type, param.dnnl_format);
    auto result = onednn::convert_memory_desc_to_traits(test_desc, true, param.grouped);
    ASSERT_EQ(result.str, param.expected.str);
    ASSERT_EQ(result.batch_num, param.expected.batch_num);
    ASSERT_EQ(result.feature_num, param.expected.feature_num);
    ASSERT_EQ(result.spatial_num, param.expected.spatial_num);
    ASSERT_EQ(result.group_num, param.expected.group_num);
    ASSERT_EQ(result._order, param.expected._order);
    ASSERT_EQ(result.order, param.expected.order);
    ASSERT_EQ(result.internal_order, param.expected.internal_order);
    ASSERT_EQ(result.block_sizes, param.expected.block_sizes);
    ASSERT_EQ(result.logic_block_sizes, param.expected.logic_block_sizes);
}

INSTANTIATE_TEST_SUITE_P(smoke, memory_desc_to_fmt_conversion_test,
    testing::ValuesIn(std::vector<memory_desc_to_fmt_conversion_params>{
        {
            {1, 3, 8, 8}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::abcd, false,
            format_traits{
                "custom", 1, 1, 2, 0, {0, 1, 2, 3}, "oiyx", "oixy?", {}, {}
            }
        },
        {
            {16, 16, 8, 8}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::ABcd16b16a, false,
            format_traits{
                "custom", 1, 1, 2, 0, {0, 1, 2, 3}, "oiyx", "oixy?", {{1, 16}, {0, 16}}, {{1, 16}, {0, 16}}
            }
        },
        {
            {8, 4, 16, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::ABcde8a4b, false,
            format_traits{
                "custom", 1, 1, 3, 0, {0, 1, 2, 3, 4}, "oizyx", "oixyz", {{0, 8}, {1, 4}}, {{0, 8}, {1, 4}}
            }
        },
        {
            {16, 16, 8, 8}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::ABcd2a8b8a2b, false,
            format_traits{
                "custom", 1, 1, 2, 0, {0, 1, 2, 3}, "oiyx", "oixy?", {{0, 2}, {1, 8}, {0, 8}, {1, 2}}, {{0, 2}, {1, 8}, {0, 8}, {1, 2}}
            }
        },
        {
            {64, 64, 8, 8}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::BAcd4b8a8b4a, false,
            format_traits{
                "custom", 1, 1, 2, 0, {1, 0, 2, 3}, "ioyx", "oixy?", {{1, 4}, {0, 8}, {1, 8}, {0, 4}}, {{1, 4}, {0, 8}, {1, 8}, {0, 4}}
            }
        },
        {
            {32, 32, 8, 8}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::BAcd4b8a8b4a, false,  // same format as above but different sizes
            format_traits{
                "custom", 1, 1, 2, 0, {0, 1, 2, 3}, "oiyx", "oixy?", {{1, 4}, {0, 8}, {1, 8}, {0, 4}}, {{1, 4}, {0, 8}, {1, 8}, {0, 4}} // Order is different due to A=B=1 (strides are equal)
            }
        },
        {
            {65, 32, 8, 8}, dnnl::memory::data_type::u8, dnnl::memory::format_tag::BAcd4b8a8b4a, false,  // same format as above but different sizes
            format_traits{
                "custom", 1, 1, 2, 0, {1, 0, 2, 3}, "ioyx", "oixy?", {{1, 4}, {0, 8}, {1, 8}, {0, 4}}, {{1, 4}, {0, 8}, {1, 8}, {0, 4}}
            }
        },
        {
            {10, 20, 30, 40, 50}, dnnl::memory::data_type::f32, dnnl::memory::format_tag::AcdeB24a2b, true,
            format_traits{
                "custom", 1, 1, 2, 1, {0, 2, 3, 4, 1}, "giyxo", "oixy????g", {{8, 24}, {0, 2}}, {{0, 24}, {1, 2}}
            }
        },
        {
            {10, 60, 70, 40, 50}, dnnl::memory::data_type::f32, dnnl::memory::format_tag::aCBde4c8b8c4b, true,
            format_traits{
                "custom", 1, 1, 2, 1, {0, 2, 1, 3, 4}, "gioyx", "oixy????g", {{1, 4}, {0, 8}, {1, 8}, {0, 4}}, {{2, 4}, {1, 8}, {2, 8}, {1, 4}}
            }
        },
    }),
    memory_desc_to_fmt_conversion_test::PrintToString);

// Test for keep_weights_reorder_shape_consistent function
TEST(keep_weights_reorder_shape_consistent_test, simple_data_formats) {
    // Test case 1: bfyx format with matching shapes
    {
        auto layout = cldnn::layout{ov::PartialShape{512, 1024, 16}, data_types::f16, format::bfyx};
        dnnl::memory::desc desc({512, 1024, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::abc);
        bool result = onednn::keep_weights_reorder_shape_consistent(layout, desc);
        EXPECT_TRUE(result);
        EXPECT_EQ(layout.get_shape(), ov::Shape({512, 1024, 16}));
    }
    // Test case 2: fbyx format with matching shapes
    {
        auto layout = cldnn::layout{ov::PartialShape{512, 1024, 16}, data_types::f16, format::fbyx};
        dnnl::memory::desc desc({512, 1024, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::abc);
        bool result = onednn::keep_weights_reorder_shape_consistent(layout, desc);
        EXPECT_TRUE(result);
        EXPECT_EQ(layout.get_shape(), ov::Shape({512, 1024, 16}));
    }
    // Test case 3: ioyx format with matching shapes
    {
        auto layout = cldnn::layout{ov::PartialShape{512, 1024, 16}, data_types::f16, format::ioyx};
        dnnl::memory::desc desc({512, 1024, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::abc);
        bool result = onednn::keep_weights_reorder_shape_consistent(layout, desc);
        EXPECT_TRUE(result);
        EXPECT_EQ(layout.get_partial_shape(), ov::Shape({512, 1024, 16}));
    }
    // Test case 4: oizyx format with matching shapes
    {
        auto layout = cldnn::layout{ov::PartialShape{16, 1, 16, 16}, data_types::f16, format::oiyx};
        dnnl::memory::desc desc({16, 1, 1, 16, 16}, dnnl::memory::data_type::f16, dnnl::memory::format_tag::abcde);
        bool result = onednn::keep_weights_reorder_shape_consistent(layout, desc);
        EXPECT_TRUE(result);
        EXPECT_EQ(layout.format, cldnn::format::oizyx);
        EXPECT_EQ(layout.get_partial_shape(), ov::Shape({16, 1, 1, 16, 16}));
    }
}

class test_layout_to_memory_desc : public ::testing::Test {
protected:
    void SetUp() override {
        // Common test setup
    }

    dnnl::memory::format_tag get_format_tag_from_desc(const dnnl::memory::desc& desc) {
        if (desc.get_format_kind() != dnnl::memory::format_kind::blocked) {
            return dnnl::memory::format_tag::undef;
        }
        for (int tag = 0; tag < static_cast<int>(dnnl::memory::format_tag::format_tag_last); ++tag) {
            auto format_tag = static_cast<dnnl::memory::format_tag>(tag);
            try {
                dnnl::memory::desc test_desc(desc.get_dims(), desc.get_data_type(), format_tag);
                if (test_desc == desc) {
                    return format_tag;
                }
            } catch (...) {
                continue;
            }
        }
        return dnnl::memory::format_tag::undef;
    }
};

TEST_F(test_layout_to_memory_desc, basic_bfyx) {
    layout l = layout{ov::PartialShape{1, 3, 224, 224}, data_types::f32, format::bfyx};

    auto desc = layout_to_memory_desc(l, dnnl::memory::format_tag::undef);

    EXPECT_EQ(desc.get_ndims(), 4);
    EXPECT_EQ(desc.get_dims()[0], 1);
    EXPECT_EQ(desc.get_dims()[1], 3);
    EXPECT_EQ(desc.get_dims()[2], 224);
    EXPECT_EQ(desc.get_dims()[3], 224);
    EXPECT_EQ(desc.get_data_type(), dnnl::memory::data_type::f32);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::abcd);
}

TEST_F(test_layout_to_memory_desc, 3d_shape) {
    // 3D shape should be handled specially
    layout l = layout{ov::PartialShape{1, 256, 64}, data_types::f16, format::bfyx};

    auto desc = layout_to_memory_desc(l, dnnl::memory::format_tag::undef);

    EXPECT_EQ(desc.get_ndims(), 3);
    EXPECT_EQ(desc.get_data_type(), dnnl::memory::data_type::f16);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::abc);
}

TEST_F(test_layout_to_memory_desc, 3d_shape_with_ioyx) {
    // 3D shape should be handled specially
    layout l = layout{ov::PartialShape{1, 256, 64}, data_types::f16, format::ioyx};

    auto desc = layout_to_memory_desc(l, dnnl::memory::format_tag::any);

    // dnnl::memory::desc represents dnnl::memory::format_tag::bac using strides.
    // refer to src/plugins/intel_gpu/thirdparty/onednn_gpu/src/common/memory_desc_wrapper.cpp
    EXPECT_EQ(desc.get_ndims(), 3);
    EXPECT_EQ(desc.get_dims()[0], 1);
    EXPECT_EQ(desc.get_dims()[1], 256);
    EXPECT_EQ(desc.get_dims()[2], 64);
    EXPECT_EQ(desc.get_data_type(), dnnl::memory::data_type::f16);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::abc);
    EXPECT_EQ(desc.get_strides()[0], 64);
    EXPECT_EQ(desc.get_strides()[1], 64);
    EXPECT_EQ(desc.get_strides()[2], 1);
}

TEST_F(test_layout_to_memory_desc, with_target_format) {
    layout l = layout{ov::PartialShape{1, 256}, data_types::i32, format::bfyx};
    auto desc = layout_to_memory_desc(l, dnnl::memory::format_tag::ab);
    EXPECT_EQ(desc.get_ndims(), 2);
    EXPECT_EQ(desc.get_dims()[0], 1);
    EXPECT_EQ(desc.get_dims()[1], 256);
    EXPECT_EQ(desc.get_data_type(), dnnl::memory::data_type::s32);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::ab);
}


TEST_F(test_layout_to_memory_desc, flatten_basic) {
    layout l = layout{ov::PartialShape{2, 3, 4, 5}, data_types::f16, format::bfyx};
    auto desc = layout_to_memory_desc_flatten(l, dnnl::memory::format_tag::any);
    EXPECT_EQ(desc.get_ndims(), 1);
    EXPECT_EQ(desc.get_dims()[0], 2 * 3 * 4 * 5);
}

TEST_F(test_layout_to_memory_desc, flatten_with_ab_format) {
    layout l = layout{ov::PartialShape{2, 128}, data_types::f16, format::bfyx};
    auto desc = layout_to_memory_desc_flatten(l, dnnl::memory::format_tag::ab);

    EXPECT_EQ(desc.get_ndims(), 2);
    EXPECT_EQ(desc.get_dims()[0], 1);  // Prepended dimension
    EXPECT_EQ(desc.get_dims()[1], 2 * 128);
}

TEST_F(test_layout_to_memory_desc, flatten_bias) {
    // Common use case: bias flattening for convolution
    layout l = layout{ov::PartialShape{64}, data_types::f16, format::bfyx};
    auto desc = layout_to_memory_desc_flatten(l, dnnl::memory::format_tag::a);
    EXPECT_EQ(desc.get_ndims(), 1);
    EXPECT_EQ(desc.get_dims()[0], 64);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::a);
}

TEST_F(test_layout_to_memory_desc, blocked_format) {
    layout l = layout{ov::PartialShape{1, 64, 56, 56}, data_types::f16, format::b_fs_yx_fsv16};
    auto desc = layout_to_memory_desc_blocked(l, dnnl::memory::format_tag::undef);
    EXPECT_EQ(desc.get_ndims(), 4);
    EXPECT_GT(desc.get_inner_nblks(), 0);  // Should have blocking
}

TEST_F(test_layout_to_memory_desc, strides_with_padding) {
    layout l = layout{ov::PartialShape{1, 64, 56, 56}, data_types::f16, format::bfyx};
    l.data_padding = padding{{0, 0, 2, 2}, {0, 0, 2, 2}};
    auto desc = layout_to_memory_desc_strides(l, dnnl::memory::format_tag::abcd);
    EXPECT_EQ(desc.get_ndims(), 4);
    auto strides = desc.get_strides();
    EXPECT_GT(strides[2], 56);  // Stride should account for padding
}

TEST_F(test_layout_to_memory_desc, strides_cannot_flatten) {
    layout l = layout{ov::PartialShape{1, 64, 56, 56}, data_types::f32, format::bfyx};
    l.data_padding = padding{{0, 0, 1, 1}, {0, 0, 1, 1}};
    EXPECT_NO_THROW({
        auto desc = layout_to_memory_desc_strides(l, dnnl::memory::format_tag::abcd);
    });
}

TEST_F(test_layout_to_memory_desc, quantize_use_default) {
    layout l = layout{ov::PartialShape{1, 256}, data_types::f32, format::bfyx};
    auto desc = layout_to_memory_desc(l, true, false);
    EXPECT_EQ(desc.get_ndims(), 2);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::ab);
}

TEST_F(test_layout_to_memory_desc, quantize_output_blocked) {
    layout l = layout{ov::PartialShape{1, 64, 56, 56}, data_types::f16, format::bfyx};
    auto desc = layout_to_memory_desc(l, false, true);
    EXPECT_EQ(desc.get_ndims(), 4);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::abcd);
}

TEST_F(test_layout_to_memory_desc, quantize_gemm_case) {
    // Simulating gemm/fc case where use_default_format=true
    layout scale_layout = layout{ov::PartialShape{1, 256}, data_types::f16, format::bfyx};
    auto desc = layout_to_memory_desc(scale_layout, true, false);
    EXPECT_EQ(desc.get_ndims(), 2);
    EXPECT_EQ(desc.get_dims()[0], 1);
    EXPECT_EQ(desc.get_dims()[1], 256);
}

TEST_F(test_layout_to_memory_desc, quantize_conv_case) {
    // Simulating conv case where use_default_format=false, is_output_blocked=true
    layout scale_layout = layout{ov::PartialShape{1, 64}, data_types::f16, format::bfyx};
    auto desc = layout_to_memory_desc(scale_layout, false, true);
    EXPECT_EQ(desc.get_ndims(), 4);
    EXPECT_EQ(desc.get_dims()[0], 1);
    EXPECT_EQ(desc.get_dims()[1], 64);
    EXPECT_EQ(desc.get_dims()[2], 1);
    EXPECT_EQ(desc.get_dims()[3], 1);
}

TEST_F(test_layout_to_memory_desc, builder_method_chaining_flatten_and_format) {
    layout l = layout{ov::PartialShape{2, 3, 4, 5}, data_types::f32, format::bfyx};
    auto desc = layout_to_memory_desc_flatten(l, dnnl::memory::format_tag::ab);
    EXPECT_EQ(desc.get_ndims(), 2);
    EXPECT_EQ(desc.get_dims()[0], 1);
    EXPECT_EQ(desc.get_dims()[1], 2 * 3 * 4 * 5);
}

TEST_F(test_layout_to_memory_desc, builder_blocked_with_undef_format) {
    layout l = layout{ov::PartialShape{1, 64, 56, 56}, data_types::f32, format::bfyx};
    auto desc = layout_to_memory_desc_blocked(l, dnnl::memory::format_tag::undef);
    EXPECT_EQ(desc.get_ndims(), 4);
    EXPECT_EQ(get_format_tag_from_desc(desc), dnnl::memory::format_tag::abcd);
}

TEST_F(test_layout_to_memory_desc, get_conv_memory_descs_query_formats) {
    layout input_layout = layout{ov::PartialShape{1, 3, 224, 224}, data_types::f32, format::bfyx};
    layout weights_layout = layout{ov::PartialShape{64, 3, 7, 7}, data_types::f32, format::oiyx};
    layout output_layout = layout{ov::PartialShape{1, 64, 112, 112}, data_types::f32, format::bfyx};

    // format_tag::any is used in ConvolutionImplementationManager::query_formats().
    auto [input_desc, weights_desc, output_desc] = get_conv_memory_descs(
        input_layout, weights_layout, output_layout, dnnl::memory::format_tag::any);

    EXPECT_EQ(input_desc.get_ndims(), 4);
    EXPECT_EQ(weights_desc.get_ndims(), 4);
    EXPECT_EQ(output_desc.get_ndims(), 4);
    EXPECT_EQ(get_format_tag_from_desc(input_desc), dnnl::memory::format_tag::undef);
    EXPECT_EQ(get_format_tag_from_desc(weights_desc), dnnl::memory::format_tag::undef);
    EXPECT_EQ(get_format_tag_from_desc(output_desc), dnnl::memory::format_tag::undef);
}

TEST_F(test_layout_to_memory_desc, get_conv_memory_descs_blocked_format) {
    layout input_layout = layout{ov::PartialShape{1, 64, 56, 56}, data_types::f16, format::b_fs_yx_fsv16};
    layout weights_layout = layout{ov::PartialShape{128, 64, 3, 3}, data_types::f16, format::os_is_yx_osv16_isv16};
    layout output_layout = layout{ov::PartialShape{1, 128, 56, 56}, data_types::f16, format::b_fs_yx_fsv16};

    auto [input_desc, weights_desc, output_desc] = get_conv_memory_descs(
        input_layout, weights_layout, output_layout, dnnl::memory::format_tag::undef);

    EXPECT_EQ(input_desc.get_ndims(), 4);
    EXPECT_EQ(output_desc.get_ndims(), 4);
    // Should be blocked
    EXPECT_EQ(input_desc.get_inner_nblks(), 1);
    EXPECT_EQ(output_desc.get_inner_nblks(), 1);
}

TEST_F(test_layout_to_memory_desc, get_conv_memory_descs_grouped_weights) {
    layout input_layout = layout{ov::PartialShape{1, 256, 28, 28}, data_types::f32, format::bfyx};
    layout weights_layout = layout{ov::PartialShape{32, 128, 8, 3, 3}, data_types::f32, format::goiyx};
    layout output_layout = layout{ov::PartialShape{1, 256, 28, 28}, data_types::f32, format::bfyx};

    auto [input_desc, weights_desc, output_desc] = get_conv_memory_descs(
        input_layout, weights_layout, output_layout, dnnl::memory::format_tag::undef);

    EXPECT_EQ(weights_desc.get_ndims(), 5);  // Grouped weights have 5 dims
}

TEST_F(test_layout_to_memory_desc, regression_3d_shape_format_selection) {
    // Regression test: 3D shapes should use abc/acb format correctly
    layout l_bfyx = layout{ov::PartialShape{1, 256, 64}, data_types::f16, format::bfyx};
    layout l_byxf = layout{ov::PartialShape{1, 256, 64}, data_types::f16, format::byxf};

    auto desc_bfyx = layout_to_memory_desc(l_bfyx, dnnl::memory::format_tag::undef);
    auto desc_byxf = layout_to_memory_desc(l_byxf, dnnl::memory::format_tag::undef);

    EXPECT_EQ(desc_bfyx.get_ndims(), 3);
    EXPECT_EQ(desc_byxf.get_ndims(), 3);
    EXPECT_EQ(get_format_tag_from_desc(desc_bfyx), dnnl::memory::format_tag::abc);
    EXPECT_EQ(get_format_tag_from_desc(desc_byxf), dnnl::memory::format_tag::acb);
    // Format tags should be different (abc vs acb)
}
