#include <cstddef>
#include <cstdint>
#include <iostream>
#include <type_traits>
#include <vector>

#include "dpf.hpp"
#include "dpf/app_flow.hpp"
#include "dpf/app_plans.hpp"

// 3-party Duoram, the DPF steps only (Vadapalli, Henry, Goldberg, USENIX
// Security 2023). Online update: expand with `leaf_later`, rotate value and
// control together, then `apply_leaf_correction` once F is known.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty \
//       examples/applications/duoram3.cpp

namespace
{

constexpr std::size_t n = 256;
using output_t = simde_uint128;

template <typename Key>
output_t leaf_cw_of(const Key & key)
{
    using exterior = typename Key::exterior_node;
    return dpf::extract_leaf<exterior, output_t>(key.template leaf<0>(), 0);
}

}  // namespace

int main()
{
    constexpr std::uint8_t r = 10;
    constexpr std::uint8_t i_star = 42;
    constexpr unsigned shift = static_cast<unsigned>(i_star - r);
    const output_t message = output_t{7};

    std::vector<output_t> memory(n);
    memory[i_star] = output_t{100};
    memory[r] = output_t{5};

    auto [u0, u1] = dpf::make_dpf(r, output_t{1});
    const auto read = dpf::reconstruct(
        dpf::eval_full_inner_product(dpf::paired, u0, memory, dpf::rotate{shift}),
        dpf::eval_full_inner_product(dpf::paired, u1, memory, dpf::rotate{shift}));
    if (read != memory[i_star])
    {
        std::cerr << "duoram read\n";
        return 1;
    }

    auto [w0, w1] = dpf::make_dpf(r, message);
    const output_t F = leaf_cw_of(w0);
    std::vector<output_t> d0 = memory;
    std::vector<output_t> d1(n);
    std::vector<std::uint8_t> t0(n), t1(n);
    dpf::eval_full_add_into(d0, t0, w0, dpf::leaf_later{}, dpf::rotate{shift});
    dpf::eval_full_add_into(d1, t1, w1, dpf::leaf_later{}, dpf::rotate{shift});
    dpf::apply_leaf_correction(d0, t0, F);
    dpf::apply_leaf_correction(d1, t1, F);

    if ((d0[i_star] - d1[i_star]) != memory[i_star] + message
        || (d0[r] - d1[r]) != memory[r])
    {
        std::cerr << "duoram update\n";
        return 1;
    }

    {
        if (int rc = dpf::app::run_measured("duoram3",
                dpf::protocol::duoram_update_plan(0), 8))
            return rc;
    }

    std::cout << static_cast<unsigned long long>(d0[i_star] - d1[i_star]) << "\n";
    return 0;
}