/// @file dpf/json.hpp
/// @brief nlohmann::json serializers for DPF keys and beaver triples.
/// @details ADL `adl_serializer` specializations so `nlohmann::json` can
///          convert the library's key and triple types. Works with the
///          vendored nlohmann 3.12 headers and with a newer nlohmann already
///          included by the caller. Classic keys, multi-level `at<>` keys,
///          comparison channels (including payloads wider than 64 bits),
///          wildcard coefficients, and verifiable correction seeds round-trip.
/// @author Ryan Henry <ryan.henry@ucalgary.ca>
/// @copyright Copyright (c) 2019-2024 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_JSON_HPP__
#define LIBDPF_INCLUDE_DPF_JSON_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <tuple>
#include <array>
#include <type_traits>
#include <utility>

#if !defined(NLOHMANN_JSON_VERSION_MAJOR)
#include "json/include/nlohmann/json.hpp"
#endif

#include "portable-snippets/exact-int/exact-int.h"

#include "dpf/dpf_key.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{
namespace json
{
namespace codec
{

template <typename T>
struct is_std_array : std::false_type {};
template <typename T, std::size_t N>
struct is_std_array<std::array<T, N>> : std::true_type {};

template <typename T>
struct is_std_tuple : std::false_type {};
template <typename ...Ts>
struct is_std_tuple<std::tuple<Ts...>> : std::true_type {};

inline std::uint64_t as_u64(const nlohmann::json & j)
{
    if (j.is_number_unsigned())
        return j.get<std::uint64_t>();
    if (j.is_number_integer())
        return static_cast<std::uint64_t>(j.get<std::int64_t>());
    throw std::invalid_argument("dpf::json: expected an integer");
}

inline std::string hex_encode(const void * data, std::size_t n)
{
    static constexpr char digits[] = "0123456789abcdef";
    const auto * bytes = static_cast<const unsigned char *>(data);
    std::string out(n * 2, '\0');
    for (std::size_t i = 0; i < n; ++i)
    {
        out[2 * i] = digits[bytes[i] >> 4];
        out[2 * i + 1] = digits[bytes[i] & 0x0f];
    }
    return out;
}

inline int hex_nybble(char c)
{
    if (c >= '0' && c <= '9')
        return c - '0';
    if (c >= 'a' && c <= 'f')
        return c - 'a' + 10;
    if (c >= 'A' && c <= 'F')
        return c - 'A' + 10;
    return -1;
}

template <typename T>
nlohmann::json dump(const T & value);

template <typename T>
T load(const nlohmann::json & j);

template <typename Tuple, std::size_t ...Is>
nlohmann::json dump_tuple(const Tuple & value, std::index_sequence<Is...>)
{
    return nlohmann::json::array({dump(std::get<Is>(value))...});
}

template <typename Tuple, std::size_t ...Is>
Tuple load_tuple(const nlohmann::json & j, std::index_sequence<Is...>)
{
    return Tuple{load<std::tuple_element_t<Is, Tuple>>(j.at(Is))...};
}

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
template <typename T>
nlohmann::json dump(const T & value)
{
    using U = std::remove_cv_t<T>;
    if constexpr (std::is_same_v<U, simde__m128i>)
    {
        HEDLEY_PRAGMA(GCC diagnostic push)
        HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        std::uint64_t lane[2];
        std::memcpy(lane, &value, sizeof(lane));
        HEDLEY_PRAGMA(GCC diagnostic pop)
        nlohmann::json out = nlohmann::json::array();
        out.push_back(lane[0]);
        out.push_back(lane[1]);
        return out;
    }
    else if constexpr (std::is_same_v<U, simde__m256i>)
    {
        HEDLEY_PRAGMA(GCC diagnostic push)
        HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        std::uint64_t lane[4];
        std::memcpy(lane, &value, sizeof(lane));
        HEDLEY_PRAGMA(GCC diagnostic pop)
        nlohmann::json out = nlohmann::json::array();
        for (std::uint64_t limb : lane)
            out.push_back(limb);
        return out;
    }
    else if constexpr (std::is_same_v<U, simde_uint128>)
    {
        nlohmann::json out = nlohmann::json::array();
        out.push_back(static_cast<std::uint64_t>(value));
        out.push_back(static_cast<std::uint64_t>(value >> 64));
        return out;
    }
    else if constexpr (std::is_same_v<U, uint128_t>)
    {
        nlohmann::json out = nlohmann::json::array();
        out.push_back(value.lower());
        out.push_back(value.upper());
        return out;
    }
    else if constexpr (std::is_same_v<U, uint256_t>)
    {
        nlohmann::json out = nlohmann::json::array();
        out.push_back(value.lower().lower());
        out.push_back(value.lower().upper());
        out.push_back(value.upper().lower());
        out.push_back(value.upper().upper());
        return out;
    }
    else if constexpr (std::is_same_v<U, dpf::detail::cmp_meta>)
    {
        nlohmann::json out = nlohmann::json{
            {"nbits", value.nbits},
            {"mask", value.mask},
            {"kind", static_cast<psnip_uint8_t>(value.kind)},
            {"trivial", static_cast<psnip_uint8_t>(value.trivial)},
            {"eval_as_ge", value.eval_as_ge},
            {"include_eq", value.include_eq},
            {"active", value.active}
        };
        if (value.incremental)
            out["incremental"] = true;
        if (value.block_width != 0)
        {
            out["block_width"] = value.block_width;
            out["tail_bits"] = value.tail_bits;
        }
        return out;
    }
    else if constexpr (is_std_array<U>::value)
    {
        nlohmann::json out = nlohmann::json::array();
        for (const auto & elem : value)
            out.push_back(dump(elem));
        return out;
    }
    else if constexpr (is_std_tuple<U>::value)
    {
        return dump_tuple(value, std::make_index_sequence<std::tuple_size_v<U>>{});
    }
    else if constexpr (dpf::is_wildcard_v<U>)
    {
        return nlohmann::json(nullptr);
    }
    else if constexpr (std::is_enum_v<U>)
    {
        return dump(static_cast<std::underlying_type_t<U>>(value));
    }
    else if constexpr (std::is_same_v<U, bool>)
    {
        return nlohmann::json(value);
    }
    else if constexpr (std::is_integral_v<U> && sizeof(U) <= sizeof(std::uint64_t))
    {
        if constexpr (std::is_signed_v<U>)
            return nlohmann::json(static_cast<std::int64_t>(value));
        else
            return nlohmann::json(static_cast<std::uint64_t>(value));
    }
    else if constexpr (std::is_same_v<U, float> || std::is_same_v<U, double>)
    {
        return nlohmann::json(value);
    }
    else if constexpr (std::is_trivially_copyable_v<U>)
    {
        nlohmann::json out = nlohmann::json::object();
        out["$bytes"] = hex_encode(&value, sizeof(U));
        return out;
    }
    else
    {
        static_assert(sizeof(U) == 0, "dpf::json: no conversion for this type");
        return nlohmann::json(nullptr);
    }
}

template <typename T>
T load(const nlohmann::json & j)
{
    using U = std::remove_cv_t<T>;
    if constexpr (std::is_same_v<U, simde__m128i>)
    {
        std::uint64_t lane[2] = {as_u64(j.at(0)), as_u64(j.at(1))};
        HEDLEY_PRAGMA(GCC diagnostic push)
        HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        simde__m128i out;
        std::memcpy(&out, lane, sizeof(out));
        HEDLEY_PRAGMA(GCC diagnostic pop)
        return out;
    }
    else if constexpr (std::is_same_v<U, simde__m256i>)
    {
        std::uint64_t lane[4] = {
            as_u64(j.at(0)), as_u64(j.at(1)), as_u64(j.at(2)), as_u64(j.at(3))
        };
        HEDLEY_PRAGMA(GCC diagnostic push)
        HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        simde__m256i out;
        std::memcpy(&out, lane, sizeof(out));
        HEDLEY_PRAGMA(GCC diagnostic pop)
        return out;
    }
    else if constexpr (std::is_same_v<U, simde_uint128>)
    {
        if (j.is_number())
            return static_cast<simde_uint128>(as_u64(j));
        const simde_uint128 lo = as_u64(j.at(0));
        const simde_uint128 hi = as_u64(j.at(1));
        return lo | (hi << 64);
    }
    else if constexpr (std::is_same_v<U, uint128_t>)
    {
        if (j.is_number())
            return uint128_t{as_u64(j)};
        return uint128_t{as_u64(j.at(1)), as_u64(j.at(0))};
    }
    else if constexpr (std::is_same_v<U, uint256_t>)
    {
        if (j.is_number())
            return uint256_t{as_u64(j)};
        const uint128_t lo{as_u64(j.at(1)), as_u64(j.at(0))};
        const uint128_t hi{as_u64(j.at(3)), as_u64(j.at(2))};
        return uint256_t{hi, lo};
    }
    else if constexpr (std::is_same_v<U, dpf::detail::cmp_meta>)
    {
        dpf::detail::cmp_meta c;
        j.at("nbits").get_to(c.nbits);
        j.at("mask").get_to(c.mask);
        c.kind = static_cast<dpf::cmp_kind>(j.at("kind").get<psnip_uint8_t>());
        c.trivial =
            static_cast<dpf::cmp_trivial>(j.at("trivial").get<psnip_uint8_t>());
        j.at("eval_as_ge").get_to(c.eval_as_ge);
        j.at("include_eq").get_to(c.include_eq);
        j.at("active").get_to(c.active);
        c.incremental = j.value("incremental", false);
        c.block_width = j.value("block_width", 0);
        c.tail_bits = j.value("tail_bits", 0);
        return c;
    }
    else if constexpr (is_std_array<U>::value)
    {
        U out{};
        if (j.size() != out.size())
            throw std::invalid_argument("dpf::json: array length mismatch");
        for (std::size_t i = 0; i < out.size(); ++i)
            out[i] = load<typename U::value_type>(j.at(i));
        return out;
    }
    else if constexpr (is_std_tuple<U>::value)
    {
        if (j.size() != std::tuple_size_v<U>)
            throw std::invalid_argument("dpf::json: tuple length mismatch");
        return load_tuple<U>(j, std::make_index_sequence<std::tuple_size_v<U>>{});
    }
    else if constexpr (dpf::is_wildcard_v<U>)
    {
        return U{};
    }
    else if constexpr (std::is_enum_v<U>)
    {
        return static_cast<U>(load<std::underlying_type_t<U>>(j));
    }
    else if constexpr (std::is_same_v<U, bool>)
    {
        if (j.is_boolean())
            return j.get<bool>();
        return as_u64(j) != 0;
    }
    else if constexpr (std::is_integral_v<U> && sizeof(U) <= sizeof(std::uint64_t))
    {
        if constexpr (std::is_signed_v<U>)
            return static_cast<U>(static_cast<std::int64_t>(as_u64(j)));
        else
            return static_cast<U>(as_u64(j));
    }
    else if constexpr (std::is_same_v<U, float> || std::is_same_v<U, double>)
    {
        return j.get<U>();
    }
    else if constexpr (std::is_trivially_copyable_v<U>)
    {
        if (!j.is_object() || !j.contains("$bytes"))
            throw std::invalid_argument("dpf::json: expected a byte blob");
        const auto hex = j.at("$bytes").get<std::string>();
        if (hex.size() != sizeof(U) * 2)
            throw std::invalid_argument("dpf::json: byte blob has the wrong size");
        U out{};
        auto * bytes = reinterpret_cast<unsigned char *>(&out);
        for (std::size_t i = 0; i < sizeof(U); ++i)
        {
            const int hi = hex_nybble(hex[2 * i]);
            const int lo = hex_nybble(hex[2 * i + 1]);
            if (hi < 0 || lo < 0)
                throw std::invalid_argument("dpf::json: bad hex");
            bytes[i] = static_cast<unsigned char>((hi << 4) | lo);
        }
        return out;
    }
    else
    {
        static_assert(sizeof(U) == 0, "dpf::json: no conversion for this type");
        return U{};
    }
}
HEDLEY_PRAGMA(GCC diagnostic pop)

template <typename Word>
std::uint64_t low64(const Word & word)
{
    if constexpr (std::is_same_v<Word, uint256_t>)
        return static_cast<std::uint64_t>(word.lower().lower());
    else if constexpr (std::is_same_v<Word, uint128_t>)
        return word.lower();
    else if constexpr (sizeof(Word) <= sizeof(std::uint64_t))
        return static_cast<std::uint64_t>(word);
    else
        return static_cast<std::uint64_t>(word);
}

}  // namespace codec

template <typename DpfKey>
static std::string to_json(const DpfKey & dpf)
{
    nlohmann::json json = dpf;
    return json.dump();
}

template <typename DpfType>
static auto from_json(const std::string & json_string)
{
    nlohmann::json json = nlohmann::json::parse(json_string);
    return static_cast<DpfType>(json);
}

}  // namespace json
}  // namespace dpf

NLOHMANN_JSON_NAMESPACE_BEGIN

template <typename NodeT,
          typename OutputT>
struct adl_serializer<dpf::beaver<true, NodeT, OutputT>, void>
{
    static void from_json(const nlohmann::json & j, dpf::beaver<true, NodeT, OutputT> & beaver)  // NOLINT(runtime/references)
    {
        using beaver_type = dpf::beaver<true, NodeT, OutputT>;
        beaver.output_blind = dpf::json::codec::load<OutputT>(j.at("output_blind"));
        beaver.vector_blind = dpf::json::codec::load<typename beaver_type::LeafT>(j.at("vector_blind"));
        beaver.blinded_vector = dpf::json::codec::load<typename beaver_type::LeafT>(j.at("blinded_vector"));
        if (j.contains("assign_pad"))
            beaver.assign_pad = dpf::json::codec::load<typename beaver_type::LeafT>(j.at("assign_pad"));
    }

    static void to_json(nlohmann::json & j, const dpf::beaver<true, NodeT, OutputT> & beaver)  // NOLINT(runtime/references)
    {
        j = nlohmann::json{
            {"output_blind", dpf::json::codec::dump(beaver.output_blind)},
            {"vector_blind", dpf::json::codec::dump(beaver.vector_blind)},
            {"blinded_vector", dpf::json::codec::dump(beaver.blinded_vector)},
            {"assign_pad", dpf::json::codec::dump(beaver.assign_pad)}
        };
    }
};

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
template <>
struct adl_serializer<simde__m128i, void>
{
    static simde__m128i from_json(const nlohmann::json & j)
    {
        return dpf::json::codec::load<simde__m128i>(j);
    }

    static void to_json(nlohmann::json & j, const simde__m128i & a)  // NOLINT(runtime/references)
    {
        j = dpf::json::codec::dump(a);
    }
};

template <>
struct adl_serializer<simde__m256i, void>
{
    static simde__m256i from_json(const nlohmann::json & j)
    {
        return dpf::json::codec::load<simde__m256i>(j);
    }

    static void to_json(nlohmann::json & j, const simde__m256i & a)  // NOLINT(runtime/references)
    {
        j = dpf::json::codec::dump(a);
    }
};
HEDLEY_PRAGMA(GCC diagnostic pop)

template <>
struct adl_serializer<simde_uint128, void>
{
    static simde_uint128 from_json(const nlohmann::json & j)
    {
        return dpf::json::codec::load<simde_uint128>(j);
    }

    static void to_json(nlohmann::json & j, const simde_uint128 & a)  // NOLINT(runtime/references)
    {
        j = dpf::json::codec::dump(a);
    }
};

template <>
struct adl_serializer<uint128_t, void>
{
    static uint128_t from_json(const nlohmann::json & j)
    {
        return dpf::json::codec::load<uint128_t>(j);
    }

    static void to_json(nlohmann::json & j, const uint128_t & a)  // NOLINT(runtime/references)
    {
        j = dpf::json::codec::dump(a);
    }
};

template <>
struct adl_serializer<uint256_t, void>
{
    static uint256_t from_json(const nlohmann::json & j)
    {
        return dpf::json::codec::load<uint256_t>(j);
    }

    static void to_json(nlohmann::json & j, const uint256_t & a)  // NOLINT(runtime/references)
    {
        j = dpf::json::codec::dump(a);
    }
};

template <typename InteriorPRG,
          typename ExteriorPRG,
          typename InputT,
          typename OutputT,
          typename ...OutputTs>
struct adl_serializer<dpf::dpf_key<InteriorPRG, ExteriorPRG, InputT, OutputT, OutputTs...>, void>
{
    using dpf_type = dpf::dpf_key<InteriorPRG, ExteriorPRG, InputT, OutputT, OutputTs...>;
    using input_type = typename dpf_type::input_type;
    using leaf_tuple = typename dpf_type::leaf_tuple;
    using leaf_wrapper_tuple = typename dpf_type::leaf_wrapper_tuple;
    static constexpr bool classic =
        dpf::detail::incr::is_classic_pack_v<OutputT, OutputTs...>;

    template <std::size_t I>
    static nlohmann::json dump_leaf(const std::tuple_element_t<I, leaf_wrapper_tuple> & wrapper)
    {
        nlohmann::json entry = nlohmann::json::object();
        entry["leaf"] = dpf::json::codec::dump(wrapper.raw_leaf());
        if constexpr (dpf::is_wildcard_v<typename dpf_type::template output_type_t<I>>)
        {
            const auto & beaver = wrapper.beaver();
            entry["beaver"] = nlohmann::json{
                {"output_blind", dpf::json::codec::dump(beaver.output_blind)},
                {"vector_blind", dpf::json::codec::dump(beaver.vector_blind)},
                {"blinded_vector", dpf::json::codec::dump(beaver.blinded_vector)},
                {"assign_pad", dpf::json::codec::dump(beaver.assign_pad)}
            };
            entry["output_share"] = dpf::json::codec::dump(wrapper.output_share());
            entry["state"] = wrapper.state();
        }
        return entry;
    }

    template <std::size_t ...Is>
    static nlohmann::json dump_leaves(const leaf_wrapper_tuple & leaves,
        std::index_sequence<Is...>)
    {
        return nlohmann::json::array({dump_leaf<Is>(std::get<Is>(leaves))...});
    }

    template <std::size_t I>
    static auto load_beaver(const nlohmann::json & entry)
    {
        using beaver_type = std::tuple_element_t<I, typename dpf_type::beaver_tuple>;
        if constexpr (dpf::is_wildcard_v<typename dpf_type::template output_type_t<I>>)
        {
            beaver_type beaver{};
            const auto & stored = entry.at("beaver");
            beaver.output_blind = dpf::json::codec::load<decltype(beaver.output_blind)>(
                stored.at("output_blind"));
            beaver.vector_blind = dpf::json::codec::load<decltype(beaver.vector_blind)>(
                stored.at("vector_blind"));
            beaver.blinded_vector = dpf::json::codec::load<decltype(beaver.blinded_vector)>(
                stored.at("blinded_vector"));
            if (stored.contains("assign_pad"))
                beaver.assign_pad = dpf::json::codec::load<decltype(beaver.assign_pad)>(
                    stored.at("assign_pad"));
            return beaver;
        }
        else
        {
            return beaver_type{};
        }
    }

    template <std::size_t ...Is>
    static leaf_tuple load_leaf_tuple(const nlohmann::json & leaves,
        std::index_sequence<Is...>)
    {
        return leaf_tuple{
            dpf::json::codec::load<std::tuple_element_t<Is, leaf_tuple>>(
                leaves.at(Is).at("leaf"))...
        };
    }

    template <std::size_t ...Is>
    static auto load_beaver_tuple(const nlohmann::json & leaves,
        std::index_sequence<Is...>)
    {
        return typename dpf_type::beaver_tuple{load_beaver<Is>(leaves.at(Is))...};
    }

    template <std::size_t I, typename Wrapper>
    static void restore_leaf(Wrapper & wrapper, const nlohmann::json & entry)
    {
        if constexpr (dpf::is_wildcard_v<typename dpf_type::template output_type_t<I>>)
        {
            if (!entry.contains("state"))
                return;
            using output_type = typename Wrapper::output_type;
            output_type share{};
            if (entry.contains("output_share"))
                share = dpf::json::codec::load<output_type>(entry.at("output_share"));
            wrapper.restore_state(std::move(share),
                entry.at("state").template get<psnip_uint8_t>());
        }
        else
        {
            (void)wrapper;
            (void)entry;
        }
    }

    template <std::size_t ...Is>
    static void restore_leaves(dpf_type & key, const nlohmann::json & leaves,
        std::index_sequence<Is...>)
    {
        (restore_leaf<Is>(std::get<Is>(key.leaf_nodes), leaves.at(Is)), ...);
    }

    template <std::size_t I>
    static auto load_wrapper(const nlohmann::json & entry)
    {
        using wrapper = std::tuple_element_t<I, leaf_wrapper_tuple>;
        auto leaf = dpf::json::codec::load<typename wrapper::leaf_type>(entry.at("leaf"));
        if constexpr (dpf::is_wildcard_v<typename dpf_type::template output_type_t<I>>)
        {
            using beaver_type = typename wrapper::beaver_type;
            beaver_type beaver{};
            const auto & stored = entry.at("beaver");
            beaver.output_blind = dpf::json::codec::load<decltype(beaver.output_blind)>(
                stored.at("output_blind"));
            beaver.vector_blind = dpf::json::codec::load<decltype(beaver.vector_blind)>(
                stored.at("vector_blind"));
            beaver.blinded_vector = dpf::json::codec::load<decltype(beaver.blinded_vector)>(
                stored.at("blinded_vector"));
            if (stored.contains("assign_pad"))
                beaver.assign_pad = dpf::json::codec::load<decltype(beaver.assign_pad)>(
                    stored.at("assign_pad"));
            wrapper out{std::move(leaf), std::move(beaver)};
            restore_leaf<I>(out, entry);
            return out;
        }
        else
        {
            return wrapper{std::move(leaf)};
        }
    }

    template <std::size_t ...Is>
    static leaf_wrapper_tuple load_wrappers(const nlohmann::json & leaves,
        std::index_sequence<Is...>)
    {
        return leaf_wrapper_tuple{load_wrapper<Is>(leaves.at(Is))...};
    }

    static void restore_offset(dpf_type & key, const nlohmann::json & j)
    {
        if constexpr (dpf::is_wildcard_v<InputT>)
        {
            if (j.contains("offset_state"))
            {
                key.offset_x.restore(
                    dpf::json::codec::load<input_type>(j.at("offset")),
                    j.at("offset_state").template get<std::uint8_t>());
            }
        }
        else
        {
            (void)key;
            (void)j;
        }
    }

    static dpf_type from_json(const nlohmann::json & j)
    {
        const auto root = dpf::json::codec::load<typename dpf_type::interior_node>(j.at("root"));
        const auto correction_words =
            dpf::json::codec::load<typename dpf_type::correction_words_array>(j.at("correction_words"));
        const auto correction_advice =
            dpf::json::codec::load<typename dpf_type::correction_advice_array>(j.at("correction_advice"));
        input_type offset{};
        if (j.contains("offset"))
            offset = dpf::json::codec::load<input_type>(j.at("offset"));

        if constexpr (classic)
        {
            constexpr auto idx = std::make_index_sequence<dpf_type::num_outputs>{};
            const auto & leaves_json = j.at("leaves");
            dpf_type key{root, correction_words, correction_advice,
                load_leaf_tuple(leaves_json, idx),
                load_beaver_tuple(leaves_json, idx),
                offset};
            restore_leaves(key, leaves_json, idx);
            restore_offset(key, j);
            return key;
        }
        else
        {
            auto leaves = [&]() {
                if constexpr (dpf_type::num_outputs == 0)
                    return leaf_wrapper_tuple{};
                else
                    return load_wrappers(j.at("leaves"),
                        std::make_index_sequence<dpf_type::num_outputs>{});
            }();

            dpf::detail::cmp_meta cmp{};
            if (j.contains("cmp"))
                cmp = dpf::json::codec::load<dpf::detail::cmp_meta>(j.at("cmp"));

            typename dpf_type::value_cw_array value_cws{};
            if (j.contains("value_cw"))
                value_cws = dpf::json::codec::load<typename dpf_type::value_cw_array>(j.at("value_cw"));

            using word = typename dpf_type::value_cw_word;
            word cw_last{};
            if (j.contains("cw_last"))
                cw_last = dpf::json::codec::load<word>(j.at("cw_last"));
            word cmp_addend{};
            if (j.contains("cmp_addend"))
                cmp_addend = dpf::json::codec::load<word>(j.at("cmp_addend"));

            typename dpf_type::addend_tuple addends{};
            if constexpr (dpf_type::num_outputs > 0)
            {
                if (j.contains("addends"))
                    addends = dpf::json::codec::load<typename dpf_type::addend_tuple>(j.at("addends"));
            }

            typename dpf_type::value_cw_array value_cw_coeff{};
            if (j.contains("value_cw_coeff"))
                value_cw_coeff = dpf::json::codec::load<typename dpf_type::value_cw_array>(
                    j.at("value_cw_coeff"));
            word cw_last_coeff{};
            if (j.contains("cw_last_coeff"))
                cw_last_coeff = dpf::json::codec::load<word>(j.at("cw_last_coeff"));

            typename dpf_type::tail_array tail{};
            typename dpf_type::tail_array tail_coeff{};
            if constexpr (dpf_type::cmp_block > 0)
            {
                if (j.contains("tail_cw"))
                    tail = dpf::json::codec::load<typename dpf_type::tail_array>(j.at("tail_cw"));
                if (j.contains("tail_coeff"))
                    tail_coeff = dpf::json::codec::load<typename dpf_type::tail_array>(j.at("tail_coeff"));
            }

            typename dpf_type::prefix_cw_array prefix{};
            typename dpf_type::prefix_cw_array prefix_coeff{};
            if constexpr (dpf_type::cmp_idcf)
            {
                if (j.contains("prefix_cw"))
                    prefix = dpf::json::codec::load<typename dpf_type::prefix_cw_array>(j.at("prefix_cw"));
                if (j.contains("prefix_coeff"))
                    prefix_coeff = dpf::json::codec::load<typename dpf_type::prefix_cw_array>(
                        j.at("prefix_coeff"));
            }

            typename dpf_type::correction_seeds_array seeds{};
            if constexpr (dpf_type::is_verifiable)
            {
                if (j.contains("correction_seeds"))
                    seeds = dpf::json::codec::load<typename dpf_type::correction_seeds_array>(
                        j.at("correction_seeds"));
            }

            dpf_type key{root, correction_words, correction_advice,
                std::move(leaves), offset, cmp, value_cws,
                dpf::json::codec::low64(cw_last), dpf::json::codec::low64(cmp_addend),
                std::move(addends), value_cw_coeff, dpf::json::codec::low64(cw_last_coeff),
                tail, tail_coeff, prefix, prefix_coeff, seeds};
            key.set_cmp_scalars(cw_last, cmp_addend, cw_last_coeff);
            if (j.contains("cmp_assigned"))
                key.set_cmp_assigned(j.at("cmp_assigned").template get<bool>());
            restore_offset(key, j);
            return key;
        }
    }

    static void to_json(nlohmann::json & j, const dpf_type & dpf)  // NOLINT(runtime/references)
    {
        j = nlohmann::json::object();
        j["root"] = dpf::json::codec::dump(dpf.root());
        j["correction_words"] = dpf::json::codec::dump(dpf.correction_words());
        j["correction_advice"] = dpf::json::codec::dump(dpf.correction_advice());
        if constexpr (dpf_type::num_outputs > 0)
        {
            j["leaves"] = dump_leaves(dpf.leaf_nodes,
                std::make_index_sequence<dpf_type::num_outputs>{});
        }
        j["offset"] = dpf::json::codec::dump(dpf.offset_x.raw());
        if constexpr (dpf::is_wildcard_v<InputT>)
            j["offset_state"] = dpf.offset_x.state();

        if constexpr (!classic)
        {
            if constexpr (dpf_type::cmp_depth > 0)
            {
                j["cmp"] = dpf::json::codec::dump(dpf.cmp());
                j["value_cw"] = dpf::json::codec::dump(dpf.value_cw());
                j["cw_last"] = dpf::json::codec::dump(dpf.cw_last_word());
                j["cmp_addend"] = dpf::json::codec::dump(dpf.cmp_addend_word());
                if constexpr (dpf_type::cmp_block > 0)
                    j["tail_cw"] = dpf::json::codec::dump(dpf.tail_cw());
                if constexpr (dpf_type::cmp_idcf)
                    j["prefix_cw"] = dpf::json::codec::dump(dpf.prefix_cws());
                if constexpr (dpf_type::cmp_is_wildcard)
                {
                    j["value_cw_coeff"] = dpf::json::codec::dump(dpf.value_cw_coeff());
                    j["cw_last_coeff"] = dpf::json::codec::dump(dpf.cw_last_coeff_word());
                    if constexpr (dpf_type::cmp_block > 0)
                        j["tail_coeff"] = dpf::json::codec::dump(dpf.tail_coeff());
                    if constexpr (dpf_type::cmp_idcf)
                        j["prefix_coeff"] = dpf::json::codec::dump(dpf.prefix_cw_coeff());
                    j["cmp_assigned"] = dpf.cmp_assigned();
                }
            }
            if constexpr (dpf_type::num_outputs > 0)
                j["addends"] = dpf::json::codec::dump(dpf.public_addends);
            if constexpr (dpf_type::is_verifiable)
                j["correction_seeds"] = dpf::json::codec::dump(dpf.correction_seeds());
        }
    }
};

template <std::size_t Party, typename Key>
struct adl_serializer<dpf::party_key<Party, Key>, void>
{
    static dpf::party_key<Party, Key> from_json(const nlohmann::json & j)
    {
        return dpf::party_key<Party, Key>(j.template get<Key>());
    }

    static void to_json(nlohmann::json & j, const dpf::party_key<Party, Key> & key)  // NOLINT(runtime/references)
    {
        j = key.key();
    }
};

NLOHMANN_JSON_NAMESPACE_END

#endif  // LIBDPF_INCLUDE_DPF_JSON_HPP__