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

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "pydpf_types.hpp"
#include "pydpf_eval.hpp"

// Header-only libdpf symbols are not `inline`; keep one TU.
#include "pydpf_eval.inc"

namespace py = pybind11;

PYBIND11_MODULE(pydpf, m)
{
    m.doc() = "libdpf bindings: point/interval/full/sequence/recipe, "
              "multileaf, wildcards, eval_until, it_dpf3";

    py::class_<pydpf_detail::PointKeyPair>(m, "PointKeyPair");
    py::class_<pydpf_detail::IdpfKeyPair>(m, "IdpfKeyPair");
    py::class_<pydpf_detail::ItDpf3Keys>(m, "ItDpf3Keys");
    py::class_<dpf::idpf_eval_ctx<pydpf_detail::idpf_key0>>(m, "IdpfEvalCtx0");
    py::class_<dpf::idpf_eval_ctx<pydpf_detail::idpf_key1>>(m, "IdpfEvalCtx1");

    m.def("make_dpf", [](std::uint8_t alpha, std::uint64_t beta) {
        auto [a, b] = dpf::make_dpf(alpha, beta);
        return pydpf_detail::PointKeyPair{std::move(a), std::move(b)};
    }, py::arg("alpha"), py::arg("beta"));

    m.def("eval_point", [](const pydpf_detail::PointKeyPair & keys, int party,
        std::uint8_t x) {
        pydpf_detail::require_party(party);
        if (party == 0)
            return pydpf_detail::share_limb(*dpf::eval_point(keys.k0, x));
        return pydpf_detail::share_limb(*dpf::eval_point(keys.k1, x));
    }, py::arg("keys"), py::arg("party"), py::arg("x"));

    m.def("reconstruct", [](std::uint64_t s0, std::uint64_t s1) {
        return s0 - s1;
    }, py::arg("share0"), py::arg("share1"));

    m.def("make_idpf16", [](std::uint16_t alpha) {
        auto [a, b] = dpf::make_dpf(alpha, dpf::idpf_ones<16>());
        return pydpf_detail::IdpfKeyPair{std::move(a), std::move(b)};
    }, py::arg("alpha"));

    m.def("make_eval_ctx", [](const pydpf_detail::IdpfKeyPair & keys, int party) {
        pydpf_detail::require_party(party);
        if (party == 0)
            return py::cast(dpf::idpf_eval_ctx<pydpf_detail::idpf_key0>(keys.k0));
        return py::cast(dpf::idpf_eval_ctx<pydpf_detail::idpf_key1>(keys.k1));
    }, py::arg("keys"), py::arg("party"));

    m.def("eval_until", [](py::object ctx, std::size_t level,
        const std::vector<std::uint16_t> & prefixes) {
        if (py::isinstance<dpf::idpf_eval_ctx<pydpf_detail::idpf_key0>>(ctx))
        {
            auto & c = ctx.cast<dpf::idpf_eval_ctx<pydpf_detail::idpf_key0> &>();
            auto shares = dpf::eval_until(c, level, prefixes);
            std::vector<std::uint64_t> out;
            out.reserve(shares.size());
            for (auto & s : shares)
                out.push_back(s.raw());
            return out;
        }
        if (py::isinstance<dpf::idpf_eval_ctx<pydpf_detail::idpf_key1>>(ctx))
        {
            auto & c = ctx.cast<dpf::idpf_eval_ctx<pydpf_detail::idpf_key1> &>();
            auto shares = dpf::eval_until(c, level, prefixes);
            std::vector<std::uint64_t> out;
            out.reserve(shares.size());
            for (auto & s : shares)
                out.push_back(s.raw());
            return out;
        }
        throw std::invalid_argument("eval_until: unknown context type");
    }, py::arg("ctx"), py::arg("level"), py::arg("prefixes"));

    m.def("ctx_node_count", [](py::object ctx) {
        if (py::isinstance<dpf::idpf_eval_ctx<pydpf_detail::idpf_key0>>(ctx))
            return ctx.cast<dpf::idpf_eval_ctx<pydpf_detail::idpf_key0> &>()
                .node_count();
        if (py::isinstance<dpf::idpf_eval_ctx<pydpf_detail::idpf_key1>>(ctx))
            return ctx.cast<dpf::idpf_eval_ctx<pydpf_detail::idpf_key1> &>()
                .node_count();
        throw std::invalid_argument("ctx_node_count: unknown context type");
    }, py::arg("ctx"));

    m.def("make_it_dpf3", [](std::uint8_t alpha, std::uint64_t beta) {
        auto [a, b, c] = dpf::make_it_dpf3(alpha, beta);
        return pydpf_detail::ItDpf3Keys{std::move(a), std::move(b), std::move(c)};
    }, py::arg("alpha"), py::arg("beta"));

    m.def("eval_it_dpf3", [](const pydpf_detail::ItDpf3Keys & keys, int party,
        std::uint8_t x) {
        if (party == 0)
            return dpf::eval_it_dpf3(keys.k0, x);
        if (party == 1)
            return dpf::eval_it_dpf3(keys.k1, x);
        if (party == 2)
            return dpf::eval_it_dpf3(keys.k2, x);
        throw std::invalid_argument("eval_it_dpf3: party must be 0, 1, or 2");
    }, py::arg("keys"), py::arg("party"), py::arg("x"));

    pydpf_detail::register_eval(m);
}