// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license.

#include "seal/serialization.h"
#include "seal/util/defines.h"
#include <fstream>
#include <functional>
#include <sstream>
#include <string>
#include "gtest/gtest.h"

using namespace seal;
using namespace std;

namespace sealtest
{
    namespace
    {
        struct test_struct
        {
            int a;
            int b;
            double c;

            void save_members(ostream &stream)
            {
                stream.write(reinterpret_cast<const char *>(&a), sizeof(int));
                stream.write(reinterpret_cast<const char *>(&b), sizeof(int));
                stream.write(reinterpret_cast<const char *>(&c), sizeof(double));
            }

            void load_members(istream &stream)
            {
                stream.read(reinterpret_cast<char *>(&a), sizeof(int));
                stream.read(reinterpret_cast<char *>(&b), sizeof(int));
                stream.read(reinterpret_cast<char *>(&c), sizeof(double));
            }

            streamoff save_size(compr_mode_type compr_mode) const
            {
                size_t members_size = Serialization::ComprSizeEstimate(sizeof(test_struct), compr_mode);

                return static_cast<streamoff>(sizeof(Serialization::SEALHeader) + members_size);
            }
        };
    } // namespace

    TEST(SerializationTest, IsValidHeader)
    {
        ASSERT_EQ(sizeof(Serialization::SEALHeader), Serialization::seal_header_size);

        Serialization::SEALHeader header;
        ASSERT_TRUE(Serialization::IsValidHeader(header));

#ifdef SEAL_USE_ZLIB
        header.compr_mode = compr_mode_type::zlib;
        ASSERT_TRUE(Serialization::IsValidHeader(header));
#endif

#ifdef SEAL_USE_ZSTD
        header.compr_mode = compr_mode_type::zstd;
        ASSERT_TRUE(Serialization::IsValidHeader(header));
#endif

        Serialization::SEALHeader invalid_header;
        invalid_header.magic = 0x1212;
        ASSERT_FALSE(Serialization::IsValidHeader(invalid_header));
        invalid_header.magic = Serialization::seal_magic;
        ASSERT_EQ(invalid_header.header_size, Serialization::seal_header_size);
        invalid_header.version_major = 0x02;
        ASSERT_FALSE(Serialization::IsValidHeader(invalid_header));
        invalid_header.version_major = SEAL_VERSION_MAJOR;
        invalid_header.compr_mode = (compr_mode_type)0x03;
        ASSERT_FALSE(Serialization::IsValidHeader(invalid_header));
    }

    TEST(SerializationTest, SEALHeaderSaveLoad)
    {
        {
            // Serialize to stream
            Serialization::SEALHeader header, loaded_header;
            header.compr_mode = Serialization::compr_mode_default;
            header.size = 256;

            stringstream stream;
            Serialization::SaveHeader(header, stream);
            ASSERT_TRUE(Serialization::IsValidHeader(header));
            Serialization::LoadHeader(stream, loaded_header);
            ASSERT_EQ(Serialization::seal_magic, loaded_header.magic);
            ASSERT_EQ(Serialization::seal_header_size, loaded_header.header_size);
            ASSERT_EQ(SEAL_VERSION_MAJOR, loaded_header.version_major);
            ASSERT_EQ(SEAL_VERSION_MINOR, loaded_header.version_minor);
            ASSERT_EQ(Serialization::compr_mode_default, loaded_header.compr_mode);
            ASSERT_EQ(0x00, loaded_header.reserved);
            ASSERT_EQ(256, loaded_header.size);
        }
        {
            // Serialize to buffer
            Serialization::SEALHeader header, loaded_header;
            header.compr_mode = Serialization::compr_mode_default;
            header.size = 256;

            vector<seal_byte> buffer(16);
            Serialization::SaveHeader(header, buffer.data(), buffer.size());
            ASSERT_TRUE(Serialization::IsValidHeader(header));
            Serialization::LoadHeader(buffer.data(), buffer.size(), loaded_header);
            ASSERT_EQ(Serialization::seal_magic, loaded_header.magic);
            ASSERT_EQ(Serialization::seal_header_size, loaded_header.header_size);
            ASSERT_EQ(SEAL_VERSION_MAJOR, loaded_header.version_major);
            ASSERT_EQ(SEAL_VERSION_MINOR, loaded_header.version_minor);
            ASSERT_EQ(Serialization::compr_mode_default, loaded_header.compr_mode);
            ASSERT_EQ(0x00, loaded_header.reserved);
            ASSERT_EQ(256, loaded_header.size);
        }
    }
    /*
        TEST(SerializationTest, SEALHeaderUpgrade)
        {
            legacy_headers::SEALHeader_3_4 header_3_4;
            header_3_4.compr_mode = Serialization::compr_mode_default;
            header_3_4.size = 0xF3F3;

            {
                Serialization::SEALHeader header;
                Serialization::LoadHeader(
                    reinterpret_cast<const seal_byte *>(&header_3_4), sizeof(legacy_headers::SEALHeader_3_4), header);
                ASSERT_TRUE(Serialization::IsValidHeader(header));
                ASSERT_EQ(header_3_4.compr_mode, header.compr_mode);
                ASSERT_EQ(header_3_4.size, header.size);
            }
            {
                Serialization::SEALHeader header;
                Serialization::LoadHeader(
                    reinterpret_cast<const seal_byte *>(&header_3_4), sizeof(legacy_headers::SEALHeader_3_4), header,
                    false);

                // No upgrade requested
                ASSERT_FALSE(Serialization::IsValidHeader(header));
            }
        }
    */
    TEST(SerializationTest, SaveLoadToStream)
    {
        test_struct st{ 3, ~0, 3.14159 }, st2;
        using namespace placeholders;
        stringstream stream;

        auto out_size = Serialization::Save(
            bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::none), stream,
            compr_mode_type::none, false);
        auto in_size = Serialization::Load(bind(&test_struct::load_members, &st2, _1), stream, false);
        ASSERT_EQ(out_size, in_size);
        ASSERT_EQ(st.a, st2.a);
        ASSERT_EQ(st.b, st2.b);
        ASSERT_EQ(st.c, st2.c);
#ifdef SEAL_USE_ZSTD
        {
            test_struct st3;
            out_size = Serialization::Save(
                bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::zstd), stream,
                compr_mode_type::zstd, false);
            in_size = Serialization::Load(bind(&test_struct::load_members, &st3, _1), stream, false);
            ASSERT_EQ(out_size, in_size);
            ASSERT_EQ(st.a, st3.a);
            ASSERT_EQ(st.b, st3.b);
            ASSERT_EQ(st.c, st3.c);
        }
#endif
#ifdef SEAL_USE_ZLIB
        {
            test_struct st3;
            out_size = Serialization::Save(
                bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::zlib), stream,
                compr_mode_type::zlib, false);
            in_size = Serialization::Load(bind(&test_struct::load_members, &st3, _1), stream, false);
            ASSERT_EQ(out_size, in_size);
            ASSERT_EQ(st.a, st3.a);
            ASSERT_EQ(st.b, st3.b);
            ASSERT_EQ(st.c, st3.c);
        }
#endif
    }

    TEST(SerializationTest, SaveLoadToBuffer)
    {
        test_struct st{ 3, ~0, 3.14159 }, st2;
        using namespace placeholders;

        constexpr size_t arr_size = 1024;
        seal_byte buffer[arr_size]{};

        stringstream ss;
        auto test_out_size = Serialization::Save(
            bind(&test_struct::save_members, &st, _1), st.save_size(Serialization::compr_mode_default), ss,
            Serialization::compr_mode_default, false);
        auto out_size = Serialization::Save(
            bind(&test_struct::save_members, &st, _1), st.save_size(Serialization::compr_mode_default), buffer,
            arr_size, Serialization::compr_mode_default, false);
        ASSERT_EQ(test_out_size, out_size);
        for (size_t i = static_cast<size_t>(out_size); i < arr_size; i++)
        {
            ASSERT_TRUE(seal_byte{} == buffer[i]);
        }

        auto in_size = Serialization::Load(bind(&test_struct::load_members, &st2, _1), buffer, arr_size, false);
        ASSERT_EQ(out_size, in_size);
        ASSERT_EQ(st.a, st2.a);
        ASSERT_EQ(st.b, st2.b);
        ASSERT_EQ(st.c, st2.c);
#ifdef SEAL_USE_ZSTD
        {
            // Reset buffer back to zero
            memset(buffer, 0, arr_size);

            test_struct st3;
            ss.seekp(0);
            test_out_size = Serialization::Save(
                bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::zstd), ss,
                compr_mode_type::zstd, false);
            out_size = Serialization::Save(
                bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::zstd), buffer, arr_size,
                compr_mode_type::zstd, false);
            ASSERT_EQ(test_out_size, out_size);
            for (size_t i = static_cast<size_t>(out_size); i < arr_size; i++)
            {
                ASSERT_EQ(seal_byte{}, buffer[i]);
            }

            in_size = Serialization::Load(bind(&test_struct::load_members, &st3, _1), buffer, arr_size, false);
            ASSERT_EQ(out_size, in_size);
            ASSERT_EQ(st.a, st3.a);
            ASSERT_EQ(st.b, st3.b);
            ASSERT_EQ(st.c, st3.c);
        }
#endif
#ifdef SEAL_USE_ZLIB
        {
            // Reset buffer back to zero
            memset(buffer, 0, arr_size);

            test_struct st3;
            ss.seekp(0);
            test_out_size = Serialization::Save(
                bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::zlib), ss,
                compr_mode_type::zlib, false);
            out_size = Serialization::Save(
                bind(&test_struct::save_members, &st, _1), st.save_size(compr_mode_type::zlib), buffer, arr_size,
                compr_mode_type::zlib, false);
            ASSERT_EQ(test_out_size, out_size);
            for (size_t i = static_cast<size_t>(out_size); i < arr_size; i++)
            {
                ASSERT_EQ(seal_byte{}, buffer[i]);
            }

            in_size = Serialization::Load(bind(&test_struct::load_members, &st3, _1), buffer, arr_size, false);
            ASSERT_EQ(out_size, in_size);
            ASSERT_EQ(st.a, st3.a);
            ASSERT_EQ(st.b, st3.b);
            ASSERT_EQ(st.c, st3.c);
        }
#endif
    }
} // namespace sealtest