#pragma once

#include <cstdint>
#include <iterator>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

#include "dpf.hpp"

namespace pydpf_detail
{

using point_key0 = decltype(dpf::make_dpf(std::uint8_t{0},
    std::uint64_t{0}).first);
using point_key1 = decltype(dpf::make_dpf(std::uint8_t{0},
    std::uint64_t{0}).second);

using multi_key0 = decltype(dpf::make_dpf(std::uint8_t{0},
    std::uint64_t{0}, std::uint64_t{0}).first);
using multi_key1 = decltype(dpf::make_dpf(std::uint8_t{0},
    std::uint64_t{0}, std::uint64_t{0}).second);

using wild_key0 = decltype(dpf::make_dpf(std::uint8_t{0},
    dpf::wildcard_value<std::uint64_t>{}).first);
using wild_key1 = decltype(dpf::make_dpf(std::uint8_t{0},
    dpf::wildcard_value<std::uint64_t>{}).second);

using idpf_key0 = decltype(dpf::make_dpf(std::uint16_t{0},
    dpf::idpf_ones<16>()).first);
using idpf_key1 = decltype(dpf::make_dpf(std::uint16_t{0},
    dpf::idpf_ones<16>()).second);

struct PointKeyPair
{
    point_key0 k0;
    point_key1 k1;
};

struct MultiKeyPair
{
    multi_key0 k0;
    multi_key1 k1;
};

struct WildcardKeyPair
{
    wild_key0 k0;
    wild_key1 k1;
};

struct IdpfKeyPair
{
    idpf_key0 k0;
    idpf_key1 k1;
};

struct ItDpf3Keys
{
    dpf::it_dpf3_key k0;
    dpf::it_dpf3_key k1;
    dpf::it_dpf3_key k2;
};

template <typename Share, typename = void>
struct has_raw_member : std::false_type {};
template <typename Share>
struct has_raw_member<Share,
    std::void_t<decltype(std::declval<const Share &>().raw())>>
    : std::true_type {};

template <typename Share>
std::uint64_t share_limb(const Share & s)
{
    if constexpr (has_raw_member<Share>::value)
        return static_cast<std::uint64_t>(s.raw());
    else
        return static_cast<std::uint64_t>(s);
}

template <typename Iterable>
std::vector<std::uint64_t> collect_limbs(Iterable && iterable)
{
    std::vector<std::uint64_t> out;
    for (auto it = std::begin(iterable); it != std::end(iterable); ++it)
        out.push_back(share_limb(*it));
    return out;
}

/// In-process wildcard leaf assign (same messages as asio, no socket).
template <std::size_t I = 0, typename DpfKey0, typename DpfKey1, typename ShareT>
void assign_leaf_local(DpfKey0 & dpf0, DpfKey1 & dpf1, const ShareT & shr0,
    const ShareT & shr1)
{
    auto & w0 = std::get<I>(dpf0.leaf_nodes);
    auto & w1 = std::get<I>(dpf1.leaf_nodes);
    if (w0.is_ready())
        w0.begin_update();
    if (w1.is_ready())
        w1.begin_update();
    const auto b0 = w0.compute_and_get_blinded_output_share(shr0);
    const auto b1 = w1.compute_and_get_blinded_output_share(shr1);
    const auto l0 = w0.compute_and_get_leaf_share(b1);
    const auto l1 = w1.compute_and_get_leaf_share(b0);
    w0.reconstruct_correction_word(l1);
    w1.reconstruct_correction_word(l0);
}

inline void require_party(int party)
{
    if (party != 0 && party != 1)
        throw std::invalid_argument("party must be 0 or 1");
}

inline void require_sorted(const std::vector<std::uint8_t> & xs)
{
    for (std::size_t i = 1; i < xs.size(); ++i)
    {
        if (xs[i] < xs[i - 1])
            throw std::invalid_argument("sequence points must be nondecreasing");
    }
}

}  // namespace pydpf_detail