/// @file dpf/app_plans.hpp
/// @brief Named compose plans and star drivers for the application sketches.
#ifndef LIBDPF_INCLUDE_DPF_APP_PLANS_HPP__
#define LIBDPF_INCLUDE_DPF_APP_PLANS_HPP__

#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <stdexcept>
#include <utility>
#include <vector>

#include "dpf/compose.hpp"
#include "dpf/mesh_apps.hpp"
#include "dpf/net/edge_mesh.hpp"
#include "dpf/pad_graphs.hpp"
#include "dpf/protocol.hpp"
#include "dpf/session_host.hpp"

namespace dpf
{
namespace protocol
{

/// @brief Drive every session until all instances are done (interleaved).
/// @details Callers must `submit` each live index first (see `submit_and_drive`).
inline void drive_sessions(std::vector<schedule_session *> sessions,
    unsigned max_spins = 100000u)
{
    if (sessions.empty())
        return;
    for (auto * s : sessions)
    {
        if (s == nullptr)
            throw std::invalid_argument("drive_sessions null");
    }
    for (unsigned spins = 0;; ++spins)
    {
        if (spins > max_spins)
            throw std::runtime_error("drive_sessions: peer not ready");
        bool all_done = true;
        for (auto * s : sessions)
        {
            s->drive();
            for (std::size_t i = 0; i < s->count(); ++i)
                all_done = all_done && s->done(i);
        }
        if (all_done)
            return;
    }
}

/// @brief Submit every unfinished index, then `drive_sessions`.
inline void submit_and_drive(std::vector<schedule_session *> sessions,
    unsigned max_spins = 100000u)
{
    for (auto * s : sessions)
    {
        if (s == nullptr)
            throw std::invalid_argument("submit_and_drive null");
        for (std::size_t i = 0; i < s->count(); ++i)
            if (!s->done(i))
                s->submit(i);
    }
    drive_sessions(std::move(sessions), max_spins);
}

/// @brief Client↔N-server star: build sessions from round lists and drive.
inline void drive_star(net::memory_star & star,
    std::vector<schedule_round> client_rounds,
    const std::function<std::vector<schedule_round>(std::size_t server_i)> &
        server_rounds_for,
    unsigned max_spins = 100000u)
{
    schedule_session client(1, star.client_mesh(), std::move(client_rounds),
        false);
    std::vector<schedule_session> servers;
    servers.reserve(star.servers);
    for (std::size_t i = 0; i < star.servers; ++i)
        servers.emplace_back(1, star.server_edge(i), server_rounds_for(i),
            false);
    std::vector<schedule_session *> ptrs;
    ptrs.reserve(1 + servers.size());
    ptrs.push_back(&client);
    for (auto & s : servers)
        ptrs.push_back(&s);
    submit_and_drive(std::move(ptrs), max_spins);
}

/// @brief N-server PIR / keyword PIR compose plan (`client_servers` waves).
inline plan n_server_pir_plan(std::size_t party, std::size_t n_servers,
    std::size_t query_bytes, std::size_t answer_bytes)
{
    composer c(party);
    auto q = c.client_servers(n_servers, query_bytes, answer_bytes);
    (void)q;
    return c.default_plan();
}

inline plan keyword_pir_compose_plan(std::size_t party, std::size_t depth,
    std::size_t answer_bytes = sizeof(int))
{
    const std::size_t query_bytes = 16 + depth * 16;
    return n_server_pir_plan(party, 2, query_bytes, answer_bytes);
}

/// @brief Express / Sabre online walk: fused CW + trailer (sketch/proof).
inline plan mailbox_write_fused_plan(std::size_t party, std::size_t depth = 8,
    std::size_t cw_bytes = 16, std::size_t trailer_bytes = 8)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto trailer = c.input(domain::a, trailer_bytes);
    auto walk = c.fss_point_fused(seed, depth, cw_bytes, trailer);
    (void)walk;
    return c.default_plan();
}

/// @brief BitMore: L parallel bit-point walks (CSE → one wave per depth).
inline plan bitmore_fan_plan(std::size_t party, std::size_t L = 4,
    std::size_t depth = 6)
{
    composer c(party);
    auto seeds = c.fan(L, [&](std::size_t) {
        return c.input(domain::fss, 16);
    });
    (void)c.fan(L, [&](std::size_t i) {
        return c.fss_point(seeds[i], depth, 16);
    });
    return c.default_plan();
}

/// @brief BitMore / PIRsona star key-ship + answer schedule rounds.
inline std::vector<schedule_round> bitmore_star_fetch(std::size_t L,
    std::size_t seed_bytes, std::size_t answer_bytes,
    const std::shared_ptr<std::vector<std::vector<std::uint8_t>>> & seeds,
    const std::shared_ptr<std::vector<std::vector<std::uint8_t>>> & answers)
{
    return pirsona_bitmore_fetch(L, seed_bytes, answer_bytes, seeds, answers);
}

/// @brief SUBLEQ one instruction: prepaid expand + cmp branch skeleton.
inline plan subleq_instruction_plan(std::size_t party, std::size_t addr_depth = 8,
    std::size_t slot_bytes = 16)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto prepaid = c.defer_expand(seed, addr_depth, slot_bytes);
    auto branch = c.fss_cmp(seed, addr_depth, slot_bytes);
    (void)prepaid;
    (void)branch;
    return c.default_plan();
}

/// @brief Queue `n_insns` SUBLEQ instruction plans on a peer `session_host`.
inline void subleq_run_instructions(session_host & host, std::size_t n_insns,
    std::size_t addr_depth = 8, std::size_t slot_bytes = 16)
{
    for (std::size_t i = 0; i < n_insns; ++i)
        host.push(subleq_instruction_plan(host.party(), addr_depth, slot_bytes));
    host.drive_until_idle();
}

/// @brief Pika early-stop unit walk (`early_stop` bits pack into the leaf).
inline plan pika_lookup_plan(std::size_t party, std::size_t full_depth = 8,
    std::size_t early_stop = 3)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto tip = c.fss_point_early_stop(seed, full_depth, early_stop, 16);
    (void)tip;
    return c.default_plan();
}

/// @brief Pika with dealer pad delivery of the unit key before the walk.
inline plan pika_dealer_lookup_plan(std::size_t party, std::size_t full_depth = 8,
    std::size_t early_stop = 3)
{
    composer c(party);
    auto key = c.input(domain::fss, 16);
    auto delivered = c.dealer_deliver(key);
    auto tip = c.fss_point_early_stop(delivered, full_depth, early_stop, 16);
    (void)tip;
    return c.default_plan();
}

/// @brief Duoram write: path CWs with leaf applied later.
inline plan duoram_update_plan(std::size_t party, std::size_t depth = 8,
    std::size_t slot_bytes = 16)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto walk = c.leaf_later_walk(seed, depth, slot_bytes);
    auto leaf = c.input(domain::a, slot_bytes);
    auto applied = c.apply_leaf_correction(walk.values, walk.control, leaf);
    (void)applied;
    return c.default_plan();
}

/// @brief Poplar / Prio / Mastic: prefix checkpoints each level.
inline plan poplar_prefix_plan(std::size_t party, std::size_t depth = 8,
    std::size_t slot_bytes = 16, std::size_t prefix_bytes = 8)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto prefixes = c.level_walk_prefixes(seed, depth, slot_bytes, prefix_bytes);
    (void)prefixes;
    return c.default_plan();
}

/// @brief Weighted / multi-seed prefix fan (Prio multi-client shape).
inline plan poplar_prefix_fan_plan(std::size_t party, std::size_t n_keys,
    std::size_t depth = 8, std::size_t slot_bytes = 16,
    std::size_t prefix_bytes = 8)
{
    composer c(party);
    (void)c.fan(n_keys, [&](std::size_t) {
        auto seed = c.input(domain::fss, 16);
        return c.level_walk_prefixes(seed, depth, slot_bytes, prefix_bytes)
            .at_level.back();
    });
    return c.default_plan();
}

/// @brief Ledger (2,3): verifiable DPF3 append — upload + proof exchange.
inline plan ledger23_append_plan(std::size_t party, std::size_t depth = 8,
    std::size_t proof_bytes = 32)
{
    composer c(party);
    const std::size_t query_bytes = 16 + depth * 16;
    auto up = c.client_servers(3, query_bytes, proof_bytes);
    (void)up;
    return c.default_plan();
}

/// @brief Floram DS keygen walk (read/write expand stay local).
inline plan floram_ds_plan(std::size_t party, std::size_t depth = 8,
    std::size_t slot_bytes = 16, bool oh = false)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto tip = c.level_walk_ds(seed, depth, slot_bytes, oh);
    (void)tip;
    return c.default_plan();
}

/// @brief Splinter-style point walk (server expand cost).
inline plan fss_point_plan(std::size_t party, std::size_t depth = 8,
    std::size_t slot_bytes = 16)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto tip = c.fss_point(seed, depth, slot_bytes);
    (void)tip;
    return c.default_plan();
}

/// @brief Waldo-style comparison walk.
inline plan fss_cmp_plan(std::size_t party, std::size_t depth = 8,
    std::size_t slot_bytes = 16)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    auto tip = c.fss_cmp(seed, depth, slot_bytes);
    (void)tip;
    return c.default_plan();
}

/// @brief Range-count: two parallel cmp walks.
inline plan range_count_plan(std::size_t party, std::size_t depth = 8,
    std::size_t slot_bytes = 16)
{
    composer c(party);
    auto seed = c.input(domain::fss, 16);
    (void)c.fan(2, [&](std::size_t) {
        return c.fss_cmp(seed, depth, slot_bytes);
    });
    return c.default_plan();
}

/// @brief PSI cuckoo probes → multipoint fan (occupied bucket ids).
inline plan psi_cuckoo_plan(std::size_t party,
    const std::vector<std::size_t> & probes, std::size_t depth = 8,
    std::size_t slot_bytes = 16, std::size_t answer_bytes = 8)
{
    composer c(party);
    auto mr = c.schedule_cuckoo_probes(probes,
        [&](std::size_t) { return c.input(domain::fss, 16); }, depth,
        slot_bytes, answer_bytes);
    (void)mr;
    return c.default_plan();
}

/// @brief idpf_agg: adaptive prefix retain loop (`n_bits` online rounds).
inline plan idpf_agg_plan(std::size_t party, std::size_t n_bits = 16,
    std::size_t slot_bytes = 16, std::size_t prefix_bytes = 8)
{
    composer c(party);
    auto f = c.begin_adaptive_prefix(c.input(domain::fss, 16));
    for (std::size_t bit = 0; bit < n_bits; ++bit)
    {
        f = c.step_adaptive_prefix(f, slot_bytes, prefix_bytes);
        f = c.retain_adaptive_prefix(f, 0);
    }
    return c.default_plan();
}

}  // namespace protocol
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_APP_PLANS_HPP__