/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 * Copyright 2018 Danny Robson <danny@nerdcruft.net>
 */

#pragma once

#include <variant>

#include "../cast.hpp"
#include "../maths.hpp"
#include "../tuple/type.hpp"

#include <tuple>


///////////////////////////////////////////////////////////////////////////////
namespace cruft {
    /// A tagged union of trivial types with a user-defined 'tag' type.
    ///
    /// This class supports far less functionality than does std::variant, but
    /// it simultaneously enforces  far less complexity and allows for a known
    /// tag type. If you need something more flexible:
    ///     don't extend this, just use std::variant.
    ///
    /// Storage is through a simple array of bytes (with suitable alignment).
    ///
    /// It should be impossible to construct an object in an uninitialised
    /// state.
    ///
    /// Requirements:
    ///   *  All member types must have a static constexpr 'tag' member of
    ///      identical type that uniquely maps to a member type.
    ///
    /// \tparam ValueT  A pack of possible member types to store.
    template <typename ...ValueT>
    class tagged {
    public:
        using tuple_t = std::tuple<ValueT...>;
        using first_t = cruft::tuple::type::nth_t<tuple_t,0>;
        using tag_t = std::decay_t<decltype (first_t::tag)>;

        /// All child types should have identical tag types
        static_assert ((std::is_same_v<tag_t, std::decay_t<decltype(ValueT::tag)>> && ...));

        /// All child types must be trivial, given we provide limited
        /// construction and destruction support.
        static_assert ((std::is_trivially_copyable_v<ValueT> && ...));
        static_assert ((std::is_trivially_destructible_v<ValueT> && ...));


        //---------------------------------------------------------------------
        template <typename InitialT>
        requires (cruft::concepts::same_as<std::remove_cvref_t<InitialT>, ValueT> || ...)
        tagged (InitialT &&initial)
        {
            set<InitialT> (std::forward<InitialT> (initial));
        }

        ~tagged () = default;

        tagged (tagged const&) = default;
        tagged (tagged &&) = default;

        tagged& operator= (tagged const&) = default;
        tagged& operator= (tagged &) = default;
        tagged& operator= (tagged &&) = default;


        //---------------------------------------------------------------------
        template <typename NextT>
        NextT&
        operator= (NextT &&next)
        {
            return set<NextT> (std::forward<NextT> (next));
        }


        ///--------------------------------------------------------------------
        /// Return the type code associated with the stored value.
        auto tag (void) const { return m_tag; }
        auto tag (void)       { return m_tag; }


        /// Returns true if the contained type is the same as the type
        /// specified in the template parameter.
        template <typename QueryT>
        bool is (void) const noexcept
        {
            return QueryT::tag == m_tag;
        }


        ///--------------------------------------------------------------------
        /// Return a reference to a stored value of known type.
        ///
        /// Requesting a reference to an incorrect type has undefined
        /// behaviour, but trigger should an assertion if the build has been
        /// configured with them enabled.
        template <typename InnerT>
        InnerT&
        get (void)&
        {
            CHECK (InnerT::tag == m_tag);
            return *cruft::cast::alignment<std::decay_t<InnerT>*> (m_data + 0);
        }


        template <typename InnerT>
        InnerT const&
        get (void) const&
        {
            CHECK (InnerT::tag == m_tag);
            return *cruft::cast::alignment<std::decay_t<InnerT> const*> (m_data + 0);
        }


        /// Set the inner member to a supplied value and store the associated
        /// type code.
        template <typename InnerT>
        InnerT&
        set (InnerT &&inner)&
        {
            m_tag = inner.tag;
            return *cruft::cast::alignment<std::decay_t<InnerT>*> (m_data + 0) = std::forward<InnerT> (inner);
        }

    private:
        /// The tag value for the currently active type
        tag_t m_tag;

        /// The storage area for all desired types.
        alignas (alignof (ValueT)...) std::byte m_data[
            cruft::max (sizeof (ValueT)...)
        ];
    };
}


namespace std {
    template <typename ...ComponentT>
    struct tuple_size<::cruft::tagged<ComponentT...>>
        : std::integral_constant<::std::size_t, sizeof...(ComponentT)>
    { ; };


    template <std::size_t I, typename ...ComponentT>
    struct tuple_element<
        I,
        ::cruft::tagged<ComponentT...>
    >
    : ::std::tuple_element<
        I,
        std::tuple<ComponentT...>
    >
    { ; };
}


namespace cruft {
    namespace detail {
        // If the tagged object matches the IndexV'th type in ComponentsT then
        // invoke visitor with the object as the argument. Else, advance to
        // the next index and try with that type.
        //
        // We assume that outside callers only call this function with an
        // IndexV of 0 (so that we have assurances that every type is visited).
        template <
            std::size_t IndexV,
            typename VisitorT,
            typename TaggedT
        >
        decltype (auto)
        visit (
            VisitorT &&visitor,
            TaggedT &&arg
        ) {
            static_assert (IndexV < std::tuple_size_v<std::remove_cvref_t<TaggedT>>);
            using nth_t = std::tuple_element_t<IndexV, std::remove_reference_t<TaggedT>>;

            // If we're the last valid index then we must be invoked because
            // there's no other option. Do this, and (statically) avoid
            // further recursion.
            if constexpr (IndexV + 1 == std::tuple_size_v<std::remove_reference_t<TaggedT>>) {
                return std::invoke (visitor, arg.template get<nth_t> ());
            } else {
                // If the tag matches, then dispatch, else recurse with the
                // next available type index.
                if (arg.tag () == nth_t::tag) {
                    return std::invoke (visitor, arg.template get<nth_t> ());
                } else {
                    return visit<IndexV+1> (
                        std::forward<VisitorT> (visitor),
                        arg
                    );
                }
            }
        }
    }


    template <
        typename VisitorT,
        template <typename...> typename TaggedT,
        typename ...ComponentsT
    >
    requires std::is_same_v<
        std::remove_cvref_t<TaggedT<ComponentsT...>>,
        tagged<ComponentsT...>
    >
    decltype (auto)
    visit (VisitorT &&visitor, TaggedT<ComponentsT...> &arg)
    {
        static_assert (sizeof...(ComponentsT));

        return detail::visit<0> (
            std::forward<VisitorT> (visitor),
            arg
        );
    }


    template <
        typename VisitorT,
        template <typename...> typename TaggedT,
        typename ...ComponentsT
    >
    requires std::is_same_v<
        std::remove_cvref_t<TaggedT<ComponentsT...>>,
        tagged<ComponentsT...>
    >
    decltype (auto)
    visit (VisitorT &&visitor, TaggedT<ComponentsT...> const &arg)
    {
        static_assert (sizeof...(ComponentsT));

        return detail::visit<0> (
            std::forward<VisitorT> (visitor),
            arg
        );
    }


    template <typename ...ComponentsT>
    bool operator== (tagged<ComponentsT...> const &lhs, tagged<ComponentsT...> const &rhs)
    {
        // Check the tags actually match to start with. This lets us use
        // short-circuiting in a second for the actual equality check.
        if (lhs.tag () != rhs.tag ())
            return false;

        // Construct a fold-expression that tests every component type for
        // equality. Use short-circuiting with `is` to protect against `get`
        // queries for the wrong types.
        return (
            (
                lhs.template is<ComponentsT> () &&
                (lhs.template get<ComponentsT> () == rhs.template get<ComponentsT> ())
            ) || ...
        );
    }

    template <typename ...ComponentsT>
    bool operator!= (tagged<ComponentsT...> const &lhs, tagged<ComponentsT...> const &rhs)
    {
        if (lhs.tag () != rhs.tag ())
            return false;

        return (
            (
                lhs.template is<ComponentsT> () &&
                (lhs.template get<ComponentsT> () != rhs.template get<ComponentsT> ())
            ) || ...
        );
    }
}