//-----------------------------------------------------------------------------
//     ___                      __ _       _
//    / _ \__ _ _ __ ___  ___  / /(_)_ __ | | __
//   / /_)/ _` | '__/ __|/ _ \/ / | | '_ \| |/ /
//  / ___/ (_| | |  \__ \  __/ /__| | | | |   <
//  \/    \__,_|_|  |___/\___\____/_|_| |_|_|\_\ .
//
//-----------------------------------------------------------------------------
// Author: Kurt Sassenrath
// Module: msgpack
//
// Default packer implementation, which aims to deduce the best format to use
// for a given value. For example, if a 32-bit unsigned integer type only
// contains the value 5, a uint32 format would serialize into:
//
//     0xce, 0x00, 0x00, 0x00, 0x05
//
// Instead, the packer will note that this value could be stored in a positive
// fixint, which is simply:
//
//     0x05
//
// The same optimization will be applied to variable-length types, like strings,
// bytes, arrays, and maps.
//
// This flexibility comes at the cost of CPU instructions. For embedded targets,
// writer (to be renamed verbatim_packer in the future) may be a better choice.
//
// Future goals for this particular packer:
//     1. Support containers/ranges seamlessly.
//     2. Support packing of trivial POD structures without an explicit
//        pack_adapter.
//
// Copyright (c) 2023 Kurt Sassenrath.
//
// License TBD.
//-----------------------------------------------------------------------------
#ifndef msgpack_core_packer_1d5939e9c1498568
#define msgpack_core_packer_1d5939e9c1498568

#include "parselink/msgpack/core/error.h"
#include "parselink/msgpack/core/format.h"
#include "parselink/msgpack/util/endianness.h"

#include <tl/expected.hpp>

#include <limits>
#include <type_traits>

namespace msgpack {

namespace detail {

// This is a generic helper function for writing integral bytes.
template <typename T, typename Itr>
constexpr auto write_integral(T value, Itr out, std::size_t sz) noexcept {
    auto bytes = ::detail::as_bytes(host_to_be(value));
    for (std::size_t i = 0; i < sz; ++i) {
        *out++ = *(bytes.end() - sz + i);
    }
    return out;
}

// Depending on the format, a number of bytes will be necessary to represent
// either the value (integer formats) or the length (variable length formats).
template <format::type T>
struct pack_helper {};

template <>
struct pack_helper<format::type::unsigned_int> {
    static constexpr std::size_t num_bytes(std::uint64_t value) noexcept {
#ifdef PARSELINK_CFG_MSGPACK_PACKER_USE_BRANCHES
        if (count <= 0x7f) return 0;
        if (count < 0x100) return 1;
        if (count < 0x10000) return 2;
        if (count < 0x100000000) return 4;
        return 8;
#else
        if (value <= std::uint64_t(format::positive_fixint::mask)) return 0;
        return std::bit_ceil(std::uint64_t((std::bit_width(value) + 7) >> 3));
#endif
    }

    static constexpr std::byte marker(std::uint64_t value) noexcept {
        switch (num_bytes(value)) {
            case 0: return static_cast<std::byte>(value);
            case 1: return format::uint8::marker;
            case 2: return format::uint16::marker;
            case 4: return format::uint32::marker;
            default: return format::uint64::marker;
        }
    }
};

template <>
struct pack_helper<format::type::signed_int> {
    static constexpr std::size_t num_bytes(std::int64_t value) noexcept {
        // Probably a better way to do this.
        if (value < 0 && value >= -32) return 0;
        auto underlying = static_cast<std::uint64_t>(value);

        // save a branch; these should be cheap on modern hardware.
        std::uint64_t counts[2] = {
                static_cast<std::uint64_t>(std::countl_zero(underlying)),
                static_cast<std::uint64_t>(std::countl_one(underlying))};

        std::uint64_t width = 1 + std::numeric_limits<std::uint64_t>::digits
                            - counts[underlying >> 63];

        return std::bit_ceil((width + 7) >> 3);
    }

    static constexpr std::byte marker(std::int64_t value) noexcept {
        switch (num_bytes(value)) {
            case 0: return static_cast<std::byte>(value);
            case 1: return format::int8::marker;
            case 2: return format::int16::marker;
            case 4: return format::int32::marker;
            default: return format::int64::marker;
        }
    }
};

template <>
struct pack_helper<format::type::binary> {
    static constexpr uint32_t max_supported_payload = 4;

    static constexpr std::size_t num_bytes(
            std::span<std::byte const> value) noexcept {
        return std::min(max_supported_payload,
                std::bit_ceil(std::uint32_t(
                        ((std::bit_width(value.size()) + 7) >> 3))));
    }

    static constexpr auto marker(std::span<std::byte const> value) noexcept {
        auto const len = num_bytes(value);
        switch (len) {
            case 1: return format::bin8::marker;
            case 2: return format::bin16::marker;
            case 4:
            default: return format::bin32::marker;
        }
    }
};

template <>
struct pack_helper<format::type::array> {
    static constexpr uint32_t max_supported_payload = 4;

    static constexpr std::size_t num_bytes(std::size_t count) noexcept {
        if (count < 16) return 0;
        return std::min(max_supported_payload,
                std::bit_ceil(std::uint32_t(
                        ((std::bit_width(count) + 15) >> 4) << 1)));
    }

    static constexpr auto marker(std::size_t value) noexcept {
        auto const len = num_bytes(value);
        switch (len) {
            case 0:
            case 2: return format::array16::marker;
            case 4:
            default: return format::bin32::marker;
        }
    }
};

constexpr auto x = pack_helper<format::type::array>::num_bytes(15);
constexpr auto y = pack_helper<format::type::array>::num_bytes(16);
constexpr auto z = pack_helper<format::type::array>::num_bytes(256);
constexpr auto a = pack_helper<format::type::array>::num_bytes(65536);

template <>
struct pack_helper<format::type::string> {
    static constexpr uint32_t max_supported_payload = 4;

    static constexpr std::size_t num_bytes(std::string_view value) noexcept {
        auto const len = value.size();
        if (len <= std::uint32_t(format::fixstr::mask)) return 0;
        return std::min(max_supported_payload,
                std::bit_ceil(std::uint32_t((std::bit_width(len) + 7) >> 3)));
    }

    static constexpr std::size_t payload_size(std::string_view value) noexcept {
        return value.size();
    }

    static constexpr auto marker(std::string_view value) noexcept {
        auto const len = num_bytes(value);
        switch (len) {
            case 0: return format::fixstr::marker | std::byte(value.size());
            case 1: return format::str8::marker;
            case 2: return format::str16::marker;
            case 4:
            default: return format::str32::marker;
        }
    }
};

} // namespace detail

// Pack adapter is the basis for packing native values into MessagePack format.
template <typename T>
struct pack_adapter {};

template <typename T>
concept basic_packable_type = requires(T const& t, std::byte* b) {
    { pack_adapter<T>::total_size(t) } -> std::same_as<std::size_t>;
    { pack_adapter<T>::write(t, b) } -> std::same_as<decltype(b)>;
    { pack_adapter<T>::marker(t) } -> std::same_as<std::byte>;
};

template <typename T>
concept packable_type = basic_packable_type<T>;

template <format::type F>
struct builtin_pack_adapter {
    static constexpr auto format_type = F;
    using pack_helper = detail::pack_helper<F>;

    // Retrieves the number of bytes to store the format, aside from the marker
    // byte.
    static constexpr auto format_size(auto value) noexcept {
        return pack_helper::num_bytes(value);
    }

    // Retrieves the number of bytes to store the payload, if it is
    // variable-length.
    static constexpr std::size_t payload_size(auto value) noexcept {
        if constexpr (requires { pack_helper::payload_size(value); }) {
            return pack_helper::payload_size(value);
        } else {
            return 0;
        }
    }

    static constexpr auto total_size(auto value) noexcept {
        return 1 + format_size(value) + payload_size(value);
    }

    static constexpr auto marker(auto value) noexcept {
        return detail::pack_helper<F>::marker(value);
    }
};

template <std::signed_integral T>
struct pack_adapter<T> : builtin_pack_adapter<format::type::signed_int> {
    template <typename Itr>
    static constexpr Itr write(T value, Itr out) noexcept {
        return detail::write_integral(value, out, format_size(value));
    }
};

template <std::unsigned_integral T>
struct pack_adapter<T> : builtin_pack_adapter<format::type::unsigned_int> {
    template <typename Itr>
    static constexpr Itr write(T value, Itr out) noexcept {
        return detail::write_integral(value, out, format_size(value));
    }
};

template <>
struct pack_adapter<std::string_view>
        : builtin_pack_adapter<format::type::string> {
    template <typename Itr>
    static constexpr Itr write(std::string_view value, Itr out) noexcept {
        out = detail::write_integral(value.size(), out, format_size(value));
        auto const* beg = reinterpret_cast<std::byte const*>(&*value.begin());
        std::copy(beg, beg + value.size(), out);
        return out + value.size();
    }
};

template <std::convertible_to<std::string_view> T>
struct pack_adapter<T> : pack_adapter<std::string_view> {};

template <>
struct pack_adapter<std::span<std::byte const>>
        : builtin_pack_adapter<format::type::binary> {
    template <typename Itr>
    static constexpr Itr write(
            std::span<std::byte const> value, Itr out) noexcept {
        out = detail::write_integral(value.size(), out, format_size(value));
        std::copy(value.begin(), value.end(), out);
        return out + value.size();
    }
};

class packer {
    using BufferType = std::span<std::byte>;

public:
    constexpr packer(BufferType buff)
            : buff_(buff)
            , curr_(std::begin(buff_))
            , end_(std::end(buff_)) {}

    template <packable_type T>
    constexpr tl::expected<tl::monostate, error> pack(T&& v) noexcept {
        using diff_type =
                std::iterator_traits<decltype(curr_)>::difference_type;
        diff_type const space_needed = pack_adapter<T>::total_size(v);
        if (space_needed > std::distance(curr_, end_)) {
            return tl::make_unexpected(error::out_of_space);
        }
        *curr_++ = pack_adapter<T>::marker(v);
        if (space_needed > 1) {
            curr_ = pack_adapter<T>::write(v, curr_);
        }
        return tl::monostate{};
    }

    constexpr auto tell() const noexcept {
        return std::distance(std::begin(buff_), curr_);
    }

    constexpr auto subspan() const noexcept { return buff_.subspan(0, tell()); }

private:
    BufferType buff_;
    decltype(std::begin(buff_)) curr_;
    decltype(std::end(buff_)) end_;
};

} // namespace msgpack

#endif // msgpack_core_packer_1d5939e9c1498568