#include <boost/ut.hpp>
#include <magic_enum.hpp>
#include <msgpack/token.h>

#include <fmt/format.h>

using namespace boost::ut;

namespace {
template <typename T>
constexpr bool operator==(std::span<T const> a, std::span<T const> b) noexcept {
    return std::equal(a.begin(), a.end(), b.begin(), b.end());
}

template <typename T>
constexpr std::array<std::byte, sizeof(T)> as_bytes(T&& t) {
    return std::bit_cast<std::array<std::byte, sizeof(T)>>(std::forward<T>(t));
}

template <typename... Bytes>
constexpr std::array<std::byte, sizeof...(Bytes)> make_bytes(Bytes&&... bytes) {
    return {std::byte(std::forward<Bytes>(bytes))...};
}

template <typename T, std::size_t C = 1024 * 1024>
struct oversized_array {
    std::array<T, C> data;
    std::size_t size;
};

constexpr auto to_bytes_array_oversized(auto const& container) {
    using value_type = std::decay_t<decltype(container[0])>;
    oversized_array<value_type> arr;
    std::copy(std::begin(container), std::end(container), std::begin(arr.data));
    arr.size = std::distance(std::begin(container), std::end(container));
    return arr;
}

consteval auto generate_bytes(auto callable) {
    constexpr auto oversized = to_bytes_array_oversized(callable());
    using value_type = std::decay_t<decltype(oversized.data[0])>;
    std::array<value_type, oversized.size> out;
    std::copy(std::begin(oversized.data),
            std::next(std::begin(oversized.data), oversized.size),
            std::begin(out));
    return out;
}

consteval auto build_string(auto callable) {
    constexpr auto string_array = generate_bytes(callable);
    return string_array;
}

template <std::size_t... Sizes>
constexpr auto cat(std::array<std::byte, Sizes> const&... a) noexcept {
    std::array<std::byte, (Sizes + ...)> out;
    std::size_t index{};
    ((std::copy_n(a.begin(), Sizes, out.begin() + index), index += Sizes), ...);
    return out;
}

constexpr auto repeat(std::span<std::byte const> sv, std::size_t count) {
    std::vector<std::byte> range;
    range.reserve(sv.size() * count);
    for (decltype(count) i = 0; i < count; ++i) {
        std::copy_n(sv.begin(), sv.size(), std::back_inserter(range));
    }
    return range;
}

constexpr auto repeat(std::string_view sv, std::size_t count) {
    std::vector<char> range;
    range.reserve(sv.size() * count);
    for (decltype(count) i = 0; i < count; ++i) {
        std::copy_n(sv.begin(), sv.size(), std::back_inserter(range));
    }
    return range;
}

constexpr auto from_string_view(std::string_view sv) {
    std::vector<std::byte> range;
    range.resize(sv.size());
    auto itr = range.begin();
    for (auto c : sv) {
        *itr = std::byte(c);
        ++itr;
    }
    return range;
}

template <typename T, typename E>
std::ostream& operator<<(std::ostream& os, tl::expected<T, E> const& exp) {
    if (exp.has_value()) {
        os << "Value: '" << *exp << "'";
    } else {
        os << "Error";
    }
    return os;
}

bool test_incomplete_message(auto const& payload) {
    // Test incomplete message.
    for (decltype(payload.size()) i = 1; i < payload.size() - 1; ++i) {
        // Test incomplete message.
        msgpack::token_reader reader(std::span(payload.data(), i));
        auto token = reader.read_one();
        if (token != tl::make_unexpected(msgpack::error::incomplete_message)) {
            fmt::print("Got the wrong response reading subview[0,{}] of "
                       "payload: {}\n",
                    i, token->get<std::string_view>().value());
            return false;
        }
    }
    return true;
}

bool test_end_of_message(auto& reader) {
    return reader.read_one()
        == tl::make_unexpected(msgpack::error::end_of_message);
}

} // namespace

suite reader_tests = [] {
    //--------------------------------------------------------------------------
    // Unsigned integer format tests
    //--------------------------------------------------------------------------
    "read format::positive_fixint"_test = [] {
        static constexpr auto payload = make_bytes(0x35);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::unsigned_int);
        expect(token->get<std::uint8_t>() == 0x35);
        expect(token->get<std::uint16_t>() == 0x35);
        expect(token->get<std::uint32_t>() == 0x35);
        expect(token->get<std::uint64_t>() == 0x35);
        expect(token == 0x35u);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::uint8"_test = [] {
        static constexpr auto payload = make_bytes(0xcc, 0x02);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::unsigned_int);
        expect(token->get<std::uint8_t>() == 0x02);
        expect(token->get<std::uint16_t>() == 0x02);
        expect(token->get<std::uint32_t>() == 0x02);
        expect(token->get<std::uint64_t>() == 0x02);
        expect(token == 0x02u);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::uint16"_test = [] {
        constexpr auto payload = make_bytes(0xcd, 0x01, 0x02);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::unsigned_int);
        expect(token->get<std::uint8_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::uint16_t>() == 0x0102);
        expect(token->get<std::uint32_t>() == 0x0102);
        expect(token->get<std::uint64_t>() == 0x0102);
        expect(token == 0x0102u);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::uint32"_test = [] {
        constexpr auto payload = make_bytes(0xce, 0x01, 0x02, 0x03, 0x04);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::unsigned_int);
        expect(token->get<std::uint8_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::uint16_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::uint32_t>() == 0x01020304);
        expect(token->get<std::uint64_t>() == 0x01020304);
        expect(token == 0x01020304u);
        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::uint64"_test = [] {
        constexpr auto payload = make_bytes(
                0xcf, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::unsigned_int);
        expect(token->get<std::uint8_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::uint16_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::uint32_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::uint64_t>() == 0x0102030405060708);
        expect(token == 0x0102030405060708u);
        // EOM
        expect(test_end_of_message(reader));
    };

    //--------------------------------------------------------------------------
    // Signed integer format tests
    //--------------------------------------------------------------------------
    "read format::negative_fixint"_test = [] {
        static constexpr auto payload = make_bytes(0xe5);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::signed_int);
        expect(token->get<std::int8_t>() == -27);
        expect(token->get<std::int16_t>() == -27);
        expect(token->get<std::int32_t>() == -27);
        expect(token->get<std::int64_t>() == -27);
        expect(token == -27);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::int8"_test = [] {
        static constexpr auto payload = make_bytes(0xd0, 0xff);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::signed_int);
        expect(token->get<std::int8_t>() == -1);
        expect(token->get<std::int16_t>() == -1);
        expect(token->get<std::int32_t>() == -1);
        expect(token->get<std::int64_t>() == -1);
        expect(token == -1);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::int16"_test = [] {
        constexpr auto payload = make_bytes(0xd1, 0xff, 0x00);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::signed_int);
        expect(token->get<std::int8_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::int16_t>() == -256);
        expect(token->get<std::int32_t>() == -256);
        expect(token->get<std::int64_t>() == -256);
        expect(token == -256);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::int32"_test = [] {
        constexpr auto payload = make_bytes(0xd2, 0xff, 0xff, 0x00, 0x00);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::signed_int);
        expect(token->get<std::int8_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::int16_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::int32_t>() == -65536);
        expect(token->get<std::int64_t>() == -65536);
        expect(token == -65536);
        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::int64"_test = [] {
        constexpr auto payload = make_bytes(
                0xd3, 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::signed_int);
        expect(token->get<std::int8_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::int16_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::int32_t>()
                == tl::make_unexpected(msgpack::error::will_truncate));
        expect(token->get<std::int64_t>() == -4294967296);
        expect(token == -4294967296);
        // EOM
        expect(test_end_of_message(reader));
    };

    //--------------------------------------------------------------------------
    // String format tests
    //--------------------------------------------------------------------------
    "read format::fixstr"_test = [] {
        constexpr std::string_view sv = "hello";
        constexpr auto payload = cat(make_bytes(0xa5),
                generate_bytes([sv] { return from_string_view(sv); }));

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::string);
        expect(token->get<std::string_view>() == sv);
        expect(token == sv);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::str8"_test = [] {
        constexpr std::string_view sv = "hello world";
        constexpr auto payload = cat(make_bytes(0xd9, sv.size()),
                generate_bytes([sv] { return from_string_view(sv); }));

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::string);
        expect(token->get<std::string_view>() == sv);
        expect(token == sv);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::str16"_test = [] {
        static constexpr auto sa =
                build_string([] { return repeat("0123456789abcdef", 17); });
        constexpr auto sv = std::string_view(sa.data(), sa.size());
        constexpr auto sv_size = host_to_be(std::uint16_t(sv.size()));

        auto payload = cat(make_bytes(0xda), as_bytes(sv_size),
                generate_bytes([] { return from_string_view(sv); }));

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::string);
        expect(token->get<std::string_view>() == sv);
        expect(token == sv);

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::str32"_test = [] {
        static constexpr auto sa =
                build_string([] { return repeat("0123456789abcdef", 4097); });
        constexpr auto sv = std::string_view(sa.data(), sa.size());
        constexpr auto sv_size = host_to_be(std::uint32_t(sv.size()));

        constexpr auto payload = cat(make_bytes(0xdb), as_bytes(sv_size),
                generate_bytes([sv] { return from_string_view(sv); }));

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::string);
        expect(token->get<std::string_view>() == sv);
        expect(token == sv);

        // EOM
        expect(test_end_of_message(reader));
    };

    //--------------------------------------------------------------------------
    // Binary format tests
    //--------------------------------------------------------------------------
    "read format::bin8"_test = [] {
        constexpr auto bytes = make_bytes(0x01, 0x02, 0x03, 0x04, 0x05, 0x06,
                0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10);
        constexpr auto bytes_size = host_to_be(std::uint8_t(bytes.size()));
        constexpr auto payload =
                cat(make_bytes(0xc4), as_bytes(bytes_size), bytes);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::binary);
        auto value = token->get<std::span<std::byte const>>().value();
        expect(value == std::span<std::byte const>(bytes));
        expect(token == std::span<std::byte const>(bytes));

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::bin16"_test = [] {
        static constexpr auto bytes = generate_bytes([] {
            return repeat(
                    make_bytes(0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08),
                    16);
        });
        constexpr auto bytes_size = host_to_be(std::uint16_t(bytes.size()));

        auto payload = cat(make_bytes(0xc5), as_bytes(bytes_size), bytes);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::binary);
        auto value = token->get<std::span<std::byte const>>().value();
        expect(value == std::span<std::byte const>(bytes));
        expect(token == std::span<std::byte const>(bytes));

        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::bin32"_test = [] {
        static constexpr auto bytes = generate_bytes([] {
            return repeat(
                    make_bytes(0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08),
                    8193);
        });
        constexpr auto bytes_size = host_to_be(std::uint32_t(bytes.size()));

        constexpr auto payload =
                cat(make_bytes(0xc6), as_bytes(bytes_size), bytes);

        // Test incomplete message.
        expect(test_incomplete_message(payload));

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();

        expect(token && token->type() == msgpack::format::type::binary);
        auto value = token->get<std::span<std::byte const>>().value();
        expect(value == std::span<std::byte const>(bytes));
        expect(token == std::span<std::byte const>(bytes));

        // EOM
        expect(test_end_of_message(reader));
    };
    //--------------------------------------------------------------------------
    // Structural format tests (map/array)
    //--------------------------------------------------------------------------
    "read format::fixarray"_test = [] {
        // A MessagePack array of 5 8-bit unsigned integers.
        constexpr auto payload = make_bytes(0x95, 0xcc, 0x85, 0xcc, 0x84, 0xcc,
                0x83, 0xcc, 0x82, 0xcc, 0x81);

        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::array);
        auto array = token->get<msgpack::array_desc>();
        expect(array && array->count == 5);

        for (std::size_t i = 0; i < array->count; ++i) {
            auto v = reader.read_one();
            expect(v && v->get<std::uint8_t>() == 0x85 - i);
        }

        expect(test_end_of_message(reader));
    };
    "read format::array16"_test = [] {
        // A MessagePack array of 5 8-bit unsigned integers.
        constexpr auto payload = make_bytes(0xdc, 0x00, 0x05, 0xcc, 0x85, 0xcc,
                0x84, 0xcc, 0x83, 0xcc, 0x82, 0xcc, 0x81);

        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::array);
        auto array = token->get<msgpack::array_desc>();
        expect(array && array->count == 5);

        for (std::size_t i = 0; i < array->count; ++i) {
            auto v = reader.read_one();
            expect(v && v == 0x85u - i);
        }

        expect(test_end_of_message(reader));
    };
    "read format::array32"_test = [] {
        // A MessagePack array of 5 8-bit unsigned integers.
        constexpr auto payload = make_bytes(0xdd, 0x00, 0x00, 0x00, 0x05, 0xcc,
                0x85, 0xcc, 0x84, 0xcc, 0x83, 0xcc, 0x82, 0xcc, 0x81);

        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::array);
        auto array = token->get<msgpack::array_desc>();
        expect(array && array->count == 5);

        for (std::size_t i = 0; i < array->count; ++i) {
            auto v = reader.read_one();
            expect(v && v->get<std::uint8_t>() == 0x85 - i);
        }

        expect(test_end_of_message(reader));
    };
    "read format::fixmap"_test = [] {
        // A MessagePack map of 3 strings to 8-bit unsigned integers.
        static constexpr std::array<std::string_view, 3> strings = {
                "one", "two", "three"};

        constexpr auto payload = cat(make_bytes(0x83, 0xa3),
                generate_bytes([] { return from_string_view(strings[0]); }),
                make_bytes(0x01, 0xa3),
                generate_bytes([] { return from_string_view(strings[1]); }),
                make_bytes(0x02, 0xa5),
                generate_bytes([] { return from_string_view(strings[2]); }),
                make_bytes(0x03));

        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::map);
        auto map = token->get<msgpack::map_desc>();
        expect(map && map->count == 3);

        for (std::size_t i = 0; i < map->count; ++i) {
            auto str = reader.read_one();
            expect(str && str->get<std::string_view>() == strings[i]);
            auto value = reader.read_one();
            expect(value && value->get<std::uint8_t>() == i + 1);
        }

        expect(test_end_of_message(reader));
    };
    "read format::map16"_test = [] {
        // A MessagePack map of 3 strings to 8-bit unsigned integers.
        static constexpr std::array<std::string_view, 3> strings = {
                "one", "two", "three"};

        constexpr auto payload = cat(make_bytes(0xde, 0x00, 0x03, 0xa3),
                generate_bytes([] { return from_string_view(strings[0]); }),
                make_bytes(0x01, 0xa3),
                generate_bytes([] { return from_string_view(strings[1]); }),
                make_bytes(0x02, 0xa5),
                generate_bytes([] { return from_string_view(strings[2]); }),
                make_bytes(0x03));

        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::map);
        auto map = token->get<msgpack::map_desc>();
        expect(map && map->count == 3);

        for (std::size_t i = 0; i < map->count; ++i) {
            auto str = reader.read_one();
            expect(str && str->get<std::string_view>() == strings[i]);
            expect(str == strings[i]);
            auto value = reader.read_one();
            expect(value && value->get<std::uint8_t>() == i + 1);
            expect(value == i + 1);
        }

        expect(test_end_of_message(reader));
    };
    "read format::map32"_test = [] {
        // A MessagePack map of 3 strings to 8-bit unsigned integers.
        static constexpr std::array<std::string_view, 3> strings = {
                "one", "two", "three"};

        constexpr auto payload = cat(
                make_bytes(0xdf, 0x00, 0x00, 0x00, 0x03, 0xa3),
                generate_bytes([] { return from_string_view(strings[0]); }),
                make_bytes(0x01, 0xa3),
                generate_bytes([] { return from_string_view(strings[1]); }),
                make_bytes(0x02, 0xa5),
                generate_bytes([] { return from_string_view(strings[2]); }),
                make_bytes(0x03));

        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::map);
        auto map = token->get<msgpack::map_desc>();
        expect(map && map->count == 3);

        for (std::size_t i = 0; i < map->count; ++i) {
            auto str = reader.read_one();
            expect(str && str->get<std::string_view>() == strings[i]);
            auto value = reader.read_one();
            expect(value && value->get<std::uint8_t>() == i + 1);
            expect(value == i + 1);
        }

        expect(test_end_of_message(reader));
    };
    //--------------------------------------------------------------------------
    // Assorted format tests
    //--------------------------------------------------------------------------
    "read format::boolean"_test = [] {
        constexpr auto payload = make_bytes(0xc2, 0xc3);

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::boolean);
        expect(token->get<bool>() == false);
        expect(token == false);
        token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::boolean);
        expect(token->get<bool>() == true);
        expect(token == true);
        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::nil"_test = [] {
        constexpr auto payload = make_bytes(0xc0);

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::nil);
        expect(token->get<msgpack::nil>().has_value());
        // EOM
        expect(test_end_of_message(reader));
    };
    "read format::invalid"_test = [] {
        constexpr auto payload = make_bytes(0xc1);

        // Test complete message.
        msgpack::token_reader reader(payload);
        auto token = reader.read_one();
        expect(token && token->type() == msgpack::format::type::invalid);
        expect(token->get<msgpack::invalid>().has_value());
        // EOM
        expect(test_end_of_message(reader));
    };

    //--------------------------------------------------------------------------
    // Assorted format tests
    //--------------------------------------------------------------------------
    "read multiple tokens"_test = [] {
        std::array<msgpack::token, 5> arr;
        constexpr auto payload = make_bytes(0xc2, 0xc3);
        msgpack::token_reader reader(payload);
        auto output = reader.read_some(arr);
        expect(output && output->size() == 2);
        expect((*output)[0] == false);
        expect((*output)[1] == true);
        output = reader.read_some(arr);
        expect(output == tl::make_unexpected(msgpack::error::end_of_message));
    };
};