/// @file dpf/offset_wrapper.hpp
/// @brief An output value shifted by a public offset.
/// @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_OFFSET_WRAPPER_HPP__
#define LIBDPF_INCLUDE_DPF_OFFSET_WRAPPER_HPP__

#include <cstdint>
#include <utility>

#include "hedley/hedley.h"

namespace dpf
{

template <typename InputT>
struct offset_wrapper final
{
  public:
    using input_type = dpf::concrete_type_t<InputT>;

    offset_wrapper(input_type = input_type{})
        : offset_{} { }

    template <typename InputType>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    constexpr auto operator()(InputType && x) const noexcept
    {
        static_assert(std::is_convertible_v<InputType, input_type>);
        return input_type(std::forward<InputType>(x));
    }

    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; }

    /// @brief Stored share. Concrete wrappers ignore it at evaluation.
    /// @return the stored share
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    constexpr const input_type & raw() const noexcept { return offset_; }

  private:
    input_type offset_;  // waste an `input_type` to make `sizeof` match up
};

template <typename ConcreteInputT>
struct offset_wrapper<dpf::wildcard_value<ConcreteInputT>>
{
  public:
    using input_type = ConcreteInputT;

    offset_wrapper(input_type x)
        : offset_{x},
          offset_state_{offset_status::notset}
    { }

    template <typename InputType>
    HEDLEY_INLINE
    input_type operator()(InputType && x) const
    {
        static_assert(std::is_convertible_v<InputType, input_type>);
        if (HEDLEY_UNLIKELY(offset_state_ != offset_status::ready))
        {
            throw std::runtime_error("offset not set");
        }
        return input_type(x) + offset_;
    }

    template <typename InputType>
    const input_type & compute_and_get_share(InputType && input_share)
    {
        static_assert(std::is_convertible_v<InputType, input_type>);
        begin_transition(offset_status::notset);
        offset_ -= input_share;
        offset_state_ = offset_status::waiting;
        return offset_; 
    }

    template <typename InputType>
    const input_type & reconstruct(InputType && other_share)
    {
        static_assert(std::is_convertible_v<InputType, input_type>);
        begin_transition(offset_status::waiting);
        offset_ += other_share;
        offset_state_ = offset_status::ready;
        return offset_;
    }

    template <typename InputType>
    const input_type & set(InputType && offset)
    {
        static_assert(std::is_convertible_v<InputType, input_type>);
        begin_transition(offset_status::notset);
        offset_ += offset;
        offset_state_ = offset_status::ready;
        return offset_;
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    bool is_ready() const noexcept { return offset_state_ == offset_status::ready; }

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

    /// @brief Party share, including before the offset is marked ready.
    /// @return the party share
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    constexpr const input_type & raw() const noexcept { return offset_; }

    /// @brief `offset_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>(offset_state_);
    }

    /// @brief Restore a serialized share and status byte.
    /// @param offset the party share
    /// @param state the status byte
    HEDLEY_ALWAYS_INLINE
    void restore(input_type offset, std::uint8_t state) noexcept
    {
        offset_ = std::move(offset);
        offset_state_ = static_cast<offset_status>(state);
    }

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

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

    input_type offset_;
    offset_status offset_state_;
};

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_OFFSET_WRAPPER_HPP__