#pragma once

#include <new>

#include <assert.h>
#include <bit>
#include <span>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <utility>

#ifndef SHOW_MEMORY
#define SHOW_MEMORY 0
#endif

#ifndef DEBUG_VERBOSE
#define DEBUG_VERBOSE 0
#endif

// Use to toggle debug verbose dynamically
inline bool debugVerboseEnabled = true;

// This header contains code that we want to reuse outside of ConflictSet.cpp or
// want to exclude from coverage since it's only testing related.

// GCOVR_EXCL_START

#if SHOW_MEMORY
inline __attribute__((visibility("default"))) int64_t mallocBytes = 0;
inline __attribute__((visibility("default"))) int64_t peakMallocBytes = 0;
#endif

inline thread_local int64_t mallocBytesDelta = 0;

#ifndef NDEBUG
constexpr auto kMallocHeaderSize = 16;
#endif

// malloc that aborts on OOM and thus always returns a non-null pointer. Must be
// paired with `safe_free`.
__attribute__((always_inline)) inline void *safe_malloc(size_t s) {
  mallocBytesDelta += s;
#if SHOW_MEMORY
  mallocBytes += s;
  if (mallocBytes > peakMallocBytes) {
    peakMallocBytes = mallocBytes;
  }
#endif
  void *p = malloc(s
#ifndef NDEBUG
                   + kMallocHeaderSize
#endif
  );
  if (p == nullptr) {
    abort();
  }
#ifndef NDEBUG
  memcpy(p, &s, sizeof(s));
  (char *&)p += kMallocHeaderSize;
#endif
  return p;
}

// calloc that aborts on OOM and thus always returns a non-null pointer. Must be
// paired with `safe_free`.
__attribute__((always_inline)) inline void *safe_calloc(size_t count,
                                                        size_t size) {
  size_t s = count * size;
  mallocBytesDelta += s;
#if SHOW_MEMORY
  mallocBytes += s;
  if (mallocBytes > peakMallocBytes) {
    peakMallocBytes = mallocBytes;
  }
#endif
  void *p = calloc(s
#ifndef NDEBUG
                       + kMallocHeaderSize
#endif
                   ,
                   1);
  if (p == nullptr) {
    abort();
  }
#ifndef NDEBUG
  memcpy(p, &s, sizeof(s));
  (char *&)p += kMallocHeaderSize;
#endif
  return p;
}

// Must be paired with `safe_malloc` or `safe_calloc`.
//
// There's nothing safer about this than free. Only called safe_free for
// symmetry with safe_malloc.
__attribute__((always_inline)) inline void safe_free(void *p, size_t s) {
  mallocBytesDelta -= s;
#if SHOW_MEMORY
  mallocBytes -= s;
#endif
#ifndef NDEBUG
  (char *&)p -= kMallocHeaderSize;
  size_t expected;
  memcpy(&expected, p, sizeof(expected));
  assert(s == expected);
#endif
  free(p);
}

#if SHOW_MEMORY
inline __attribute__((visibility("default"))) bool showedMemory = false;
static struct __attribute__((visibility("default"))) PeakPrinter {
  ~PeakPrinter() {
    if (!std::exchange(showedMemory, true)) {
      printf("--- memory usage ---\n");
      printf("malloc bytes: %g\n", double(mallocBytes));
      printf("Peak malloc bytes: %g\n", double(peakMallocBytes));
    }
  }
} peakPrinter;
#endif

// ==================== BEGIN ARENA IMPL ====================

/// Group allocations with similar lifetimes to amortize the cost of malloc/free
struct Arena {
  explicit Arena(int initialSize = 0);
  /// O(log n) in the number of allocations
  ~Arena();
  struct ArenaImpl;
  Arena(const Arena &) = delete;
  Arena &operator=(const Arena &) = delete;
  Arena(Arena &&other) noexcept;
  Arena &operator=(Arena &&other) noexcept;

  ArenaImpl *impl = nullptr;
};

[[maybe_unused]] inline void operator delete(void *, std::align_val_t,
                                             Arena &) {}
inline void *operator new(size_t size, std::align_val_t align, Arena &arena);
void *operator new(size_t size, std::align_val_t align, Arena *arena) = delete;

[[maybe_unused]] inline void operator delete(void *, Arena &) {}
inline void *operator new(size_t size, Arena &arena) {
  return operator new(size, std::align_val_t(alignof(std::max_align_t)), arena);
}
inline void *operator new(size_t size, Arena *arena) = delete;

[[maybe_unused]] inline void operator delete[](void *, Arena &) {}
inline void *operator new[](size_t size, Arena &arena) {
  return operator new(size, arena);
}
inline void *operator new[](size_t size, Arena *arena) = delete;

[[maybe_unused]] inline void operator delete[](void *, std::align_val_t,
                                               Arena &) {}
inline void *operator new[](size_t size, std::align_val_t align, Arena &arena) {
  return operator new(size, align, arena);
}
inline void *operator new[](size_t size, std::align_val_t align,
                            Arena *arena) = delete;

/// align must be a power of two
template <class T> T *align_up(T *t, size_t align) {
  assert(std::popcount(align) == 1);
  auto unaligned = uintptr_t(t);
  auto aligned = (unaligned + align - 1) & ~(align - 1);
  return reinterpret_cast<T *>(reinterpret_cast<char *>(t) + aligned -
                               unaligned);
}

/// align must be a power of two
constexpr inline int align_up(uint32_t unaligned, uint32_t align) {
  assert(std::popcount(align) == 1);
  return (unaligned + align - 1) & ~(align - 1);
}

/// Returns the smallest power of two >= x
[[maybe_unused]] constexpr inline uint32_t nextPowerOfTwo(uint32_t x) {
  return x <= 1 ? 1 : 1 << (32 - std::countl_zero(x - 1));
}

struct Arena::ArenaImpl {
  Arena::ArenaImpl *prev;
  int capacity;
  int used;
  uint8_t *begin() { return reinterpret_cast<uint8_t *>(this + 1); }
};

inline Arena::Arena(int initialSize) : impl(nullptr) {
  if (initialSize > 0) {
    auto allocationSize =
        align_up(initialSize + sizeof(ArenaImpl), alignof(ArenaImpl));
    impl = (Arena::ArenaImpl *)safe_malloc(allocationSize);
    impl->prev = nullptr;
    impl->capacity = allocationSize - sizeof(ArenaImpl);
    impl->used = 0;
  }
}

inline void onDestroy(Arena::ArenaImpl *impl) {
  while (impl) {
    auto *prev = impl->prev;
    safe_free(impl, sizeof(Arena::ArenaImpl) + impl->capacity);
    impl = prev;
  }
}

[[maybe_unused]] inline Arena::Arena(Arena &&other) noexcept
    : impl(std::exchange(other.impl, nullptr)) {}
[[maybe_unused]] inline Arena &Arena::operator=(Arena &&other) noexcept {
  onDestroy(impl);
  impl = std::exchange(other.impl, nullptr);
  return *this;
}

inline Arena::~Arena() { onDestroy(impl); }

inline void *operator new(size_t size, std::align_val_t align, Arena &arena) {
  int64_t aligned_size = size + size_t(align) - 1;
  if (arena.impl == nullptr ||
      (arena.impl->capacity - arena.impl->used) < aligned_size) {
    auto allocationSize = align_up(
        sizeof(Arena::ArenaImpl) +
            std::max<int>(aligned_size,
                          (arena.impl ? std::max<int>(sizeof(Arena::ArenaImpl),
                                                      arena.impl->capacity * 2)
                                      : 0)),
        alignof(Arena::ArenaImpl));
    auto *impl = (Arena::ArenaImpl *)safe_malloc(allocationSize);
    impl->prev = arena.impl;
    impl->capacity = allocationSize - sizeof(Arena::ArenaImpl);
    impl->used = 0;
    arena.impl = impl;
  }
  auto *result =
      align_up(arena.impl->begin() + arena.impl->used, size_t(align));
  auto usedDelta = (result - arena.impl->begin()) + size - arena.impl->used;
  arena.impl->used += usedDelta;
  return result;
}

/// STL-friendly allocator using an arena
template <class T> struct ArenaAlloc {
  typedef T value_type;

  ArenaAlloc() = delete;
  explicit ArenaAlloc(Arena *arena) : arena(arena) {}

  Arena *arena;

  template <class U> constexpr ArenaAlloc(const ArenaAlloc<U> &other) noexcept {
    arena = other.arena;
  }

  [[nodiscard]] T *allocate(size_t n) {
    if (n > 0xfffffffffffffffful / sizeof(T)) { // NOLINT
      __builtin_unreachable();
    }

    return static_cast<T *>((void *)new (std::align_val_t(alignof(T)), *arena)
                                uint8_t[n * sizeof(T)]); // NOLINT
  }

  void deallocate(T *, size_t) noexcept {}
};

template <class T, class U>
bool operator==(const ArenaAlloc<T> &lhs, const ArenaAlloc<U> &rhs) {
  return lhs.arena == rhs.arena;
}

template <class T, class U>
bool operator!=(const ArenaAlloc<T> &lhs, const ArenaAlloc<U> &rhs) {
  return !(lhs == rhs);
}

// ==================== END ARENA IMPL ====================

// ==================== BEGIN RANDOM IMPL ====================

struct Random {
  // *Really* minimal PCG32 code / (c) 2014 M.E. O'Neill / pcg-random.org
  // Licensed under Apache License 2.0 (NO WARRANTY, etc. see website)
  //
  // Modified - mostly c -> c++
  Random() = default;

  Random(uint64_t initState, uint64_t initSeq) {
    pcg32_srandom_r(initState, initSeq);
    next();
  }

  /// Draws from a uniform distribution of uint32_t's
  uint32_t next() {
    auto result = next_;
    next_ = pcg32_random_r();
    return result;
  }

  /// Draws from a uniform distribution of [0, s). From
  /// https://arxiv.org/pdf/1805.10941.pdf
  uint32_t bounded(uint32_t s) {
    assert(s != 0);
    uint32_t x = next();
    auto m = uint64_t(x) * uint64_t(s);
    auto l = uint32_t(m);
    if (l < s) {
      uint32_t t = -s % s;
      while (l < t) {
        x = next();
        m = uint64_t(x) * uint64_t(s);
        l = uint32_t(m);
      }
    }
    uint32_t result = m >> 32;
    return result;
  }

  /// Fill `bytes` with `size` random hex bytes
  void randomHex(uint8_t *bytes, int size);

private:
  uint32_t pcg32_random_r() {
    uint64_t oldState = state;
    // Advance internal state
    state = oldState * 6364136223846793005ULL + inc;
    // Calculate output function (XSH RR), uses old state for max ILP
    uint32_t xorShifted = ((oldState >> 18u) ^ oldState) >> 27u;
    uint32_t rot = oldState >> 59u;
    return (xorShifted >> rot) | (xorShifted << ((-rot) & 31));
  }

  // Seed the rng.  Specified in two parts, state initializer and a
  // sequence selection constant (a.k.a. stream id)
  void pcg32_srandom_r(uint64_t initstate, uint64_t initSeq) {
    state = 0U;
    inc = (initSeq << 1u) | 1u;
    pcg32_random_r();
    state += initstate;
    pcg32_random_r();
  }
  uint32_t next_{};
  // RNG state.  All values are possible.
  uint64_t state{};
  // Controls which RNG sequence (stream) is selected. Must *always* be odd.
  uint64_t inc{};
};

inline void Random::randomHex(uint8_t *bytes, int size) {
  int i = 0;
  while (i + 8 < size) {
    uint32_t r = next();
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
  }
  uint32_t r = next();
  while (i < size) {
    bytes[i++] = "0123456789abcdef"[r & 0b1111];
    r >>= 4;
  }
}

inline Random seededRandom() {
  FILE *f = fopen("/dev/urandom", "r");
  if (f == nullptr) {
    fprintf(stderr, "Failed to open /dev/urandom\n");
    abort();
  }
  uint64_t seed[2];
  if (fread(seed, sizeof(seed[0]), sizeof(seed) / sizeof(seed[0]), f) !=
      sizeof(seed) / sizeof(seed[0])) {
    fprintf(stderr, "Failed to read from /dev/urandom\n");
    abort();
  }
  fclose(f);
  return Random{seed[0], seed[1]};
}

inline thread_local Random gRandom = seededRandom();

// ==================== END RANDOM IMPL ====================

// ==================== BEGIN ARBITRARY IMPL ====================

/// Think of `Arbitrary` as an attacker-controlled random number generator.
/// Usually you want your random number generator to be fair, so that you can
/// sensibly analyze probabilities. E.g. The analysis that shows that quicksort
/// is expected O(n log n) with a random pivot relies on the random pivot being
/// selected uniformly from a fair distribution.
///
/// Other times you want your randomness to be diabolically unfair, like when
/// looking for bugs and fuzzing. The random-number-like interface is still
/// convenient here, but you can potentially get much better coverage by
/// allowing the possibility of e.g. flipping heads 100 times in a row.
///
/// When it runs out of entropy, it always returns 0.
struct Arbitrary {
  Arbitrary() = default;

  explicit Arbitrary(std::span<const uint8_t> bytecode) : bytecode(bytecode) {}

  /// Draws an arbitrary uint32_t
  uint32_t next() { return consume<4>(); }

  /// Draws an arbitrary element from [0, s)
  uint32_t bounded(uint32_t s);

  /// Fill `bytes` with `size` random hex bytes
  void randomHex(uint8_t *bytes, int size) {
    for (int i = 0; i < size;) {
      uint8_t arbitrary = consume<1>();
      bytes[i++] = "0123456789abcdef"[arbitrary & 0xf];
      arbitrary >>= 4;
      if (i < size) {
        bytes[i++] = "0123456789abcdef"[arbitrary & 0xf];
      }
    }
  }

  bool hasEntropy() const { return bytecode.size() != 0; }

private:
  uint8_t consumeByte() {
    if (bytecode.size() == 0) {
      return 0;
    }
    auto result = bytecode[0];
    bytecode = bytecode.subspan(1, bytecode.size() - 1);
    return result;
  }

  template <int kBytes> uint32_t consume() {
    uint32_t result = 0;
    static_assert(kBytes <= 4);
    for (int i = 0; i < kBytes; ++i) {
      result <<= 8;
      result |= consumeByte();
    }
    return result;
  }

  std::span<const uint8_t> bytecode;
};

inline uint32_t Arbitrary::bounded(uint32_t s) {
  if (s == 1) {
    return 0;
  }
  switch (32 - std::countl_zero(s - 1)) {
  case 1:
  case 2:
  case 3:
  case 4:
  case 5:
  case 6:
  case 7:
  case 8:
    return consume<1>() % s;
  case 9:
  case 10:
  case 11:
  case 12:
  case 13:
  case 14:
  case 15:
  case 16:
    return consume<2>() % s;
  case 17:
  case 18:
  case 19:
  case 20:
  case 21:
  case 22:
  case 23:
  case 24:
    return consume<3>() % s;
  default:
    return consume<4>() % s;
  }
}

inline Arbitrary gArbitrary;

inline void initFuzz(const uint8_t *data, size_t size) {
  gArbitrary = Arbitrary{{data, size}};
  uint64_t state = gArbitrary.next();
  uint64_t seq = gArbitrary.next();
  gRandom = Random{state, seq};
}

// ==================== END ARBITRARY IMPL ====================

// GCOVR_EXCL_STOP