#include "parselink/msgpack/core/reader.h"
#include "rng.h"
#include <boost/ut.hpp>
#include <string>

namespace {
constexpr static std::size_t rng_samples_count = 10;

template <typename E>
constexpr auto enum_name() noexcept {
    return __PRETTY_FUNCTION__;
}

using namespace boost::ut;
namespace format = msgpack::format;

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>
struct oversized_array {
    std::array<T, C> data;
    std::size_t size;
};

constexpr auto to_bytes_array_oversized(auto const& container) {
    oversized_array<std::byte> 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());
    std::array<std::byte, oversized.size> out;
    std::copy(std::begin(oversized.data),
            std::next(std::begin(oversized.data), oversized.size),
            std::begin(out));
    return out;
}

template <std::size_t A, std::size_t B>
consteval auto cat(
        std::array<std::byte, A> const& a, std::array<std::byte, B> const& b) {
    std::array<std::byte, A + B> out;
    std::copy(std::begin(a), std::next(std::begin(a), std::size(a)),
            std::begin(out));
    std::copy(std::begin(b), std::next(std::begin(b), std::size(b)),
            std::next(std::begin(out), std::size(a)));
    return out;
}

template <typename T, typename Itr>
constexpr auto zip(T val, Itr begin, Itr end) {
    std::vector<T> output;
    for (auto itr = begin; itr != end; ++itr) {
        output.emplace_back(val);
        output.emplace_back(*itr);
    }
    return output;
}

template <typename T, typename Iterable>
constexpr auto zip(T val, Iterable const& itr) {
    return zip(val, std::begin(itr), std::end(itr));
}

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;
}

constexpr auto make_contiguous_range(std::uint8_t start, std::uint8_t end) {
    auto count = std::size_t(end) - std::size_t(start) + 1;
    std::vector<std::byte> range;
    range.resize(count);
    for (auto i = std::size_t(start); i <= std::size_t(end); ++i) {
        range[i - std::size_t(start)] = std::byte(i);
    }
    return range;
}
} // namespace

suite reader = [] {
    "empty span"_test = [] {
        std::span<std::byte> bytes;
        msgpack::reader reader(bytes);
        auto v = reader.read<msgpack::format::uint8>();
        expect(v.error() == msgpack::error::end_of_message);
    };

    ////////////////////////////////////////////////////////////////////////////////
    //  Unsigned integers
    ////////////////////////////////////////////////////////////////////////////////
    "reader::read<format::positive_fixint>"_test = [] {
        using fmt = format::positive_fixint;
        using error = msgpack::error;
        /**
         * All bytes 0x00->0x79 are effectively literal uint8s.
         */
        {
            constexpr auto payload = generate_bytes(
                    [] { return make_contiguous_range(0, 0x79); });
            msgpack::reader reader(payload);
            static_assert(std::is_same_v<
                    typename decltype(reader.read<fmt>())::value_type,
                    std::uint8_t>);
            for (auto byte : payload) {
                auto result = reader.read<fmt>();
                expect(result == std::uint8_t(byte));
            }

            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::end_of_message));
        }
        {
            constexpr auto payload = generate_bytes(
                    [] { return make_contiguous_range(0x80, 0xff); });
            // Continually narrow the span over the range, as read will not
            // advance the internal span on failure.
            for (auto itr = payload.begin(); itr != payload.end(); ++itr) {
                msgpack::reader reader(std::span(itr, std::end(payload)));
                auto result = reader.read<fmt>();
                expect(result == tl::make_unexpected(error::wrong_type));
            }
        }
    };

    "reader::read<format::uint8>"_test = [] {
        using fmt = format::uint8;
        using error = msgpack::error;
        /**
         * The uint8 format is 0xcc followed by one byte of data.
         */
        {
            constexpr auto payload = generate_bytes([] {
                return zip(std::byte{0xcc}, make_contiguous_range(0x00, 0xff));
            });
            msgpack::reader reader(payload);
            static_assert(std::is_same_v<
                    typename decltype(reader.read<fmt>())::value_type,
                    std::uint8_t>);
            for (auto i = 0; i < 0x100; ++i) {
                auto result = reader.read<fmt>();
                expect(result == std::uint8_t(i));
            }
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::end_of_message));
        }
        {
            using error = msgpack::error;
            // Test that partial read fails.
            constexpr auto payload = make_bytes(0xcc);
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
            // Retry, ensure the error remains the same.
            result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }
    };

    "reader::read<format::uint16>"_test = [] {
        using fmt = format::uint16;
        using error = msgpack::error;
        rng rng;
        {
            auto samples = rng.get<std::uint16_t, rng_samples_count>();
            std::vector<std::byte> payload;
            for (auto sample : samples) {
                auto bytes = ::detail::raw_cast(host_to_be(sample));
                payload.push_back(std::byte{0xcd});
                for (auto byte : bytes) {
                    payload.push_back(byte);
                }
            }

            msgpack::reader reader(payload);
            for (auto sample : samples) {
                auto result = reader.read<fmt>();
                expect(result == sample);
            }

            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::end_of_message));
        }

        {
            constexpr auto payload = make_bytes(0xee);
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::wrong_type));
        }

        // Test that partial read fails.
        {
            constexpr auto payload = make_bytes(0xcd);
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
            constexpr auto payload2 = make_bytes(0xcd, 0x01);
            reader = msgpack::reader(payload2);
            expect(result == tl::make_unexpected(error::incomplete_message));
        }
    };

    "reader::read<format::uint32>"_test = [] {
        using fmt = format::uint32;
        using error = msgpack::error;
        rng rng;
        auto samples = rng.get<std::uint32_t, rng_samples_count>();
        std::vector<std::byte> payload;
        for (auto sample : samples) {
            auto bytes = ::detail::raw_cast(host_to_be(sample));
            payload.push_back(std::byte{0xce});
            for (auto byte : bytes) {
                payload.push_back(byte);
            }
        }

        msgpack::reader reader(payload);
        for (auto sample : samples) {
            auto result = reader.read<fmt>();
            expect(bool(result));
            expect(result == sample);
        }
        // Test that partial read fails.
        {
            constexpr auto payload = make_bytes(0xce, 0x01, 0x02, 0x03, 0x09);
            for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
                msgpack::reader reader({payload.begin(), itr});
                auto result = reader.read<fmt>();
                expect(result
                        == tl::make_unexpected(error::incomplete_message));
            }
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == 0x01020309);
        }
    };

    "reader::read<format::uint64>"_test = [] {
        using fmt = format::uint64;
        using error = msgpack::error;
        rng rng;
        auto samples = rng.get<std::uint64_t, rng_samples_count>();
        std::vector<std::byte> payload;
        for (auto sample : samples) {
            auto bytes = ::detail::raw_cast(host_to_be(sample));
            payload.push_back(std::byte{0xcf});
            for (auto byte : bytes) {
                payload.push_back(byte);
            }
        }

        msgpack::reader reader(payload);
        for (auto sample : samples) {
            auto result = reader.read<fmt>();
            expect(bool(result));
            expect(result == sample);
        }

        {
            constexpr auto payload = make_bytes(
                    0xcf, 0x01, 0x02, 0x03, 0x09, 0x10, 0x20, 0x30, 0x90);
            for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
                msgpack::reader reader({payload.begin(), itr});
                auto result = reader.read<fmt>();
                expect(result
                        == tl::make_unexpected(error::incomplete_message));
            }
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == 0x0102030910203090);
        }
    };

    ////////////////////////////////////////////////////////////////////////////////
    //  Signed integers
    ////////////////////////////////////////////////////////////////////////////////
    "reader::read<format::negative_fixint>"_test = [] {
        using fmt = format::negative_fixint;
        using error = msgpack::error;
        /**
         * All bytes 0x00->0x79 are effectively literal uint8s.
         */
        {
            constexpr auto payload = generate_bytes(
                    [] { return make_contiguous_range(0xe0, 0xff); });
            msgpack::reader reader(payload);
            static_assert(std::is_same_v<
                    typename decltype(reader.read<fmt>())::value_type,
                    std::int8_t>);
            for (auto byte : payload) {
                auto result = reader.read<fmt>();
                expect(result == std::int8_t(byte));
            }

            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::end_of_message));
        }
        {
            constexpr auto payload = generate_bytes(
                    [] { return make_contiguous_range(0x00, 0xdf); });
            // Continually narrow the span over the range, as read will not
            // advance the internal span on failure.
            for (auto itr = payload.begin(); itr != payload.end(); ++itr) {
                msgpack::reader reader(std::span(itr, std::end(payload)));
                auto result = reader.read<fmt>();
                expect(result == tl::make_unexpected(error::wrong_type));
            }
        }
    };

    "reader::read<format::int8>"_test = [] {
        using fmt = format::int8;
        using error = msgpack::error;
        /**
         * The uint8 format is 0xcc followed by one byte of data.
         */
        {
            constexpr auto payload = generate_bytes([] {
                return zip(std::byte{0xd0}, make_contiguous_range(0x00, 0xff));
            });
            msgpack::reader reader(payload);
            static_assert(std::is_same_v<
                    typename decltype(reader.read<fmt>())::value_type,
                    std::int8_t>);
            for (auto i = 0; i < 0x100; ++i) {
                auto result = reader.read<fmt>();
                expect(*result == std::int8_t(i));
            }
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::end_of_message));
        }
        {
            using error = msgpack::error;
            // Test that partial read fails.
            constexpr auto payload = make_bytes(0xd0);
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
            // Retry, ensure the error remains the same.
            result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }
    };

    "reader::read<format::int16>"_test = [] {
        using fmt = format::int16;
        using error = msgpack::error;
        rng rng;
        {
            auto samples = rng.get<std::int16_t, rng_samples_count>();
            std::vector<std::byte> payload;
            for (auto sample : samples) {
                auto bytes = ::detail::raw_cast(host_to_be(sample));
                payload.push_back(std::byte{0xd1});
                for (auto byte : bytes) {
                    payload.push_back(byte);
                }
            }

            msgpack::reader reader(payload);
            for (auto sample : samples) {
                auto result = reader.read<fmt>();
                expect(result == sample);
            }

            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::end_of_message));
        }

        {
            constexpr auto payload = make_bytes(0xd2);
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::wrong_type));
        }

        // Test that partial read fails.
        {
            constexpr auto payload = make_bytes(0xd1);
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
            constexpr auto payload2 = make_bytes(0xd1, 0x01);
            reader = msgpack::reader(payload2);
            expect(result == tl::make_unexpected(error::incomplete_message));
        }
    };

    "reader::read<format::int32>"_test = [] {
        using fmt = format::int32;
        using error = msgpack::error;
        rng rng;
        auto samples = rng.get<std::int32_t, rng_samples_count>();
        std::vector<std::byte> payload;
        for (auto sample : samples) {
            auto bytes = ::detail::raw_cast(host_to_be(sample));
            payload.push_back(std::byte{0xd2});
            for (auto byte : bytes) {
                payload.push_back(byte);
            }
        }

        msgpack::reader reader(payload);
        for (auto sample : samples) {
            auto result = reader.read<fmt>();
            expect(bool(result));
            expect(result == sample);
        }
        // Test that partial read fails.
        {
            constexpr auto payload = make_bytes(0xd2, 0xff, 0xff, 0xff, 0xfe);
            for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
                msgpack::reader reader({payload.begin(), itr});
                auto result = reader.read<fmt>();
                expect(result
                        == tl::make_unexpected(error::incomplete_message));
            }
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == -2);
        }
    };

    "reader::read<format::int64>"_test = [] {
        using fmt = format::int64;
        using error = msgpack::error;
        rng rng;
        auto samples = rng.get<std::int64_t, rng_samples_count>();
        std::vector<std::byte> payload;
        for (auto sample : samples) {
            auto bytes = ::detail::raw_cast(host_to_be(sample));
            payload.push_back(std::byte{0xd3});
            for (auto byte : bytes) {
                payload.push_back(byte);
            }
        }

        msgpack::reader reader(payload);
        for (auto sample : samples) {
            auto result = reader.read<fmt>();
            expect(bool(result));
            expect(result == sample);
        }

        {
            constexpr auto payload = make_bytes(
                    0xd3, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xf0);
            for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
                msgpack::reader reader({payload.begin(), itr});
                auto result = reader.read<fmt>();
                expect(result
                        == tl::make_unexpected(error::incomplete_message));
            }
            msgpack::reader reader(payload);
            auto result = reader.read<fmt>();
            expect(result == -16);
        }
    };

    ////////////////////////////////////////////////////////////////////////////////
    //  Strings
    ////////////////////////////////////////////////////////////////////////////////

    "reader::read<fmt::fixstr>"_test = [] {
        using fmt = msgpack::format::fixstr;
        using error = msgpack::error;
        constexpr std::string_view sv = "hello";
        constexpr auto payload = cat(make_bytes(0xa0 + sv.size()),
                generate_bytes([sv] { return from_string_view(sv); }));

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == sv);
    };

    "reader::read<fmt::str8>"_test = [] {
        using fmt = msgpack::format::str8;
        using error = msgpack::error;
        constexpr std::string_view sv = "hello d";
        constexpr auto payload = cat(make_bytes(0xd9, sv.size()),
                generate_bytes([sv] { return from_string_view(sv); }));

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == sv);
    };

    "reader::read<fmt::str16>"_test = [] {
        using fmt = msgpack::format::str16;
        using error = msgpack::error;
        constexpr std::string_view sv = "hello world";
        constexpr auto payload = cat(make_bytes(0xda, 0x00, sv.size()),
                generate_bytes([sv] { return from_string_view(sv); }));

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == sv);
    };

    "reader::read<fmt::str32>"_test = [] {
        using fmt = msgpack::format::str32;
        using error = msgpack::error;
        constexpr std::string_view sv = "hello world";
        constexpr auto payload =
                cat(make_bytes(0xdb, 0x00, 0x00, 0x00, sv.size()),
                        generate_bytes([sv] { return from_string_view(sv); }));

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == sv);
    };

    ////////////////////////////////////////////////////////////////////////////////
    //  Binary payloads
    ////////////////////////////////////////////////////////////////////////////////

    "reader::read<fmt::bin8>"_test = [] {
        using fmt = msgpack::format::bin8;
        using error = msgpack::error;
        constexpr auto bv = make_bytes(0x0, 0x01, 0x02, 0x04);
        constexpr auto payload = cat(make_bytes(0xc4, bv.size()), bv);

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result.has_value());
        expect(std::equal(
                (*result).begin(), (*result).end(), bv.begin(), bv.end()));
    };

    "reader::read<fmt::bin16>"_test = [] {
        using fmt = msgpack::format::bin16;
        using error = msgpack::error;
        constexpr auto bv = generate_bytes([] {
            auto rg = make_contiguous_range(0, 0xff);
            std::vector<std::byte> rg2(rg.size() * 2);
            auto itr = std::copy(rg.begin(), rg.end(), rg2.begin());
            std::copy(rg.begin(), rg.end(), itr);
            return rg2;
        });
        constexpr auto payload =
                cat(make_bytes(0xc5, bv.size() >> 8, bv.size() & 0xff), bv);

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(std::equal(
                (*result).begin(), (*result).end(), bv.begin(), bv.end()));
    };

    // TODO: Support this with a proper test payload
    "reader::read<fmt::bin32>"_test = [] {
        using fmt = msgpack::format::bin32;
        using error = msgpack::error;
        constexpr auto bv = generate_bytes([] {
            auto rg = make_contiguous_range(0, 0xff);
            std::vector<std::byte> rg2(rg.size() * 2);
            auto itr = std::copy(rg.begin(), rg.end(), rg2.begin());
            std::copy(rg.begin(), rg.end(), itr);
            return rg2;
        });
        constexpr auto payload =
                cat(make_bytes(0xc6, bv.size() >> 24, bv.size() >> 16,
                            bv.size() >> 8, bv.size() & 0xff),
                        bv);

        for (auto itr = payload.begin() + 1; itr != payload.end(); ++itr) {
            msgpack::reader reader({payload.begin(), itr});
            auto result = reader.read<fmt>();
            expect(result == tl::make_unexpected(error::incomplete_message));
        }

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(std::equal(
                (*result).begin(), (*result).end(), bv.begin(), bv.end()));
    };

    ////////////////////////////////////////////////////////////////////////////////
    //  Misc formats
    ////////////////////////////////////////////////////////////////////////////////
    "reader::read<fmt::nil>"_test = [] {
        using fmt = msgpack::format::nil;
        constexpr auto payload = make_bytes(0xc0);

        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result.has_value());
    };

    "reader::read<fmt::boolean>"_test = [] {
        using fmt = msgpack::format::boolean;
        constexpr auto payload = make_bytes(0xc2, 0xc3);
        msgpack::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == false);
        result = reader.read<fmt>();
        expect(result == true);
    };

    ////////////////////////////////////////////////////////////////////////////////
    //  Structural formats
    ////////////////////////////////////////////////////////////////////////////////
    "reader::read<fmt::fixarray>"_test = [] {
        using fmt = msgpack::format::fixarray;
        using error = msgpack::error;
        // 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::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == msgpack::array_desc{5});

        for (std::size_t i = 0; i < (*result).count; ++i) {
            auto v = reader.read<format::uint8>();
            expect(v == 0x85 - i);
        }

        auto end = reader.read<format::uint8>();
        expect(end == tl::make_unexpected(error::end_of_message));
    };

    "reader::read<fmt::array16>"_test = [] {
        using fmt = msgpack::format::array16;
        using error = msgpack::error;
        // 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::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == msgpack::array_desc{5});

        for (std::size_t i = 0; i < (*result).count; ++i) {
            auto v = reader.read<format::uint8>();
            expect(v == 0x85 - i);
        }

        auto end = reader.read<format::uint8>();
        expect(end == tl::make_unexpected(error::end_of_message));
    };

    "reader::read<fmt::array32>"_test = [] {
        using fmt = msgpack::format::array32;
        using error = msgpack::error;
        // 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::reader reader(payload);
        auto result = reader.read<fmt>();
        expect(result == msgpack::array_desc{5});

        for (std::size_t i = 0; i < (*result).count; ++i) {
            auto v = reader.read<format::uint8>();
            expect(v == 0x85 - i);
        }

        auto end = reader.read<format::uint8>();
        expect(end == tl::make_unexpected(error::end_of_message));
    };
};