//-----------------------------------------------------------------------------
//     ___                      __ _       _
//    / _ \__ _ _ __ ___  ___  / /(_)_ __ | | __
//   / /_)/ _` | '__/ __|/ _ \/ / | | '_ \| |/ /
//  / ___/ (_| | |  \__ \  __/ /__| | | | |   <
//  \/    \__,_|_|  |___/\___\____/_|_| |_|_|\_\ .
//
//-----------------------------------------------------------------------------
// Author: Kurt Sassenrath
// Module: msgpack
//
// Core writing implementation, upon which writers are built.
//
// Copyright (c) 2023 Kurt Sassenrath.
//
// License TBD.
//-----------------------------------------------------------------------------
#ifndef msgpack_core_writer_ce48a51aa6ed0858
#define msgpack_core_writer_ce48a51aa6ed0858

#include <tl/expected.hpp>

#include "../util/endianness.h"
#include "error.h"
#include "format.h"
#include <limits>
#include <type_traits>

namespace msgpack {

enum class writer_error {
    none,
    unsupported,
    not_implemented,
    bad_value,
    out_of_space
};

// Helper template for writing a non-integral datatype to an output.
// template <typename T>
// struct write_adapter {
// template <typename Itr>
// static constexpr tl::expected<Itr, error> write(T const& t);
//};

template <std::size_t N, typename Itr>
constexpr inline decltype(auto) write_bytes(
        std::array<std::byte, N>&& data, Itr out) noexcept {
    for (auto b : data) {
        *out++ = b;
    }
    return out;
}

namespace detail {

constexpr std::size_t bytes_needed_for_str(std::uint32_t v) {
    if (v <= std::uint32_t(format::fixstr::mask)) return 0;
    return std::bit_ceil(std::uint32_t((std::bit_width(v) + 7) >> 3));
}

constexpr std::size_t bytes_needed(std::uint64_t v) {
    if (v <= std::uint64_t(format::positive_fixint::mask)) return 0;
    return std::bit_ceil(std::uint64_t((std::bit_width(v) + 7) >> 3));
}

constexpr std::size_t bytes_needed(std::int64_t v) {
    if (v < 0 && v >= -32) return 0;
    auto width = 1 + std::numeric_limits<std::uint64_t>::digits
               - (v < 0 ? std::countl_one(std::uint64_t(v))
                        : std::countl_zero(std::uint64_t(v)));
    return std::bit_ceil(std::uint64_t((width + 7) >> 3));
}

constexpr std::byte get_format_signed(auto value) {
    switch (bytes_needed(std::int64_t{value})) {
        case 0: return 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;
    }
}

constexpr std::byte get_format(auto value) {
    switch (bytes_needed(std::uint64_t{value})) {
        case 0: return 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;
    }
}

} // namespace detail

template <typename T>
struct write_adapter {};

template <std::integral T>
struct write_adapter<T> {
    static constexpr auto size(T t) noexcept { return sizeof(T); }

    template <typename Itr>
    static constexpr auto write(T t, Itr out) noexcept {
        return write_bytes(::detail::raw_cast(host_to_be(t)), out);
    }
};

template <typename T>
struct deduced_write_adapter {};

template <std::signed_integral T>
struct deduced_write_adapter<T> {
    static constexpr auto size(T value) noexcept {
        return detail::bytes_needed(std::int64_t{value});
    }

    template <typename Itr>
    static constexpr auto write(T value, Itr out) noexcept {
        auto bytes = ::detail::as_bytes(host_to_be(value));
        auto sz = size(value);
        for (std::size_t i = 0; i < sz; ++i) {
            *out++ = *(bytes.end() - sz + i);
        }
        return out;
    }

    static constexpr auto format_hint(T value) noexcept {
        return detail::get_format_signed(value);
    }
};

template <std::unsigned_integral T>
struct deduced_write_adapter<T> {
    static constexpr auto size(T value) noexcept {
        return detail::bytes_needed(std::uint64_t{value});
    }

    template <typename Itr>
    static constexpr auto write(T value, Itr out) noexcept {
        auto bytes = ::detail::as_bytes(host_to_be(value));
        auto sz = size(value);
        for (std::size_t i = 0; i < sz; ++i) {
            *out++ = *(bytes.end() - sz + i);
        }
        return out;
    }

    template <typename Itr>
    static constexpr auto write(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;
    }

    static constexpr auto format_hint(T value) noexcept {
        return detail::get_format(value);
    }
};

template <>
struct deduced_write_adapter<std::string_view> {
    static constexpr auto size(std::string_view value) noexcept {
        return value.size() + detail::bytes_needed_for_str(value.size());
    }

    template <typename Itr>
    static constexpr auto write(std::string_view value, Itr out) noexcept {
        auto bytes_needed = detail::bytes_needed_for_str(value.size());
        if (bytes_needed) {
            out = deduced_write_adapter<std::uint64_t>::write(
                    value.size(), out, bytes_needed);
        }
        auto const* beg = reinterpret_cast<std::byte const*>(&*value.begin());
        std::copy(beg, beg + value.size(), out);
        return out + value.size();
    }

    static constexpr auto format_hint(std::string_view value) noexcept {
        switch (detail::bytes_needed_for_str(value.size())) {
            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;
        }
    }
};

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

template <>
struct write_adapter<std::string_view> {
    static constexpr auto size(std::string_view str) noexcept {
        return str.size();
    }

    template <typename Itr>
    static constexpr auto write(std::string_view str, Itr out) noexcept {
        std::byte const* beg =
                reinterpret_cast<std::byte const*>(&*str.begin());
        std::copy(beg, beg + str.size(), out);
        return out + str.size();
    }
};

template <>
struct write_adapter<std::span<std::byte const>> {
    static constexpr auto size(std::span<std::byte const> bytes) noexcept {
        return bytes.size();
    }

    template <typename Itr>
    static constexpr auto write(
            std::span<std::byte const> bytes, Itr out) noexcept {
        std::copy(bytes.begin(), bytes.end(), out);
        return out += bytes.size();
    }
};

template <>
struct write_adapter<map_desc> {
    static constexpr auto value(map_desc desc) noexcept { return desc.count; }
};

template <>
struct write_adapter<array_desc> {
    static constexpr auto value(array_desc desc) noexcept { return desc.count; }
};

// TODO: These could be optimized away because invalid/nil never contain real
// data.
template <>
struct write_adapter<invalid> {
    static constexpr auto value(invalid) noexcept { return 0; }
};

template <>
struct write_adapter<nil> {
    static constexpr auto value(nil) noexcept { return 0; }
};

namespace detail {
template <typename T>
using expected = tl::expected<T, error>;

template <format_type F>
constexpr inline std::size_t calculate_space(
        typename F::value_type val) noexcept {
    // At a minimum, one byte is needed to store the format.
    std::size_t size = sizeof(typename F::first_type);

    if (!is_fixtype<F>) {
        ++size; // For format
    }

    if constexpr (F::payload_type == format::payload::variable) {
        size += write_adapter<typename F::value_type>::size(val);
    }

    return size;
}

// The "first type" is either the size of the variable length payload or
// a fixed-length value. Additionally, this "first type" may be inlined
// into the marker byte if it's a fix type.

template <format_type F>
constexpr inline expected<typename F::first_type> pack_first(
        typename F::value_type value) noexcept {
    using value_type = typename F::value_type;
    if constexpr (F::payload_type == format::payload::variable) {
        return typename F::first_type(write_adapter<value_type>::size(value));
    } else {
        if constexpr (requires { write_adapter<value_type>::value; }) {
            return typename F::first_type(
                    write_adapter<value_type>::value(value));
        } else {
            return typename F::first_type(value);
        }
    }
}

template <typename T>
concept v2_adapted = requires(T const& t, std::byte* b) {
    { deduced_write_adapter<T>::size(t) } -> std::same_as<std::size_t>;
    { deduced_write_adapter<T>::write(t, b) } -> std::same_as<decltype(b)>;
    { deduced_write_adapter<T>::format_hint(t) } -> std::same_as<std::byte>;
};

template <format_type F, typename Itr>
constexpr inline expected<Itr> write(
        typename F::value_type&& value, Itr out, Itr const end) {
    using diff_type = typename std::iterator_traits<Itr>::difference_type;
    if (diff_type(calculate_space<F>(value)) > std::distance(out, end)) {
        return tl::make_unexpected(error::out_of_space);
    }

    auto marker = F::marker;

    auto result = pack_first<F>(value);
    if (!result) {
        return tl::make_unexpected(result.error());
    }

    if constexpr (is_fixtype<F>) {
        if (*result > 0xff) {
            return tl::make_unexpected(error::bad_value);
        }
        if constexpr (F::flags & format::flag::apply_mask) {
            marker |= std::byte(*result);
        } else {
            marker = std::byte(*result);
        }
        if ((marker & ~F::mask) != F::marker) {
            return tl::make_unexpected(error::bad_value);
        }
    }

    *out++ = marker;

    if constexpr (!is_fixtype<F>) {
        out = write_adapter<typename decltype(result)::value_type>::write(
                *result, out);
    }

    if constexpr (F::payload_type == format::payload::variable) {
        out = write_adapter<typename F::value_type>::write(value, out);
    }

    return out;
}

template <typename T>
struct format_hint;

template <>
struct format_hint<std::uint8_t> {
    using type = format::positive_fixint;
};

template <>
struct format_hint<std::uint16_t> {
    using type = format::uint16;
};
} // namespace detail

class writer {
public:
    template <typename T>
    using expected = detail::expected<T>;

    constexpr writer(std::span<std::byte> dest)
            : data(dest)
            , curr(std::begin(data))
            , end(std::end(data)) {}

    template <format_type F>
    constexpr expected<tl::monostate> write(
            typename F::value_type&& v) noexcept {
        using value_type = typename F::value_type;
        auto result = detail::write<F>(std::forward<value_type>(v), curr, end);
        if (!result) {
            return tl::make_unexpected(result.error());
        }

        curr = *result;
        return tl::monostate{};
    }

    // Deduced-type write, automatically chooses smallest representative format
    template <detail::v2_adapted T>
    constexpr expected<tl::monostate> write2(T&& v) {
        using diff_type = std::iterator_traits<decltype(curr)>::difference_type;
        auto const space_needed =
                diff_type(1 + deduced_write_adapter<T>::size(v));
        if (space_needed > std::distance(curr, end)) {
            return tl::make_unexpected(error::out_of_space);
        }
        *curr++ = deduced_write_adapter<T>::format_hint(v);
        if (space_needed > 1) {
            curr = deduced_write_adapter<T>::write(v, curr);
        }
        return tl::monostate{};
    }

    constexpr auto pos() const noexcept { return curr; }

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

    constexpr auto subspan() const noexcept {
        return std::span<std::byte>(std::begin(data), tell());
    }

private:
    std::span<std::byte> data;
    decltype(data)::iterator curr;
    decltype(data)::iterator end;
};

} // namespace msgpack

#endif // msgpack_core_writer_ce48a51aa6ed0858