/*
 * Copyright 2026 Rive
 */

#ifndef _RIVE_BITWISE_HPP_
#define _RIVE_BITWISE_HPP_

#include <type_traits>

#include "rive/rive_types.hpp"
#include "rive/enums.hpp"
#include "rive/math/math_types.hpp"

namespace rive::math
{

// Attempt to generate a "clz" assembly instruction.
RIVE_ALWAYS_INLINE static int clz32(uint32_t x)
{
  assert(x != 0);
#if __has_builtin(__builtin_clz)
  return __builtin_clz(x);
#else
  uint64_t doubleBits = bit_cast<uint64_t>(static_cast<double>(x));
  return 1054 - (doubleBits >> 52);
#endif
}

RIVE_ALWAYS_INLINE static int clz64(uint64_t x)
{
  assert(x != 0);
#if __has_builtin(__builtin_clzll)
  return __builtin_clzll(x);
#else
  uint32_t hi32 = x >> 32;
  return hi32 != 0 ? clz32(hi32) : 32 + clz32(x & 0xffffffff);
#endif
}

// Returns the 1-based index of the most significat bit in x.
//
//   0    -> 0
//   1    -> 1
//   2..3 -> 2
//   4..7 -> 3
//   ...
//
RIVE_ALWAYS_INLINE static uint32_t msb(uint32_t x)
{
  return x != 0 ? 32 - clz32(x) : 0;
}

// Attempt to generate a "rotl" (rotate-left) assembly instruction.
RIVE_ALWAYS_INLINE static uint32_t rotateleft32(uint32_t x, int y)
{
#if __has_builtin(__builtin_rotateleft32)
  return __builtin_rotateleft32(x, y);
#else
  return (x << y) | (x >> (32 - y));
#endif
}

namespace internal
{
// For compilers where there isn't a built-in popcount function, here's the
// simplest efficient implementation. This is done here in a namespace so that
// it can be unit tested even on systems that have built-in popcount support.
template <typename I> constexpr uint32_t count_set_bits_fallback(I i) noexcept
{
  uint32_t count = 0;
  while (i != 0)
  {
    count++;

    // remove the lowest set bit from the given value
    i &= i - 1;
  }
  return count;
}

} // namespace internal

template <typename T> constexpr uint32_t count_set_bits(T i) noexcept
{
#if __has_builtin(__builtin_popcount)
  static_assert(std::is_integral_v<T>);
  if constexpr (sizeof(T) == sizeof(uint64_t))
  {
    return uint32_t(__builtin_popcountll(uint64_t(i)));
  }
  else
  {
    return uint32_t(__builtin_popcount(uint32_t(i)));
  }
#else
  return internal::count_set_bits_fallback(i);
#endif
}

// Give a value with bits set within a given mask, compact the bits down to the
// minimum number needed to represent the values that are set in the mask.
//  For example: a mask of 11010001 has 4 set bits, and so a value of 10x0xxx1
//   (where x is any bit value not within the mask) would be compacted to 1001
constexpr uint32_t compact_bitmask_value(uint32_t value, uint32_t mask) noexcept
{
  uint32_t compacted = 0;
  for (int32_t inBitIndex = 31; inBitIndex >= 0; inBitIndex--)
  {
    auto bit = 1u << inBitIndex;
    if ((mask & bit) != 0)
    {
      compacted <<= 1;
      compacted |= ((value & bit) != 0) ? 1 : 0;
    }
  }

  return compacted;
}

// Undo the operation that compactBitmaskValue does.
//  For example: a mask of 11010001, with a given compacted vaue of 1001 would
//  expand to 10000001 (where each 1 in "compacted" means setting the
//  corresponding bit from mask)
constexpr uint32_t expand_compacted_bitmask_value(
    uint32_t compacted,
    uint32_t mask) noexcept
{
  uint32_t expanded = 0;
  for (int32_t maskBitIndex = 0; maskBitIndex < 32; maskBitIndex++)
  {
    auto maskBit = 1u << maskBitIndex;
    if ((mask & maskBit) != 0)
    {
      if (compacted & 1)
      {
        expanded |= maskBit;
      }
      compacted >>= 1;
    }
  }

  return expanded;
}

// Iteratable returned from iterateBitCombinationsInMask
template <typename T> class BitCombinationIterable
{
public:
  explicit constexpr BitCombinationIterable(T mask) : m_mask(mask) {}

  class Iterator
  {
  public:
    explicit constexpr Iterator(T mask) noexcept
        : m_currentValue(mask)
        , m_mask(mask)
    {}

    constexpr static Iterator MakeEnd(T mask) noexcept
    {
      Iterator it{mask};
      it.m_wasAdvanced = true;
      return it;
    }

    constexpr Iterator& operator++() noexcept
    {
      m_wasAdvanced = true;

      // Do this operation as an unsigned integral version of the type.
      // For enums this allows the subtraction below to work, and for
      // signed integers this allows us to do wrapping subtraction
      // (subtracting from 0) without it being undefined behavior.
      // NOTE: The odd "enable_if" in there is a trick to delay evaluation
      // of underlying_type<T> until the conditional gets chosen, because
      // otherwise this won't compile for non-enum types
      using U = std::make_unsigned_t<typename std::conditional_t<
          std::is_enum_v<T>,
          std::underlying_type<T>,
          std::enable_if<true, T>>::type>;

      // We actually iterate through these combinations from most bits to
      // least bits because it's simpler. (Iterating forward can be
      // accomplished by starting at 0, and xoring with the mask before
      // and after this decrement, but it seemed unnecessary since it's
      // unlikely that the ordering matters anywhere)

      // This subtraction and mask effectively removes the lowest set bit
      // in the mask from the current value (and re-sets and mask bits
      // below it), that is with a mask of 01001010 and a current value of
      // 01001000, subtracting would give 01000111, then the mask leaves
      // it as 01000010, which is the next-lowest valid combination of
      // mask bits from where we were.
      m_currentValue = T(U(m_currentValue) - 1) & m_mask;

      return *this;
    }

    constexpr Iterator operator++(int) noexcept
    {
      auto prev = *this;
      ++(*this);
      return prev;
    }

    constexpr T operator*() const noexcept { return m_currentValue; }

    constexpr bool operator==(const Iterator& other) const noexcept
    {
      assert(m_mask == other.m_mask);
      return (
          m_currentValue == other.m_currentValue &&
          m_wasAdvanced == other.m_wasAdvanced);
    }

    constexpr bool operator!=(const Iterator& other) const noexcept
    {
      return !(*this == other);
    }

  private:
    T m_currentValue{};
    T m_mask{};
    bool m_wasAdvanced = false;
  };

  constexpr Iterator begin() const noexcept { return Iterator{m_mask}; }
  constexpr Iterator end() const noexcept { return Iterator::MakeEnd(m_mask); }

private:
  T m_mask;
};

// Iterate through all combinations of set/unset bits for a given mask.
// That is, for a mask that is 101011, it would iterate through:
// 101011, 101010, 101001, 101000, 100011, ...,  000001, 000000
template <typename T>
BitCombinationIterable<T> iterate_bit_combinations_in_mask(T mask)
{
  return BitCombinationIterable{mask};
}

// Shift bits into the bottom of a key value (ensuring that the key value does
// not overflow and the value fits within the specified bit count)
template <typename Key, typename Bits>
[[nodiscard]] Key add_bits_to_key(
    Key key,
    Bits bits,
    uint32_t bitCount) noexcept
{
  static_assert(sizeof(Bits) <= sizeof(Key));

  assert((key << bitCount) >> bitCount == key);
  assert((Key(bits) & ((1ull << bitCount) - 1)) == Key(bits));
  return (key << bitCount) | Key(bits);
}

} // namespace rive::math
#endif