/**
 * MIT License
 *
 * Copyright (c) 2025 Huawei Technologies Co., Ltd. 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.
 * */
#include "cache/cc/trans_manager.h"
#include "detail/data_generator.h"
#include "detail/mock_store.h"
#include "detail/random.h"
#include "detail/types_helper.h"

class UCCacheTransManagerTest : public ::testing::Test {
public:
    UC::Test::Detail::Random rd;
    static UC::Detail::TaskHandle NextId()
    {
        static std::atomic<size_t> id{1};
        return id.fetch_add(1, std::memory_order_relaxed);
    }
};

TEST_F(UCCacheTransManagerTest, DumpThenLoad)
{
    using namespace UC::CacheStore;
    UC::Test::Detail::MockStore backend;
    EXPECT_CALL(backend, Dump).WillOnce(testing::Invoke(NextId));
    UC::Latch finish{};
    finish.Up();
    EXPECT_CALL(backend, Wait).WillOnce(testing::Invoke([&finish]() {
        finish.Done();
        return UC::Status::OK();
    }));
    Config config;
    config.storeBackend = &backend;
    size_t tensorSize = 32768;
    config.tensorSizes = {tensorSize};
    config.shardSize = tensorSize;
    config.blockSize = config.shardSize;
    config.deviceId = 0;
    config.bufferCapacity = config.shardSize * 1024;
    config.uniqueId = rd.RandomString(10);
    config.shareBufferEnable = true;
    TransBuffer buffer;
    auto s = buffer.Setup(config);
    ASSERT_EQ(s, UC::Status::OK());
    TransManager transMgr;
    s = transMgr.Setup(config, &buffer);
    ASSERT_EQ(s, UC::Status::OK());
    auto block = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
    constexpr size_t nBlocks = 1;
    UC::Test::Detail::DataGenerator data1{nBlocks, config.blockSize};
    data1.GenerateRandom();
    UC::Detail::TaskDesc desc1;
    desc1.brief = "Dump";
    desc1.push_back(UC::Detail::Shard{block, 0, {data1.Buffer()}});
    auto handle1 = transMgr.Submit({TransTask::Type::DUMP, desc1});
    ASSERT_TRUE(handle1.HasValue());
    s = transMgr.Wait(handle1.Value());
    ASSERT_EQ(s, UC::Status::OK());
    UC::Test::Detail::DataGenerator data2{nBlocks, config.blockSize};
    data2.Generate();
    UC::Detail::TaskDesc desc2;
    desc2.brief = "Load";
    desc2.push_back(UC::Detail::Shard{block, 0, {data2.Buffer()}});
    auto handle2 = transMgr.Submit({TransTask::Type::LOAD, desc2});
    ASSERT_TRUE(handle2.HasValue());
    s = transMgr.Wait(handle2.Value());
    ASSERT_EQ(s, UC::Status::OK());
    ASSERT_EQ(data1.Compare(data2), 0);
    finish.Wait();
}

TEST_F(UCCacheTransManagerTest, DumpThenLoadWithLayerWise)
{
    using namespace UC::Test::Detail;
    using namespace UC::CacheStore;
    constexpr size_t tensorSize = 32768;
    constexpr size_t shardSize = tensorSize;
    constexpr size_t layerNumber = 16;
    constexpr size_t blockSize = shardSize * layerNumber;
    constexpr size_t blockNumber = 5;
    MockStore backend;
    EXPECT_CALL(backend, Dump).WillRepeatedly(testing::Invoke(NextId));
    UC::Latch finish{};
    finish.Set(layerNumber);
    EXPECT_CALL(backend, Wait).WillRepeatedly(testing::Invoke([&finish]() {
        finish.Done();
        return UC::Status::OK();
    }));
    Config config;
    config.storeBackend = &backend;
    config.tensorSizes = {tensorSize};
    config.shardSize = shardSize;
    config.blockSize = blockSize;
    config.deviceId = 0;
    config.bufferCapacity = shardSize * 1024;
    config.uniqueId = rd.RandomString(10);
    config.shareBufferEnable = true;
    config.timeoutMs = 10 * 60 * 1000;
    TransBuffer buffer;
    auto s = buffer.Setup(config);
    ASSERT_EQ(s, UC::Status::OK());
    TransManager transMgr;
    s = transMgr.Setup(config, &buffer);
    ASSERT_EQ(s, UC::Status::OK());
    UC::Detail::BlockId blockIds[blockNumber];
    std::for_each_n(blockIds, blockNumber, [](auto& b) { b = TypesHelper::MakeBlockIdRandomly(); });
    DataGenerator data1{blockNumber, blockSize};
    data1.GenerateRandom();
    for (size_t i = 0; i < layerNumber; i++) {
        UC::Detail::TaskDesc desc;
        desc.brief = "Dump";
        for (size_t j = 0; j < blockNumber; j++) {
            auto addr = (void*)(((char*)data1.Buffer()) + blockSize * j + shardSize * i);
            desc.push_back(UC::Detail::Shard{blockIds[j], i, {addr}});
        }
        auto handle = transMgr.Submit({TransTask::Type::DUMP, desc});
        ASSERT_TRUE(handle.HasValue());
        s = transMgr.Wait(handle.Value());
        ASSERT_EQ(s.Underlying(), UC::Status::OK().Underlying());
    }
    finish.Wait();
    DataGenerator data2{blockNumber, blockSize};
    data2.Generate();
    for (size_t i = 0; i < layerNumber; i++) {
        UC::Detail::TaskDesc desc;
        desc.brief = "Load";
        for (size_t j = 0; j < blockNumber; j++) {
            auto addr = (void*)(((char*)data2.Buffer()) + blockSize * j + shardSize * i);
            desc.push_back(UC::Detail::Shard{blockIds[j], i, {addr}});
        }
        auto handle = transMgr.Submit({TransTask::Type::LOAD, desc});
        ASSERT_TRUE(handle.HasValue());
        s = transMgr.Wait(handle.Value());
        ASSERT_EQ(s.Underlying(), UC::Status::OK().Underlying());
    }
    ASSERT_EQ(data1.Compare(data2), 0);
}

TEST_F(UCCacheTransManagerTest, DumpThenLoadWithLayerAndChunk)
{
    using namespace UC::Test::Detail;
    using namespace UC::CacheStore;
    constexpr size_t tensorSize = 32768;
    constexpr size_t chunkNumber = 4;
    constexpr size_t shardSize = tensorSize * chunkNumber;
    constexpr size_t layerNumber = 16;
    constexpr size_t blockSize = shardSize * layerNumber;
    constexpr size_t blockNumber = 5;
    MockStore backend;
    EXPECT_CALL(backend, Dump).WillRepeatedly(testing::Invoke(NextId));
    UC::Latch finish{};
    finish.Set(layerNumber);
    EXPECT_CALL(backend, Wait).WillRepeatedly(testing::Invoke([&finish]() {
        finish.Done();
        return UC::Status::OK();
    }));
    Config config;
    config.storeBackend = &backend;
    config.tensorSizes.assign(chunkNumber, tensorSize);
    config.shardSize = shardSize;
    config.blockSize = blockSize;
    config.deviceId = 0;
    config.bufferCapacity = shardSize * 1024;
    config.uniqueId = rd.RandomString(10);
    config.shareBufferEnable = true;
    config.timeoutMs = 10 * 60 * 1000;
    TransBuffer buffer;
    auto s = buffer.Setup(config);
    ASSERT_EQ(s, UC::Status::OK());
    TransManager transMgr;
    s = transMgr.Setup(config, &buffer);
    ASSERT_EQ(s, UC::Status::OK());
    UC::Detail::BlockId blockIds[blockNumber];
    std::for_each_n(blockIds, blockNumber, [](auto& b) { b = TypesHelper::MakeBlockIdRandomly(); });
    DataGenerator data1{blockNumber, blockSize};
    data1.GenerateRandom();
    for (size_t i = 0; i < layerNumber; i++) {
        UC::Detail::TaskDesc desc;
        desc.brief = "Dump";
        for (size_t j = 0; j < blockNumber; j++) {
            UC::Detail::Shard shard;
            shard.owner = blockIds[j];
            shard.index = i;
            for (size_t k = 0; k < chunkNumber; k++) {
                auto addr = (void*)(((char*)data1.Buffer()) + blockSize * j + shardSize * i +
                                    tensorSize * k);
                shard.addrs.push_back(addr);
            }
            desc.push_back(std::move(shard));
        }
        auto handle = transMgr.Submit({TransTask::Type::DUMP, desc});
        ASSERT_TRUE(handle.HasValue());
        s = transMgr.Wait(handle.Value());
        ASSERT_EQ(s.Underlying(), UC::Status::OK().Underlying());
    }
    finish.Wait();
    DataGenerator data2{blockNumber, blockSize};
    data2.Generate();
    for (size_t i = 0; i < layerNumber; i++) {
        UC::Detail::TaskDesc desc;
        desc.brief = "Load";
        for (size_t j = 0; j < blockNumber; j++) {
            UC::Detail::Shard shard;
            shard.owner = blockIds[j];
            shard.index = i;
            for (size_t k = 0; k < chunkNumber; k++) {
                auto addr = (void*)(((char*)data2.Buffer()) + blockSize * j + shardSize * i +
                                    tensorSize * k);
                shard.addrs.push_back(addr);
            }
            desc.push_back(std::move(shard));
        }
        auto handle = transMgr.Submit({TransTask::Type::LOAD, desc});
        ASSERT_TRUE(handle.HasValue());
        s = transMgr.Wait(handle.Value());
        ASSERT_EQ(s.Underlying(), UC::Status::OK().Underlying());
    }
    ASSERT_EQ(data1.Compare(data2), 0);
}

TEST_F(UCCacheTransManagerTest, DumpThenLoadWithVariableLengthIo)
{
    using namespace UC::Test::Detail;
    using namespace UC::CacheStore;
    constexpr size_t tensorSize1 = 32768;
    constexpr size_t tensorSize2 = 4096;
    constexpr size_t shardSize = tensorSize1 + tensorSize2;
    constexpr size_t layerNumber = 16;
    constexpr size_t blockSize = shardSize * layerNumber;
    constexpr size_t blockNumber = 5;
    MockStore backend;
    EXPECT_CALL(backend, Dump).WillRepeatedly(testing::Invoke(NextId));
    UC::Latch finish{};
    finish.Set(layerNumber);
    EXPECT_CALL(backend, Wait).WillRepeatedly(testing::Invoke([&finish]() {
        finish.Done();
        return UC::Status::OK();
    }));
    Config config;
    config.storeBackend = &backend;
    config.tensorSizes = {tensorSize1, tensorSize2};
    config.shardSize = shardSize;
    config.blockSize = blockSize;
    config.deviceId = 0;
    config.bufferCapacity = shardSize * 1024;
    config.uniqueId = rd.RandomString(10);
    config.shareBufferEnable = true;
    TransBuffer buffer;
    auto s = buffer.Setup(config);
    ASSERT_EQ(s, UC::Status::OK());
    TransManager transMgr;
    s = transMgr.Setup(config, &buffer);
    ASSERT_EQ(s, UC::Status::OK());
    UC::Detail::BlockId blockIds[blockNumber];
    std::for_each_n(blockIds, blockNumber, [](auto& b) { b = TypesHelper::MakeBlockIdRandomly(); });
    DataGenerator data1{blockNumber, blockSize};
    data1.GenerateRandom();
    for (size_t i = 0; i < layerNumber; i++) {
        UC::Detail::TaskDesc desc;
        desc.brief = "Dump";
        for (size_t j = 0; j < blockNumber; j++) {
            auto addr1 = (void*)(((char*)data1.Buffer()) + blockSize * j + shardSize * i);
            auto addr2 = (void*)(((char*)addr1) + tensorSize1);
            desc.push_back(UC::Detail::Shard{
                blockIds[j], i, {addr1, addr2}
            });
        }
        auto handle = transMgr.Submit({TransTask::Type::DUMP, desc});
        ASSERT_TRUE(handle.HasValue());
        s = transMgr.Wait(handle.Value());
        ASSERT_EQ(s.Underlying(), UC::Status::OK().Underlying());
    }
    finish.Wait();
    DataGenerator data2{blockNumber, blockSize};
    data2.Generate();
    for (size_t i = 0; i < layerNumber; i++) {
        UC::Detail::TaskDesc desc;
        desc.brief = "Load";
        for (size_t j = 0; j < blockNumber; j++) {
            auto addr1 = (void*)(((char*)data2.Buffer()) + blockSize * j + shardSize * i);
            auto addr2 = (void*)(((char*)addr1) + tensorSize1);
            desc.push_back(UC::Detail::Shard{
                blockIds[j], i, {addr1, addr2}
            });
        }
        auto handle = transMgr.Submit({TransTask::Type::LOAD, desc});
        ASSERT_TRUE(handle.HasValue());
        s = transMgr.Wait(handle.Value());
        ASSERT_EQ(s.Underlying(), UC::Status::OK().Underlying());
    }
    ASSERT_EQ(data1.Compare(data2), 0);
}
