/*
 * 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 <concepts>
#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 (std::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.
        tag_t tag (void) const { return m_tag; }
        tag_t 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 {
        template <typename VisitorT, typename TaggedT, typename HeadT, typename ...TailT>
        requires (
            std::is_invocable_v<
                VisitorT,
                decltype (std::declval<TaggedT> ().template get <HeadT>())
            >
        )
        decltype (auto)
        visit (VisitorT &&visitor, TaggedT &&arg)
        {
            if constexpr (sizeof ...(TailT) == 0) {
                return std::invoke (std::forward<VisitorT> (visitor), arg.template get<HeadT> ());
            } else {
                if (arg.tag () == HeadT::tag) {
                    return std::invoke (std::forward<VisitorT> (visitor), arg.template get<HeadT> ());
                } else {
                    return visit<VisitorT, TaggedT, TailT...> (std::forward<VisitorT> (visitor), arg);
                }
            }
        }
    }


    /// Call the invokable `VisitorT` with the value in the supplied tagged
    /// union.
    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<VisitorT, TaggedT<ComponentsT...>&, ComponentsT...> (
            std::forward<VisitorT> (visitor),
            arg
        );
    }


    /// Call the invokable `VisitorT` with the value in the supplied tagged
    /// union.
    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<VisitorT, TaggedT<ComponentsT...> const&, ComponentsT...> (
            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> ())
            ) || ...
        );
    }
}