/// @file dpf/leaf_wrapper.hpp
/// @brief A concrete or wildcard leaf and the slot for its Beaver triple.
/// @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_LEAF_WRAPPER_HPP__
#define LIBDPF_INCLUDE_DPF_LEAF_WRAPPER_HPP__

#include <stdexcept>

#include "hedley/hedley.h"

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

namespace dpf
{

/// @brief Concrete leaf. `get()` is ready immediately.
/// @tparam OutputT output type
/// @tparam NodeT exterior node type
/// @see dpf::wildcard_value
template <typename OutputT,
          typename NodeT>
struct leaf_wrapper
{
  public:
    using leaf_type = dpf::leaf_node_t<NodeT, OutputT>;
    using output_type = OutputT;

    leaf_wrapper() = delete;
    leaf_wrapper(leaf_type leaf, dpf::beaver<false, NodeT, OutputT> = dpf::beaver<false, NodeT, OutputT>{})
        : leaf_{std::forward<leaf_type>(leaf)} { }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr const leaf_type & get() const noexcept { return leaf_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr leaf_type & get() noexcept { return leaf_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr const leaf_type & raw_leaf() const noexcept { return leaf_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr leaf_type & raw_leaf() noexcept { return leaf_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    const dpf::beaver<false, NodeT, OutputT> & beaver() const noexcept
    {
        static const dpf::beaver<false, NodeT, OutputT> dummy{};
        return dummy;
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    constexpr bool is_ready() const noexcept { return true; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    static constexpr bool is_wildcard() noexcept { return false; }

  private:
    leaf_type leaf_;
};

// // unpacked wildcard reconstruction
// template <typename ConcreteOutputT,
//           typename NodeT>
// struct leaf_wrapper<wildcard_value<ConcreteOutputT>, NodeT, false>
// {
//   public:
//     using leaf_type = dpf::leaf_node_t<NodeT, ConcreteOutputT>;
//     using output_type = ConcreteOutputT;

//     leaf_wrapper() = delete;
//     leaf_wrapper(leaf_type leaf_share, dpf::beaver<NodeT, output_type> = dpf::beaver<NodeT, output_type>{})
//         : leaf_{leaf_share},
//           leaf_state_(std::make_unique<std::atomic<leaf_status>>(leaf_status::notset)),
//           ready_{false} { }

//     HEDLEY_ALWAYS_INLINE
//     const leaf_type & get() const
//     {
//         if (HEDLEY_UNLIKELY(!ready_))
//         {
//             throw std::runtime_error("offset not set");
//         }
//         return leaf_;
//     }

//     const leaf_type compute_and_get_leaf_share(output_type output_share)
//     {
//         leaf_status notset = leaf_status::notset;
//         if (HEDLEY_UNLIKELY(!leaf_state_->compare_exchange_strong(notset,
//             leaf_status::computing,
//             std::memory_order_seq_cst, std::memory_order_relaxed)))
//         {
//             throw std::runtime_error("invalid state transition");
//         }
//         leaf_type tmp;
//         std::memcpy(&tmp, &output_share, sizeof(output_type));
//         leaf_ = add_leaf<output_type>(leaf_, tmp);
//         leaf_state_->store(leaf_status::waiting, std::memory_order_release);
//         return leaf_;
//     }

//     const leaf_type reconstruct_correction_word(leaf_type other_share)
//     {
//         leaf_status waiting = leaf_status::waiting;
//         if (HEDLEY_UNLIKELY(!leaf_state_->compare_exchange_strong(waiting,
//             leaf_status::computing,
//             std::memory_order_acquire, std::memory_order_relaxed)))
//         {
//             throw std::runtime_error("invalid state transition");
//         }
//         leaf_ = add_leaf<output_type>(leaf_, other_share);
//         ready_ = true;
//         leaf_state_->store(leaf_status::ready, std::memory_order_relaxed);
//         return leaf_;
//     }

//     HEDLEY_ALWAYS_INLINE
//     HEDLEY_NO_THROW
//     bool is_ready() const noexcept { return ready_; }

//     HEDLEY_ALWAYS_INLINE
//     HEDLEY_PURE
//     HEDLEY_NO_THROW
//     static constexpr bool is_wildcard() noexcept { return true; }

// //   private:
//     enum class leaf_status : psnip_uint8_t { ready = 0, waiting = 1, computing = 2, notset = 3 };
//     leaf_type leaf_;
//     std::unique_ptr<std::atomic<leaf_status>> leaf_state_;
//     bool ready_;
// };

template <typename ConcreteOutputT,
          typename NodeT>
struct leaf_wrapper<wildcard_value<ConcreteOutputT>, NodeT>
{
  public:
    using node_type = NodeT;
    using output_type = ConcreteOutputT;
    using leaf_type = dpf::leaf_node_t<node_type, output_type>;
    using beaver_type = dpf::beaver<true, node_type, output_type>;

    leaf_wrapper() = delete;
    leaf_wrapper(leaf_type leaf_share, beaver_type beaver)
        : leaf_{leaf_share},
          beaver_{beaver},
          output_share_{},
          leaf_state_{leaf_status::notset},
          updating_{false}
    { }

    HEDLEY_ALWAYS_INLINE
    const leaf_type & get() const
    {
        if (HEDLEY_UNLIKELY(leaf_state_ != leaf_status::ready))
        {
            throw std::runtime_error("offset not set");
        }
        return leaf_;
    }

    output_type compute_and_get_blinded_output_share(output_type output_share)
    {
        begin_transition(leaf_status::notset);
        output_share_ = output_share;
        // Use the leaf group (XOR of IEEE bits for float/double), not `operator+`.
        auto blinded_output_share = leaf_group_add(output_share_, beaver_.output_blind);
        leaf_state_ = leaf_status::blinded;
        return blinded_output_share;
    }

    /// @brief Accept a party-tagged share; convert to additive before Beaver math.
    /// @tparam Party party index, `0` or `1`
    /// @tparam Scheme scheme
    /// @param output_share the `output_share`
    /// @return the blinded output share
    template <std::size_t Party, sharing Scheme>
    output_type compute_and_get_blinded_output_share(
        const secret_share<output_type, Party, Scheme> & output_share)
    {
        // Beaver leaf math is a (2,2) additive absorb. (3,3) and replicated
        // shares are a different party count; fold them with `add_replicated`.
        if constexpr (is_two_party_sharing_v<Scheme>)
        {
            return compute_and_get_blinded_output_share(
                output_share.as_additive().raw());
        }
        else
        {
            static_assert(is_two_party_sharing_v<Scheme>,
                "wildcard Beaver absorb expects a (2,2) additive or "
                "subtractive share");
            return output_type{};
        }
    }

    leaf_type compute_and_get_leaf_share(output_type other_output_share)
    {
        begin_transition(leaf_status::blinded);
        leaf_ = add_leaf<output_type>(leaf_, subtract_leaf<output_type>(
            multiply_leaf(beaver_.blinded_vector, output_share_),
            multiply_leaf(beaver_.vector_blind, other_output_share)));
        leaf_state_ = leaf_status::waiting;
        return leaf_;
    }

    leaf_type reconstruct_correction_word(leaf_type other_share)
    {
        begin_transition(leaf_status::waiting);
        leaf_ = add_leaf<output_type>(leaf_, other_share);
        if (updating_)
        {
            // Opened naked delta; add it onto the previously committed payload.
            leaf_ = add_leaf<output_type>(committed_, leaf_);
            updating_ = false;
        }
        leaf_state_ = leaf_status::ready;
        return leaf_;
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    bool is_ready() const noexcept { return leaf_state_ == leaf_status::ready; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    static constexpr bool is_wildcard() noexcept { return true; }

    // Unassigned leaf / Beaver (before online payload). get() throws until ready.
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    const leaf_type & raw_leaf() const noexcept { return leaf_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    leaf_type & raw_leaf() noexcept { return leaf_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    const beaver_type & beaver() const noexcept { return beaver_; }

    /// @brief Output share captured during Beaver blinding, if any.
    /// @return the output share
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    const output_type & output_share() const noexcept { return output_share_; }

    /// @brief `leaf_status` as a byte, for serialization.
    /// @return the status byte
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    constexpr std::uint8_t state() const noexcept
    {
        return static_cast<std::uint8_t>(leaf_state_);
    }

    /// @brief Restore a serialized blinding share and status byte.
    /// @param share the output share
    /// @param state the status byte
    HEDLEY_ALWAYS_INLINE
    void restore_state(output_type share, std::uint8_t state) noexcept
    {
        output_share_ = std::move(share);
        leaf_state_ = static_cast<leaf_status>(state);
    }

    /// @brief Re-open a ready leaf so a later assign installs `β' − β`.
    /// @details Saves the committed correction word and restores the keygen
    ///          Beaver pad (not the zero-payload CW) so the next opening adds
    ///          only a naked delta. Reusing one scale Beaver for two openings
    ///          is not maliciously secure; honest parties that install `β'−β`
    ///          still get a correct leaf.
    HEDLEY_ALWAYS_INLINE
    void begin_update()
    {
        if (leaf_state_ != leaf_status::ready)
            throw std::runtime_error("begin_update: leaf is not ready");
        committed_ = leaf_;
        leaf_ = beaver_.assign_pad;
        updating_ = true;
        leaf_state_ = leaf_status::notset;
    }

  private:
    enum class leaf_status : psnip_uint8_t { ready = 0, waiting = 1, computing = 2, blinded = 3, notset = 4 };

    void begin_transition(leaf_status expected)
    {
        if (HEDLEY_UNLIKELY(leaf_state_ != expected))
        {
            throw std::runtime_error("invalid state transition");
        }
        leaf_state_ = leaf_status::computing;
    }

    leaf_type leaf_;
    leaf_type committed_{};
    beaver_type beaver_;
    output_type output_share_;
    leaf_status leaf_state_;
    bool updating_;
};

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_LEAF_WRAPPER_HPP__