/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#include <common/test.h>
#include <algorithm>

#include "ucp_datatype.h"

extern "C" {
#include <ucp/dt/dt.h>
#include <ucp/dt/datatype_iter.inl>
#include <ucp/core/ucp_rkey.h>
}

class test_ucp_dt_iov : public ucs::test {
protected:
    size_t calc_iov_offset(const ucp_dt_iov_t *iov, size_t iov_indx, size_t iov_offs) {
        size_t offset = iov_offs;
        for (size_t i = 0; i < iov_indx; ++i) {
            offset += iov[i].length;
        }
        return offset;
   }
};

UCS_TEST_F(test_ucp_dt_iov, seek)
{
    for (int count = 0; count < 100; ++count) {
        size_t iovcnt = (ucs::rand() % 20) + 1;
        std::vector<ucp_dt_iov_t> iov(iovcnt);

        size_t total_size = 0;
        for (size_t i = 0; i < iovcnt; ++i) {
            iov[i].length = (ucs::rand() % 1000) + 1;
            total_size   += iov[i].length;
        }

        ASSERT_EQ(total_size, calc_iov_offset(&iov[0], iovcnt, 0));

        size_t offset = 0;
        size_t iov_offs = 0, iov_indx = 0;
        for (int j = 0; j < 100; ++j) {
            size_t new_offset = ucs::rand() % total_size;
            ucp_dt_iov_seek(&iov[0], iovcnt,
                            (ptrdiff_t)new_offset - (ptrdiff_t)offset,
                            &iov_offs, &iov_indx);
            EXPECT_EQ(new_offset, calc_iov_offset(&iov[0], iov_indx, iov_offs));
            offset = new_offset;
        }
    }
};

class test_ucp_dt_iter : public ucs::test_with_param<ucp_datatype_t> {
protected:
    virtual void init() {
        ucp_params_t ctx_params;
        ctx_params.field_mask = UCP_PARAM_FIELD_FEATURES;
        ctx_params.features   = UCP_FEATURE_TAG;
        UCS_TEST_CREATE_HANDLE(ucp_context_h, m_ucph, ucp_cleanup, ucp_init,
                               &ctx_params, NULL);
    }

    virtual void cleanup() {
        m_ucph.reset();
    }

    void init_dt_iter(size_t size, bool is_pack)
    {
        m_dt_buffer.resize(size, 0);
        ucs::fill_random(m_dt_buffer);

        m_packed_buffer.resize(size, 0);
        ucs::fill_random(m_packed_buffer);

        size_t iovcnt;
        if (GetParam() == UCP_DATATYPE_IOV) {
            if (size == 0) {
                iovcnt = 0;
            } else {
                iovcnt = std::min(static_cast<size_t>((ucs::rand() % 20) + 1),
                                  m_dt_buffer.size());
            }
        } else {
            iovcnt = 1;
        }

        m_dt_desc.make(GetParam(), &m_dt_buffer[0], m_dt_buffer.size(), iovcnt);

        uint8_t sg_count;
        /* Pass empty param argument to disable memh initialization */
        ucp_request_param_t param;
        param.op_attr_mask = 0;

        ucs_status_t status = ucp_datatype_iter_init(
                m_ucph.get(), m_dt_desc.buf(), m_dt_desc.count(),
                m_dt_desc.dt(), m_dt_buffer.size(), is_pack, &m_dt_iter,
                &sg_count, &param);
        ASSERT_UCS_OK(status);
        if (!UCP_DT_IS_GENERIC(GetParam())) {
            EXPECT_EQ(iovcnt, sg_count);
        }

        ucp_md_map_t md_map = m_ucph->reg_md_map[UCS_MEMORY_TYPE_HOST] &
                              m_ucph->cache_md_map[UCS_MEMORY_TYPE_HOST];
        status = ucp_datatype_iter_mem_reg(m_ucph, &m_dt_iter, md_map, 0,
                                           UINT_MAX);
        ASSERT_UCS_OK(status);

        UCS_STRING_BUFFER_ONSTACK(strb, 64);
        ucp_datatype_iter_str(&m_dt_iter, &strb);
        UCS_TEST_MESSAGE << ucs_string_buffer_cstr(&strb);
    }

    void finalize_dt_iter()
    {
        EXPECT_EQ(m_dt_buffer, m_packed_buffer);
        ucp_datatype_iter_cleanup(&m_dt_iter, 1, UINT_MAX);
    }

    size_t random_seg_size() const
    {
        return ucs::rand() % (m_dt_buffer.size() / 2);
    }

    void test_pack(size_t size)
    {
        init_dt_iter(size, true);

        while (!ucp_datatype_iter_is_end(&m_dt_iter)) {
            EXPECT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
            size_t seg_size  = random_seg_size();
            void *packed_ptr = UCS_PTR_BYTE_OFFSET(&m_packed_buffer[0],
                                                   m_dt_iter.offset);
            /* TODO create non-NULL worker when using memtype */
            ucp_datatype_iter_t next_iter;
            ucp_datatype_iter_next_pack(&m_dt_iter, NULL, seg_size, &next_iter,
                                        packed_ptr);
            ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
        }

        finalize_dt_iter();
    }

    void test_unpack(size_t size)
    {
        init_dt_iter(size, false);

        typedef std::vector< std::pair<size_t, size_t> > segment_vector_t;
        segment_vector_t segments;
        size_t offset = 0;
        while (offset < m_dt_buffer.size()) {
            size_t seg_size = ucs_min(random_seg_size(),
                                      m_dt_buffer.size() - offset);
            segments.push_back(std::make_pair(offset, seg_size));
            offset += seg_size;
        }
        std::random_shuffle(segments.begin(), segments.end(), ucs::rand_range);

        for (segment_vector_t::iterator it = segments.begin();
             it != segments.end(); ++it) {
            size_t offset    = it->first;
            void *packed_ptr = UCS_PTR_BYTE_OFFSET(&m_packed_buffer[0], offset);
            /* TODO create non-NULL worker when using memtype */
            ucp_datatype_iter_unpack(&m_dt_iter, NULL, it->second, offset,
                                     packed_ptr);
        }

        finalize_dt_iter();
    }

public:
    static std::vector<ucp_datatype_t> enum_dt_generic_params() {
        ucp_datatype_t datatype;
        ucs_status_t status;

        if (dt_gen == 0) {
            status = ucp_dt_create_generic(&ucp::test_dt_copy_ops, NULL, &datatype);
            if (status != UCS_OK) {
                return std::vector<ucp_datatype_t>();
            }

            /* keep global pointer to dt_gen to silence valgrind leak checker */
            dt_gen = ucp_dt_to_generic(datatype);
        }

        return std::vector<ucp_datatype_t>(1, ucp_dt_from_generic(dt_gen));
    }

private:
    static ucp_dt_generic_t   *dt_gen;
    ucs::handle<ucp_context_h> m_ucph;
    std::string                m_dt_buffer;
    std::string                m_packed_buffer;
    ucp::data_type_desc_t      m_dt_desc;
    ucp_datatype_iter_t        m_dt_iter;
};

ucp_dt_generic_t* test_ucp_dt_iter::dt_gen = 0;

UCS_TEST_P(test_ucp_dt_iter, datatype_query) {
    ucs_status_t status;
    ucp_datatype_t datatype           = GetParam();
    ucp_datatype_attr_t datatype_attr = {};
    datatype_attr.field_mask          = UCP_DATATYPE_ATTR_FIELD_PACKED_SIZE;

    switch (datatype & UCP_DATATYPE_CLASS_MASK) {
    case UCP_DATATYPE_CONTIG:
        status = ucp_dt_query(datatype, &datatype_attr);
        ASSERT_EQ(UCS_OK, status);
        EXPECT_GT(datatype_attr.packed_size, 0);
        break;
    case UCP_DATATYPE_GENERIC:
    {
        const size_t count        = 2;
        datatype_attr.count       = count;
        datatype_attr.field_mask |= UCP_DATATYPE_ATTR_FIELD_COUNT;

        int32_t buf[count]        = {1, 2};
        datatype_attr.buffer      = buf;
        datatype_attr.field_mask |= UCP_DATATYPE_ATTR_FIELD_BUFFER;

        status = ucp_dt_query(datatype, &datatype_attr);
        ASSERT_EQ(UCS_OK, status);
        EXPECT_GT(datatype_attr.packed_size, 0);
        break;
    }
    case UCP_DATATYPE_IOV:
    {
        const size_t count        = 2;
        datatype_attr.count       = count;
        datatype_attr.field_mask |= UCP_DATATYPE_ATTR_FIELD_COUNT;

        ucp_dt_iov_t buf[count]   = {{NULL, 8}, {NULL, 4}};
        datatype_attr.buffer      = buf;
        datatype_attr.field_mask |= UCP_DATATYPE_ATTR_FIELD_BUFFER;

        status = ucp_dt_query(datatype, &datatype_attr);
        ASSERT_EQ(UCS_OK, status);
        EXPECT_EQ(12, datatype_attr.packed_size);
        break;
    }
    default:
        break;
    }
}

UCS_TEST_P(test_ucp_dt_iter, pack_100b) {
    test_pack(100);
}

UCS_TEST_P(test_ucp_dt_iter, pack_0) {
    test_pack(0);
}

UCS_TEST_P(test_ucp_dt_iter, pack_1MB) {
    test_pack(UCS_MBYTE + (ucs::rand() % UCS_KBYTE));
}

UCS_TEST_P(test_ucp_dt_iter, unpack_100b) {
    test_unpack(100);
}

UCS_TEST_P(test_ucp_dt_iter, unpack_0) {
    test_unpack(0);
}

UCS_TEST_P(test_ucp_dt_iter, unpack_1MB) {
    test_unpack(UCS_MBYTE + (ucs::rand() % UCS_KBYTE));
}

INSTANTIATE_TEST_SUITE_P(contig, test_ucp_dt_iter,
                        testing::Values(ucp_dt_make_contig(1),
                                        ucp_dt_make_contig(8),
                                        ucp_dt_make_contig(39)));

INSTANTIATE_TEST_SUITE_P(iov, test_ucp_dt_iter,
                        testing::Values(ucp_dt_make_iov()));

INSTANTIATE_TEST_SUITE_P(generic, test_ucp_dt_iter,
                        testing::ValuesIn(test_ucp_dt_iter::enum_dt_generic_params()));

class test_ucp_dt_sgl : public ucs::test {
protected:
    virtual void init() {
        ucp_params_t ctx_params;
        ctx_params.field_mask = UCP_PARAM_FIELD_FEATURES;
        ctx_params.features   = UCP_FEATURE_RMA;
        UCS_TEST_CREATE_HANDLE(ucp_context_h, m_ucph, ucp_cleanup, ucp_init,
                               &ctx_params, NULL);
    }

    virtual void cleanup() {
        m_ucph.reset();
    }

    void init_sgl_iter(size_t count) {
        m_buffers.resize(count);
        m_lengths.resize(count);
        m_remote_addrs.resize(count);
        m_rkeys.resize(count);

        m_dummy_rkey            = {};
        m_dummy_rkey.cfg_index  = 0;
#if ENABLE_PARAMS_CHECK
        m_dummy_rkey.ep         = nullptr;
#endif

        for (size_t i = 0; i < count; i++) {
            m_buffers[i]      = &m_buffers[i];
            m_lengths[i]      = (i + 1) * 64;
            m_remote_addrs[i] = 0x1000 + i * 0x100;
            m_rkeys[i]        = &m_dummy_rkey;
        }

        m_local = {};
        m_local.field_mask = UCP_DT_LOCAL_SGL_FIELD_BUFFERS |
                             UCP_DT_LOCAL_SGL_FIELD_LENGTHS;
        m_local.buffers    = m_buffers.data();
        m_local.lengths    = m_lengths.data();

        m_remote = {};
        m_remote.field_mask   = UCP_DT_REMOTE_SGL_FIELD_REMOTE_ADDRS |
                                UCP_DT_REMOTE_SGL_FIELD_RKEYS;
        m_remote.remote_addrs = m_remote_addrs.data();
        m_remote.rkeys        = m_rkeys.data();

        ucp_request_param_t param = {};
        param.op_attr_mask = UCP_OP_ATTR_FIELD_DATATYPE |
                             UCP_OP_ATTR_FIELD_REMOTE;
        param.datatype     = ucp_dt_make_sgl();
        param.remote       = &m_remote;

        ucs_status_t status = ucp_datatype_iter_sgl_init(m_ucph.get(),
                                                         &m_dt_iter, &m_local,
                                                         &m_remote, count,
                                                         &param);
        ASSERT_UCS_OK(status);
    }

private:
    ucs::handle<ucp_context_h> m_ucph;
    std::vector<void*>         m_buffers;
    std::vector<size_t>        m_lengths;
    std::vector<uint64_t>      m_remote_addrs;
    std::vector<ucp_rkey_h>    m_rkeys;
    ucp_rkey_t                 m_dummy_rkey;
    ucp_dt_local_sgl_t         m_local;
    ucp_dt_remote_sgl_t        m_remote;

protected:
    ucp_datatype_iter_t        m_dt_iter;
};

UCS_TEST_F(test_ucp_dt_sgl, iter_next_chunked) {
    static constexpr size_t NUM_ELEMS    = 10;
    static constexpr size_t MAX_PER_STEP = 3;

    init_sgl_iter(NUM_ELEMS);

    size_t total_advanced = 0;
    while (!ucp_datatype_iter_is_end(&m_dt_iter)) {
        ucp_datatype_iter_t next_iter = {};
        size_t advanced = ucp_datatype_iter_next_sgl(&m_dt_iter,
                                                     MAX_PER_STEP,
                                                     &next_iter);
        EXPECT_LE(advanced, MAX_PER_STEP);
        EXPECT_GT(advanced, 0u);
        total_advanced += advanced;
        ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
    }

    EXPECT_EQ(NUM_ELEMS, total_advanced);
    EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
}

UCS_TEST_F(test_ucp_dt_sgl, iter_next_single_step) {
    static constexpr size_t counts[] = {1, 5, 10};
    for (size_t count : counts) {
        init_sgl_iter(count);

        ucp_datatype_iter_t next_iter = {};
        size_t advanced = ucp_datatype_iter_next_sgl(&m_dt_iter, count,
                                                     &next_iter);
        EXPECT_EQ(count, advanced);
        ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
        EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
    }
}

UCS_TEST_F(test_ucp_dt_sgl, iter_next_one_by_one) {
    static constexpr size_t NUM_ELEMS = 5;

    init_sgl_iter(NUM_ELEMS);

    for (size_t i = 0; i < NUM_ELEMS; i++) {
        EXPECT_FALSE(ucp_datatype_iter_is_end(&m_dt_iter));
        ucp_datatype_iter_t next_iter = {};
        size_t advanced = ucp_datatype_iter_next_sgl(&m_dt_iter, 1,
                                                     &next_iter);
        EXPECT_EQ(1u, advanced);
        EXPECT_EQ(i + 1, next_iter.offset);
        ucp_datatype_iter_copy_position(&m_dt_iter, &next_iter, UINT_MAX);
    }

    EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
}

UCS_TEST_F(test_ucp_dt_sgl, init_zero_count) {
    init_sgl_iter(0);

    EXPECT_TRUE(ucp_datatype_iter_is_end(&m_dt_iter));
    EXPECT_EQ(UCS_MEMORY_TYPE_HOST, m_dt_iter.mem_info.type);
    EXPECT_EQ(UCS_SYS_DEVICE_ID_UNKNOWN, m_dt_iter.mem_info.sys_dev);
}
