#include "snapshot_manager.h"

#include <library/cpp/testing/unittest/registar.h>

#include <google/protobuf/util/message_differencer.h>

#include <util/generic/array_ref.h>
#include <util/string/join.h>
#include <util/system/src_location.h>

namespace {

    class TDummySnapshotManager : public IIndexSnapshotManager {
        using IIndexSnapshotManager::IIndexSnapshotManager;

        TDummySnapshotManager() = default;

    protected:
        void DoPublishSnapshot(const TSnapshot& snapshot) override {
            Snapshots.push_back(snapshot);
        }
        TSnapshots DoGetSnapshots(const std::function<bool(TInstant)>& filter, const IStopCondition&) override {
            TSnapshots filtered(Reserve(Snapshots.size()));
            for (const auto& snapshot : Snapshots) {
                if (filter(snapshot.Timestamp)) {
                    filtered.push_back(snapshot);
                }
            }
            return filtered;
        }

    private:
        TSnapshots Snapshots;
    };

    NRTYServer::TShardResource MakeResource(const TString& name, NUtil::TInterval<ui64> i, ui64 timestamp) {
        NRTYServer::TShardResource res;
        res.SetName(name);
        res.SetShardMin(i.GetMin());
        res.SetShardMax(i.GetMax());
        res.SetTimestamp(timestamp);
        return res;
    }

    NRTYServer::TShards MakeShards(const TString& baseName, TConstArrayRef<NUtil::TInterval<ui64>> intervals, ui64 timestamp) {
        NRTYServer::TShards s;
        for (auto const &i: intervals) {
            *s.AddShard() = MakeResource(Join('_', baseName, timestamp, i.GetMin(), i.GetMax()), i, timestamp);
        }
        return s;
    }

    IIndexSnapshotManager::TSnapshot MakeSnapshot(const TString& baseName, TConstArrayRef<NUtil::TInterval<ui64>> intervals, ui64 timestamp) {
        return {TInstant::Seconds(timestamp), MakeShards(baseName, intervals, timestamp)};
    }

    static const ui64 ThreeTimestamps[] = { (1551015835 - 100) , 1551015835, (1551015835 + 100) };
    static const ui64 OtherTwoTimestamps[] = { (1551015835 - 50), (1551015835 + 50) };

    static const NUtil::TInterval<ui64> OneShard[] = { {0, 65533} };
    static const NUtil::TInterval<ui64> TwoShards[] = { {0, 32765}, {32766, 65533} };

    static const IIndexSnapshotManager::TSnapshots OneShardThreeSnapshots{
        MakeSnapshot("OneShardThreeSnapshots", OneShard, ThreeTimestamps[0]),
        MakeSnapshot("OneShardThreeSnapshots", OneShard, ThreeTimestamps[1]),
        MakeSnapshot("OneShardThreeSnapshots", OneShard, ThreeTimestamps[2])
    };

    static const IIndexSnapshotManager::TSnapshots TwoShardsThreeSnapshots{
        MakeSnapshot("TwoShardsThreeSnapshots", TwoShards, ThreeTimestamps[0]),
        MakeSnapshot("TwoShardsThreeSnapshots", TwoShards, ThreeTimestamps[1]),
        MakeSnapshot("TwoShardsThreeSnapshots", TwoShards, ThreeTimestamps[2])
    };

    // indices of entries of this array in descending order: { 2, 4, 1, 3, 0 }
    static const IIndexSnapshotManager::TSnapshots MixedShardsFiveSnapshots{
        MakeSnapshot("MixedShardsFiveSnapshots", OneShard, ThreeTimestamps[0]),
        MakeSnapshot("MixedShardsFiveSnapshots", OneShard, ThreeTimestamps[1]),
        MakeSnapshot("MixedShardsFiveSnapshots", OneShard, ThreeTimestamps[2]),

        MakeSnapshot("MixedShardsFiveSnapshots", TwoShards, OtherTwoTimestamps[0]),
        MakeSnapshot("MixedShardsFiveSnapshots", TwoShards, OtherTwoTimestamps[1])
    };

    void Publish(IIndexSnapshotManager& im, const IIndexSnapshotManager::TSnapshots& snapshots) {
        for (const auto &s: snapshots) {
            im.PublishSnapshot(s);
        }
    }

    void AssertCorrectSnapshots(const IIndexSnapshotManager::TSnapshots& source, const IIndexSnapshotManager::TSnapshots& result, TConstArrayRef<size_t> indices, TSourceLocation&& loc) {
        UNIT_ASSERT_C(source.size() >= result.size(), "invalid result size: " << loc);
        UNIT_ASSERT_C(result.size() == indices.size(), "invalid result size: " << loc);
        for (size_t j = 0; j < indices.size(); ++j) {
            UNIT_ASSERT_C(source[indices[j]].Timestamp == result[j].Timestamp, loc);
            UNIT_ASSERT_C(::google::protobuf::util::MessageDifferencer::Equals(source[indices[j]].Shards, result[j].Shards), loc);
        }
    }

    void AssertCorrectShardResources(const IIndexSnapshotManager::TSnapshots& source, TConstArrayRef<IIndexSnapshotManager::TShardResource> result, TConstArrayRef<size_t> indices, NUtil::TInterval<ui64> interval, TSourceLocation&& loc) {
        UNIT_ASSERT_C(source.size() >= result.size(), "invalid result size: " << loc);
        UNIT_ASSERT_C(result.size() == indices.size(), "invalid result size: " << loc);
        for (size_t j = 0; j < indices.size(); ++j) {
            UNIT_ASSERT_C(source[indices[j]].Timestamp.Seconds() == result[j].GetTimestamp(), "invalid timestamp: " << loc);
            UNIT_ASSERT_C(result[j].GetShardMin() == interval.GetMin(), "invalid ShardMin: " << loc);
            UNIT_ASSERT_C(result[j].GetShardMax() == interval.GetMax(), "invalid ShardMax: " << loc);
            bool match = false;
            for (const auto& shardResource : source[indices[j]].Shards.GetShard()) {
                if (::google::protobuf::util::MessageDifferencer::Equals(shardResource, result[j])) {
                    match = true;
                }
            }
            UNIT_ASSERT_C(match, loc);
        }
    }
}

Y_UNIT_TEST_SUITE(SnapshotManagerSuite) {
    Y_UNIT_TEST(TestEmpty) {
        TDummySnapshotManager m(IIndexSnapshotManager::TContext{});

        UNIT_ASSERT(std::get<0>(m.GetSnapshots()).empty());
        UNIT_ASSERT(m.GetSnapshotsFixed(TInstant::Seconds(ThreeTimestamps[0])).empty());

        UNIT_ASSERT(m.GetResourcesForShard(OneShard[0]).empty());
        UNIT_ASSERT(m.GetResourcesForShard(OneShard[0]).empty());
        UNIT_ASSERT(m.GetResourcesForShard(OneShard[0], TInstant::Now()).empty());
        UNIT_ASSERT(m.GetResourcesForShard(OneShard[0], TDuration::Zero()).empty());
        UNIT_ASSERT(m.GetResourcesForShard(OneShard[0], TDuration::Max()).empty());
        UNIT_ASSERT(m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[0])).empty());
    }

    Y_UNIT_TEST(TestWrongShard) {
        TDummySnapshotManager m(IIndexSnapshotManager::TContext{});
        Publish(m, OneShardThreeSnapshots);

        UNIT_ASSERT(m.GetResourcesForShard(TwoShards[0]).empty());
        UNIT_ASSERT(m.GetResourcesForShard(TwoShards[0]).empty());
        UNIT_ASSERT(m.GetResourcesForShard(TwoShards[0], TInstant::Now()).empty());
        UNIT_ASSERT(m.GetResourcesForShard(TwoShards[0], TDuration::Zero()).empty());
        UNIT_ASSERT(m.GetResourcesForShard(TwoShards[0], TDuration::Max()).empty());

        UNIT_ASSERT(m.GetResourceForShardFixed(TwoShards[0], TInstant::Seconds(ThreeTimestamps[0])).empty());
    }

    Y_UNIT_TEST(TestOneShard) {
        TDummySnapshotManager m(IIndexSnapshotManager::TContext{});
        Publish(m, OneShardThreeSnapshots);

        AssertCorrectSnapshots(OneShardThreeSnapshots, std::get<0>(m.GetSnapshots()), {2, 1, 0}, __LOCATION__);
        AssertCorrectSnapshots(OneShardThreeSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[0]))), {2, 1}, __LOCATION__);
        AssertCorrectSnapshots(OneShardThreeSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[1]))), {2}, __LOCATION__);
        AssertCorrectSnapshots(OneShardThreeSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[2]))), {}, __LOCATION__);

        AssertCorrectSnapshots(OneShardThreeSnapshots, m.GetSnapshotsFixed(TInstant::Seconds(ThreeTimestamps[0])), {0}, __LOCATION__);
        AssertCorrectSnapshots(OneShardThreeSnapshots, m.GetSnapshotsFixed(TInstant::Seconds(ThreeTimestamps[1])), {1}, __LOCATION__);
        AssertCorrectSnapshots(OneShardThreeSnapshots, m.GetSnapshotsFixed(TInstant::Seconds(ThreeTimestamps[2])), {2}, __LOCATION__);

        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourcesForShard(OneShard[0]), {2, 1, 0}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourcesForShard(OneShard[0], TInstant::Seconds(ThreeTimestamps[0])), {2, 1}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourcesForShard(OneShard[0], TInstant::Seconds(ThreeTimestamps[1])), {2}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourcesForShard(OneShard[0], TInstant::Seconds(ThreeTimestamps[2])), {}, OneShard[0], __LOCATION__);

        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[0])), {0}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[1])), {1}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[2])), {2}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(OneShardThreeSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(OtherTwoTimestamps[0])), {}, OneShard[0], __LOCATION__);

    }

    Y_UNIT_TEST(TestTwoShards) {
        TDummySnapshotManager m(IIndexSnapshotManager::TContext{});
        Publish(m, TwoShardsThreeSnapshots);

        AssertCorrectSnapshots(TwoShardsThreeSnapshots, std::get<0>(m.GetSnapshots()), {2, 1, 0}, __LOCATION__);
        AssertCorrectSnapshots(TwoShardsThreeSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[0]))), {2, 1}, __LOCATION__);
        AssertCorrectSnapshots(TwoShardsThreeSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[1]))), {2}, __LOCATION__);
        AssertCorrectSnapshots(TwoShardsThreeSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[2]))), {}, __LOCATION__);

        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[0]), {2, 1, 0}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[0], TInstant::Seconds(ThreeTimestamps[0])), {2, 1}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[0], TInstant::Seconds(ThreeTimestamps[1])), {2}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[0], TInstant::Seconds(ThreeTimestamps[2])), {}, TwoShards[0], __LOCATION__);

        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[1]), {2, 1, 0}, TwoShards[1], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[1], TInstant::Seconds(ThreeTimestamps[0])), {2, 1}, TwoShards[1], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[1], TInstant::Seconds(ThreeTimestamps[1])), {2}, TwoShards[1], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourcesForShard(TwoShards[1], TInstant::Seconds(ThreeTimestamps[2])), {}, TwoShards[1], __LOCATION__);

        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourceForShardFixed(TwoShards[0], TInstant::Seconds(ThreeTimestamps[1])), {1}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(TwoShardsThreeSnapshots, m.GetResourceForShardFixed(TwoShards[1], TInstant::Seconds(ThreeTimestamps[1])), {1}, TwoShards[1], __LOCATION__);
    }

    Y_UNIT_TEST(TestSnapshotsForMixedShards) {
        TDummySnapshotManager m(IIndexSnapshotManager::TContext{});
        Publish(m, MixedShardsFiveSnapshots);

        AssertCorrectSnapshots(MixedShardsFiveSnapshots, std::get<0>(m.GetSnapshots()), {2, 4, 1, 3, 0}, __LOCATION__);
        AssertCorrectSnapshots(MixedShardsFiveSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[0]))), {2, 4, 1, 3}, __LOCATION__);
        AssertCorrectSnapshots(MixedShardsFiveSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(OtherTwoTimestamps[0]))), {2, 4, 1}, __LOCATION__);
        AssertCorrectSnapshots(MixedShardsFiveSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[1]))), {2, 4}, __LOCATION__);
        AssertCorrectSnapshots(MixedShardsFiveSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(OtherTwoTimestamps[1]))), {2}, __LOCATION__);
        AssertCorrectSnapshots(MixedShardsFiveSnapshots, std::get<0>(m.GetSnapshots(TInstant::Seconds(ThreeTimestamps[2]))), {}, __LOCATION__);

        AssertCorrectSnapshots(MixedShardsFiveSnapshots, m.GetSnapshotsFixed(TInstant::Seconds(ThreeTimestamps[1])), {1}, __LOCATION__);
        AssertCorrectSnapshots(MixedShardsFiveSnapshots, m.GetSnapshotsFixed(TInstant::Seconds(OtherTwoTimestamps[1])), {4}, __LOCATION__);
    }

    Y_UNIT_TEST(TestResourcesForMixedShards) {
        TDummySnapshotManager m(IIndexSnapshotManager::TContext{});
        Publish(m, MixedShardsFiveSnapshots);

        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(OneShard[0]), {2, 1, 0}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(OneShard[0], TInstant::Seconds(ThreeTimestamps[0])), {2, 1}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(OneShard[0], TInstant::Seconds(ThreeTimestamps[1])), {2}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(OneShard[0], TInstant::Seconds(ThreeTimestamps[2])), {}, OneShard[0], __LOCATION__);

        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(TwoShards[0]), {4, 3}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(TwoShards[0], TInstant::Seconds(OtherTwoTimestamps[0])), {4}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourcesForShard(TwoShards[0], TInstant::Seconds(OtherTwoTimestamps[1])), {}, TwoShards[0], __LOCATION__);

        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[0])), {0}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[1])), {1}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(ThreeTimestamps[2])), {2}, OneShard[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(TwoShards[1], TInstant::Seconds(OtherTwoTimestamps[0])), {3}, TwoShards[1], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(TwoShards[0], TInstant::Seconds(OtherTwoTimestamps[1])), {4}, TwoShards[0], __LOCATION__);

        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(TwoShards[0], TInstant::Seconds(ThreeTimestamps[0])), {}, TwoShards[0], __LOCATION__);
        AssertCorrectShardResources(MixedShardsFiveSnapshots, m.GetResourceForShardFixed(OneShard[0], TInstant::Seconds(OtherTwoTimestamps[0])), {}, OneShard[0], __LOCATION__);
    }
}
