#pragma once

#include "assert.hpp"
#include "enum.hpp"
#include "tagged_union.hpp"
#include "type_traits.hpp"
#include "types.hpp"
#include <memory>
#include <optional>
#include <concepts>
#include <type_traits>
#include <vector>

namespace util {
namespace detail {
template <int N, typename... Ts>
using nth_type = typename std::tuple_element<N, std::tuple<Ts...>>::type;

template <typename> consteval std::size_t locate(std::size_t ind) {
  return static_cast<std::size_t>(-1);
}

template <typename IndexedType, typename T, typename... Ts>
consteval std::size_t locate(std::size_t ind = 0) {
  if (std::is_same<IndexedType, T>::value) {
    return ind;
  } else {
    return locate<IndexedType, Ts...>(ind + 1);
  }
}

template <typename T, typename... Ts>
struct index_of
    : std::integral_constant<std::size_t, locate<T, std::tuple<Ts...>>> {};

template <typename T, typename... Ts>
    inline constexpr std::size_t index_of_v = index_of<T, Ts...>::value;


template <typename IndexedType, typename...> struct is_in : std::false_type {};

template <typename IndexedType, typename Item, typename... Items>
struct is_in<IndexedType, Item, Items...>
    : std::integral_constant<bool, std::is_same<IndexedType, Item>::value ||
                                       is_in<IndexedType, Items...>::value> {};


template <typename Item, typename... Items>
concept in_types = is_in<std::size_t, Item, Items...>::value;


template<typename T>
inline constexpr auto max_all_inner(T arg) -> T {
    return arg;
}

template<typename T, typename... Ts>
inline constexpr auto max_all_inner(T arg, Ts... args) -> decltype(max_all_inner(arg)){
    return std::max(arg, max_all_inner(args...));
}

template<typename... Ts>
inline constexpr auto max_all(Ts... args) -> usize {
    return max_all_inner(args...);
}

template <typename... Ts> class variadic_union {
  std::vector<u8> m_data;
  std::optional<usize> m_type_index;

public:
  variadic_union() : m_data(detail::max_all(sizeof(Ts)...)), m_type_index(std::nullopt) {}

  template<std::size_t I, typename... Us>
  variadic_union(Us&&... args) : m_data(detail::max_all(sizeof(Ts)...)), m_type_index(I) {
    *reinterpret_cast<detail::nth_type<I, Ts...>*>(m_data.data()) = detail::nth_type<I, Ts...>{std::forward<Us>(args)...};
  }

  template<std::size_t I, typename U>
  requires (std::same_as<U, detail::nth_type<I, Ts...>>)
  variadic_union(const U& value) : m_data(detail::max_all(sizeof(Ts)...)), m_type_index(I) {
    *reinterpret_cast<detail::nth_type<I, Ts...>*>(m_data.data()) = value;
  }

  template<std::size_t I, typename U>
  requires (std::same_as<U, detail::nth_type<I, Ts...>>)
  variadic_union(U&& value) : m_data(detail::max_all(sizeof(Ts)...)), m_type_index(I) {
    *reinterpret_cast<detail::nth_type<I, Ts...>*>(m_data.data()) = std::move(value);
  }

  template<typename U>
  requires (in_types<U, Ts...>)
  variadic_union(U&& value) : m_data(detail::max_all(sizeof(Ts)...)), m_type_index(index_of_v<U, Ts...>) {
    *reinterpret_cast<U*>(m_data.data()) = std::move(value);
  }

  template<typename U>
  requires (in_types<U, Ts...>)
  variadic_union(const U& value) : m_data(detail::max_all(sizeof(Ts)...)), m_type_index(index_of_v<U, Ts...>) {
    *reinterpret_cast<U*>(m_data.data()) = value;
  }

template<typename U>
constexpr auto get() -> U& {
    util::assert_eq(m_type_index.value(), index_of_v<U, Ts...>);

    return reinterpret_cast<U*>(m_data.data());
}
};

} // namespace detail

#define NTH(Is, Ts) typename detail::nth_type<, Ts...> element_##I;

template <util::enum_type Tag, typename... Ts> class TaggedUnion {
    Tag tag;
    detail::variadic_union<Ts...> inner;
public:
template <typename U>
  requires (detail::in_types<U, Ts...>)
TaggedUnion(const Tag& tag, const U& value) : tag(tag), inner(detail::variadic_union(value)) {}

};
} // namespace util