#pragma once

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <limits>
#include <new>
#include <span>
#include <type_traits>
#include <typeinfo>
#include <utility>

/**
 * @brief A high-performance arena allocator for bulk allocations.
 *
 * ArenaAllocator provides extremely fast memory allocation (~1ns per
 * allocation) by allocating large blocks and serving allocations from them
 * sequentially. It's designed for scenarios where many small objects need to be
 * allocated and can all be deallocated together.
 *
 * ## Key Features:
 * - **Ultra-fast allocation**: ~1ns per allocation vs ~20-270ns for malloc
 * - **Lazy initialization**: No memory allocated until first use
 * - **Intrusive linked list**: Minimal memory overhead using backward-linked
 * blocks
 * - **Geometric growth**: Block sizes double to minimize allocations
 * - **Memory efficient reset**: Frees unused blocks to prevent memory leaks
 * - **Proper alignment**: Respects alignment requirements for all types
 *
 * ## Performance Characteristics:
 * - Allocation: O(1) amortized
 * - Memory tracking: O(1) using accumulated counters
 * - Reset: O(n) where n is number of blocks (but frees memory)
 * - Destruction: O(n) where n is number of blocks
 *
 * ## Usage Examples:
 * ```cpp
 * ArenaAllocator arena(1024);
 * void* ptr = arena.allocate_raw(100);
 * int* num = arena.construct<int>(42);
 * arena.reset(); // Reuse arena memory
 * ```
 *
 * ## Memory Management:
 * - Individual objects cannot be freed (by design)
 * - All memory is freed when the allocator is destroyed
 * - reset() frees all blocks except the first one
 * - Move semantics transfer ownership of all blocks
 *
 * ## Thread Safety:
 * ArenaAllocator is **not thread-safe** - concurrent access from multiple
 * threads requires external synchronization. However, this design is
 * intentional for performance reasons and the WeaselDB architecture ensures
 * thread safety through ownership patterns:
 *
 * ### Safe Usage Patterns in WeaselDB:
 * - **Per-Connection Instances**: Each Connection owns its own ArenaAllocator
 *   instance, accessed only by the thread that currently owns the connection
 * - **Single Owner Principle**: Connection ownership transfers atomically
 * between threads using unique_ptr, ensuring only one thread accesses the arena
 * at a time
 *
 * ### Thread Ownership Model:
 * 1. **Network Thread**: Claims connection ownership, accesses arena for I/O
 * buffers
 * 2. **Handler Thread**: Can take ownership via unique_ptr.release(), uses
 * arena for request parsing and response generation
 * 3. **Background Thread**: Can receive ownership for async processing, uses
 * arena for temporary data structures
 * 4. **Return Path**: Connection (and its arena) safely returned via
 *    Server::release_back_to_server()
 *
 * ### Why This Design is Thread-Safe:
 * - **Exclusive Access**: Only the current owner thread should access the arena
 * - **Transfer Points**: Ownership transfers happen at well-defined
 * synchronization points with proper memory barriers.
 * - **No Shared State**: Each arena is completely isolated - no shared data
 *   between different arena instances
 *
 * @warning Do not share ArenaAllocator instances between threads. Use separate
 * instances per thread or per logical unit of work.
 */
struct ArenaAllocator {
private:
  /**
   * @brief Internal block structure for the intrusive linked list.
   *
   * Each block contains:
   * - The actual data storage immediately following the Block header
   * - Backward pointer to previous block (intrusive linked list)
   * - Accumulated counters for O(1) tracking operations
   */
  struct Block {
    uint32_t size;     ///< Size of this block's data area
    uint32_t offset;   ///< The offset of the first unused byte in the data area
    size_t total_size; ///< Accumulated size of this block + all previous blocks
    size_t total_used; ///< Accumulated offsets of previous blocks
    Block *prev;       ///< Pointer to previous block (nullptr for first block)

    /**
     * @brief Get pointer to the data area of this block.
     * @return Pointer to the start of the data area (after Block header).
     */
    char *data() { return reinterpret_cast<char *>(this + 1); }

    /**
     * @brief Create a new block with the specified size.
     * @param size Size of the data area for this block
     * @param prev Pointer to the previous block (nullptr for first block)
     * @return Pointer to the newly created block
     * @note Prints error to stderr and calls std::abort() if memory allocation
     * fails
     */
    static Block *create(size_t size, Block *prev) {
      if (size > std::numeric_limits<uint32_t>::max()) {
        std::fprintf(
            stderr,
            "ArenaAllocator: Block size %zu exceeds maximum uint32_t value\n",
            size);
        std::abort();
      }
      void *memory = std::aligned_alloc(
          alignof(Block), align_up(sizeof(Block) + size, alignof(Block)));
      if (!memory) {
        std::fprintf(
            stderr,
            "ArenaAllocator: Failed to allocate memory block of size %zu\n",
            size);
        std::abort();
      }
      size_t total_size = size + (prev ? prev->total_size : 0);
      size_t total_used = prev ? prev->total_used + prev->offset : 0;
      Block *block = new (memory)
          Block{uint32_t(size), /*offset*/ 0, total_size, total_used, prev};
      return block;
    }
  };

public:
  /**
   * @brief Construct an ArenaAllocator with the specified initial block size.
   *
   * No memory is allocated until the first allocation request (lazy
   * initialization). The initial block size is used for the first block and as
   * the baseline for geometric growth.
   *
   * @param initial_size Size in bytes for the first block (default: 1024)
   */
  explicit ArenaAllocator(size_t initial_size = 1024)
      : initial_block_size_(initial_size), current_block_(nullptr) {}

  /**
   * @brief Destructor - frees all allocated blocks.
   *
   * Traverses the intrusive linked list backwards from current_block_,
   * freeing each block. This ensures no memory leaks.
   */
  ~ArenaAllocator();

  /// Copy construction is not allowed (would be expensive and error-prone)
  ArenaAllocator(const ArenaAllocator &) = delete;
  /// Copy assignment is not allowed (would be expensive and error-prone)
  ArenaAllocator &operator=(const ArenaAllocator &) = delete;

  /**
   * @brief Move constructor - transfers ownership of all blocks.
   * @param other The ArenaAllocator to move from (will be left empty)
   */
  ArenaAllocator(ArenaAllocator &&other) noexcept;

  /**
   * @brief Move assignment operator - transfers ownership of all blocks.
   *
   * Frees any existing blocks in this allocator before taking ownership
   * of blocks from the other allocator.
   *
   * @param other The ArenaAllocator to move from (will be left empty)
   * @return Reference to this allocator
   */
  ArenaAllocator &operator=(ArenaAllocator &&other) noexcept;

  /**
   * @brief Allocate raw memory with the specified size and alignment.
   *
   * This is the core allocation method providing ~1ns allocation performance.
   * It performs lazy initialization on first use and automatically grows
   * the arena when needed using geometric growth (doubling block sizes).
   *
   * For type-safe allocation, prefer the allocate<T>() template method.
   *
   * @param size Number of bytes to allocate (0 returns nullptr)
   * @param alignment Required alignment (default: alignof(std::max_align_t))
   * @return Pointer to allocated memory, or nullptr if size is 0
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Performance:
   * - O(1) amortized allocation time
   * - Respects alignment requirements with minimal padding
   * - Automatically creates new blocks when current block is exhausted
   *
   * @note This method is kept inline for maximum performance (~1ns allocation).
   */
  void *allocate_raw(uint32_t size,
                     size_t alignment = alignof(std::max_align_t)) {
    if (size == 0) {
      return nullptr;
    }

    if (!current_block_) {
      size_t block_size = std::max(size, initial_block_size_);
      add_block(block_size);
    }

    char *block_start = current_block_->data();
    uintptr_t block_addr = reinterpret_cast<uintptr_t>(block_start);
    size_t aligned_offset =
        align_up(block_addr + current_block_->offset, alignment) - block_addr;

    if (aligned_offset + size > current_block_->size) {
      size_t next_block_size = calculate_next_block_size(size);
      add_block(next_block_size);
      block_start = current_block_->data();
      block_addr = reinterpret_cast<uintptr_t>(block_start);
      aligned_offset = align_up(block_addr, alignment) - block_addr;
    }

    void *ptr = block_start + aligned_offset;
    current_block_->offset = aligned_offset + size;

    return ptr;
  }

  /**
   * @brief Reallocate memory, extending in place if possible or copying to a
   * new location.
   *
   * This method provides realloc-like functionality for the arena allocator.
   * If the given pointer was the last allocation and there's sufficient space
   * in the current block to extend it, the allocation is grown in place.
   * Otherwise, a new allocation is made and the old data is copied.
   *
   * @param ptr Pointer to the existing allocation (must be from this allocator)
   * @param old_size Size of the existing allocation in bytes
   * @param new_size Desired new size in bytes
   * @param alignment Required alignment. Defaults to
   * `alignof(std::max_align_t)`
   * @return Pointer to the reallocated memory (may be the same as ptr or
   * different)
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Behavior:
   * - If new_size == old_size, returns ptr unchanged
   * - If new_size == 0, returns nullptr (no deallocation occurs)
   * - If ptr is null, behaves like allocate(new_size, alignment)
   * - If ptr was the last allocation and space exists, extends in place
   *
   *
   * ## Safety Notes:
   * - The caller must provide the correct old_size - this is not tracked
   * - The old pointer becomes invalid if a copy occurs
   * - Like malloc/realloc, the contents beyond old_size are uninitialized
   * - When copying to new location, uses the specified alignment
   * - **Shrinking behavior**: If `new_size < old_size` and the allocation
   *   is *not* the most recent one, this function returns the original
   *   pointer, but **no memory is reclaimed**. The arena design does not
   *   support freeing memory from the middle of a block.
   */
  void *realloc_raw(void *ptr, uint32_t old_size, uint32_t new_size,
                    uint32_t alignment = alignof(std::max_align_t));

  /**
   * @brief Type-safe version of realloc_raw for arrays of type T.
   *
   * @param ptr Pointer to the existing allocation (must be from this allocator)
   *            If nullptr, behaves like allocate<T>(new_size)
   * @param old_size Size of the existing allocation in number of T objects
   *                 Ignored if ptr is nullptr
   * @param new_size Desired new size in number of T objects
   * @return Pointer to the reallocated memory (may be the same as ptr or
   * different)
   * @note Follows standard realloc() semantics: realloc(nullptr, size) ==
   * malloc(size)
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails or size overflow occurs
   */
  template <typename T>
  T *realloc(T *ptr, uint32_t old_size, uint32_t new_size) {
    if (size_t(new_size) * sizeof(T) > std::numeric_limits<uint32_t>::max()) {
      std::fprintf(stderr,
                   "ArenaAllocator: Reallocation size overflow for type %s "
                   "(new_size=%u, sizeof(T)=%zu)\n",
                   typeid(T).name(), new_size, sizeof(T));
      std::abort();
    }
    return static_cast<T *>(realloc_raw(ptr, old_size * sizeof(T),
                                        new_size * sizeof(T), alignof(T)));
  }

  /**
   * @brief Smart pointer for arena-allocated objects with non-trivial
   * destructors.
   *
   * ArenaAllocator::Ptr calls the destructor but does not free memory (assumes
   * arena allocation). This provides RAII semantics for objects that need
   * cleanup without the overhead of individual memory deallocation.
   *
   * @tparam T The type of object being managed
   */
  template <typename T> struct Ptr {
    Ptr() noexcept : ptr_(nullptr) {}

    explicit Ptr(T *ptr) noexcept : ptr_(ptr) {}

    Ptr(const Ptr &) = delete;
    Ptr &operator=(const Ptr &) = delete;

    Ptr(Ptr &&other) noexcept : ptr_(other.ptr_) { other.ptr_ = nullptr; }

    Ptr &operator=(Ptr &&other) noexcept {
      if (this != &other) {
        reset();
        ptr_ = other.ptr_;
        other.ptr_ = nullptr;
      }
      return *this;
    }

    ~Ptr() { reset(); }

    T *operator->() const noexcept { return ptr_; }
    T &operator*() const noexcept { return *ptr_; }

    T *get() const noexcept { return ptr_; }

    explicit operator bool() const noexcept { return ptr_ != nullptr; }

    T *release() noexcept {
      T *result = ptr_;
      ptr_ = nullptr;
      return result;
    }

    void reset(T *new_ptr = nullptr) noexcept {
      if (ptr_) {
        ptr_->~T();
      }
      ptr_ = new_ptr;
    }

  private:
    T *ptr_;
  };

  /**
   * @brief Construct an object of type T in the arena using placement new.
   *
   * This method returns different types based on whether T is trivially
   * destructible:
   * - For trivially destructible types: returns T* (raw pointer)
   * - For non-trivially destructible types: returns ArenaAllocator::Ptr<T>
   * (smart pointer that calls destructor)
   *
   * @tparam T The type of object to construct
   * @tparam Args Types of constructor arguments
   * @param args Arguments to forward to T's constructor
   * @return T* for trivially destructible types, ArenaAllocator::Ptr<T>
   * otherwise
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   */
  template <typename T, typename... Args> auto construct(Args &&...args) {
    void *ptr = allocate_raw(sizeof(T), alignof(T));
    T *obj = new (ptr) T(std::forward<Args>(args)...);

    if constexpr (std::is_trivially_destructible_v<T>) {
      return obj;
    } else {
      return Ptr<T>(obj);
    }
  }

  /**
   * @brief Allocate space for an array of size T objects with proper alignment.
   *
   * This is a type-safe convenience method that combines sizing and alignment
   * calculations for allocating arrays of type T. It's preferred over calling
   * allocate_raw() directly as it prevents common errors with size calculations
   * and alignment requirements.
   *
   * @tparam T The type of objects to allocate space for (must be trivially
   * destructible)
   * @param size Number of T objects to allocate space for
   * @return Pointer to allocated memory suitable for constructing an array of T
   * objects
   * @note Prints error to stderr and calls std::abort() if memory allocation
   * fails
   *
   * ## Type Requirements:
   * T must be trivially destructible (std::is_trivially_destructible_v<T>).
   * This ensures consistency with the arena allocator's design where
   * destructors are never called.
   *
   *
   * ## Note:
   * This method only allocates memory - it does not construct objects.
   * Use placement new or other initialization methods as needed.
   */
  template <typename T> T *allocate(uint32_t size) {
    static_assert(
        std::is_trivially_destructible_v<T>,
        "ArenaAllocator::allocate requires trivially destructible types. "
        "Objects allocated in the arena will not have their destructors "
        "called.");
    if (size == 0) {
      return nullptr;
    }
    if (size_t(size) * sizeof(T) > std::numeric_limits<uint32_t>::max()) {
      std::fprintf(stderr,
                   "ArenaAllocator: Allocation size overflow for type %s "
                   "(size=%u, sizeof(T)=%zu)\n",
                   typeid(T).name(), size, sizeof(T));
      std::abort();
    }
    void *ptr = allocate_raw(sizeof(T) * size, alignof(T));
    return static_cast<T *>(ptr);
  }

  /**
   * @brief Reset the allocator to reuse the first block, freeing all others.
   *
   * This method provides memory-efficient reset behavior by:
   * 1. Keeping the first block for reuse
   * 2. Freeing all subsequent blocks to prevent memory leaks
   * 3. Resetting allocation position to the start of the first block
   *
   * If no blocks have been allocated yet, this is a no-op.
   *
   * ## Performance:
   * - O(n) where n is the number of blocks to free
   * - Prevents memory leaks by freeing unused blocks
   * - Faster than destroying and recreating the allocator
   *
   */
  void reset();

  /**
   * @brief Get the total number of bytes allocated across all blocks.
   *
   * Uses O(1) accumulated counters for fast retrieval.
   *
   * @return Total allocated bytes, or 0 if no blocks exist
   */
  size_t total_allocated() const {
    return current_block_ ? current_block_->total_size : 0;
  }

  /**
   * @brief Get the number of bytes currently used for allocations.
   *
   * This includes all fully used previous blocks plus the used portion
   * of the current block. Uses O(1) accumulated counters.
   *
   * @return Number of bytes in use
   */
  size_t used_bytes() const {
    if (!current_block_) {
      return 0;
    }
    return current_block_->total_used + current_block_->offset;
  }

  /**
   * @brief Get the number of bytes available in the current block.
   *
   * @return Available bytes in current block, or 0 if no blocks exist
   */
  size_t available_in_current_block() const {
    return current_block_ ? current_block_->size - current_block_->offset : 0;
  }

  /**
   * @brief Get all available space in the current block and claim it
   * immediately.
   *
   * This method returns a pointer to all remaining space in the current block
   * and immediately marks it as used in the arena. The caller should use
   * realloc() to shrink the allocation to the actual amount needed.
   *
   * If no block exists or current block is full, creates a new block.
   *
   * @return Pointer to allocated space and the number of bytes allocated
   * @note The caller must call realloc() to return unused space
   * @note This is designed for speculative operations like printf formatting
   * @note Postcondition: always returns at least 1 byte
   */
  struct AllocatedSpace {
    char *ptr;
    size_t allocated_bytes;
  };

  AllocatedSpace allocate_remaining_space() {
    if (!current_block_ || available_in_current_block() == 0) {
      add_block(initial_block_size_);
    }

    char *allocated_ptr = current_block_->data() + current_block_->offset;
    size_t available = available_in_current_block();

    // Claim all remaining space
    current_block_->offset = current_block_->size;

    return {allocated_ptr, available};
  }

  /**
   * @brief Get the total number of blocks in the allocator.
   *
   * @note This function is primarily used for testing and debugging.
   * It has O(n) complexity as it traverses the entire block chain.
   */
  size_t num_blocks() const {
    size_t result = 0;
    for (auto *p = current_block_; p != nullptr; p = p->prev) {
      ++result;
    }
    return result;
  }

  /**
   * @brief Debug function to visualize the arena's layout and contents.
   *
   * @note This function is intended for testing, debugging, and development
   * only. It should not be used in production code due to performance overhead.
   *
   * Prints a detailed breakdown of all blocks, memory usage, and allocation
   * patterns. This is useful for understanding memory fragmentation and
   * allocation behavior during development and debugging.
   *
   * @param out Output stream to write debug information to (default: std::cout)
   * @param show_memory_map If true, shows a visual memory map of used/free
   * space
   * @param show_content If true, shows actual memory contents in hex and ASCII
   * @param content_limit Maximum bytes of content to show per block (default:
   * 256)
   */
  void debug_dump(std::ostream &out = std::cout, bool show_memory_map = false,
                  bool show_content = false, size_t content_limit = 256) const;

private:
  /**
   * @brief Add a new block with the specified size to the allocator.
   *
   * Creates a new block and makes it the current block. Updates all
   * accumulated counters automatically through Block::create().
   *
   * @param size Size of the data area for the new block
   */
  void add_block(size_t size);

  /**
   * @brief Calculate the size for the next block using geometric growth.
   *
   * Uses a doubling strategy to minimize the number of blocks while
   * ensuring large allocations are handled efficiently.
   *
   * @param required_size Minimum size needed for the allocation
   * @return Size for the next block (max of required_size and doubled current
   * size)
   */
  size_t calculate_next_block_size(size_t required_size) const;

  /**
   * @brief Align a value up to the specified alignment boundary.
   *
   * Uses bit manipulation for efficient alignment calculation.
   * Only works with power-of-2 alignments.
   *
   * This method is kept inline in the header for maximum performance
   * as it's called in the hot allocation path and benefits from inlining.
   *
   * @param value The value to align
   * @param alignment The alignment boundary (must be power of 2)
   * @return The aligned value
   */
  static size_t align_up(size_t value, size_t alignment) {
    if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
      return value;
    }
    return (value + alignment - 1) & ~(alignment - 1);
  }

  /**
   * @brief Dump memory contents in hex/ASCII format.
   *
   * Displays memory in the classic hex dump format with 16 bytes per line,
   * showing both hexadecimal values and ASCII representation.
   *
   * @param out Output stream to write to
   * @param data Pointer to the memory to dump
   * @param size Number of bytes to dump
   */
  static void dump_memory_contents(std::ostream &out, const char *data,
                                   size_t size);

  /// Size used for the first block and baseline for geometric growth
  uint32_t initial_block_size_;
  /// Pointer to the current (most recent) block, or nullptr if no blocks exist
  Block *current_block_;
};

/**
 * @brief STL-compatible allocator that uses ArenaAllocator for memory
 * management.
 * @tparam T The type of objects to allocate
 */
template <typename T> class ArenaStlAllocator {
public:
  using value_type = T;
  using pointer = T *;
  using const_pointer = const T *;
  using reference = T &;
  using const_reference = const T &;
  using size_type = std::size_t;
  using difference_type = std::ptrdiff_t;

  template <typename U> struct rebind {
    using other = ArenaStlAllocator<U>;
  };

  explicit ArenaStlAllocator(ArenaAllocator *arena) noexcept : arena_(arena) {}

  template <typename U>
  ArenaStlAllocator(const ArenaStlAllocator<U> &other) noexcept
      : arena_(other.arena_) {}

  T *allocate(size_type n) {
    if (n == 0)
      return nullptr;
    return arena_->allocate<T>(n);
  }

  void deallocate(T *, size_type) noexcept {
    // Arena allocator doesn't support individual deallocation
  }

  template <typename U>
  bool operator==(const ArenaStlAllocator<U> &other) const noexcept {
    return arena_ == other.arena_;
  }

  template <typename U>
  bool operator!=(const ArenaStlAllocator<U> &other) const noexcept {
    return arena_ != other.arena_;
  }

  ArenaAllocator *arena_;

  template <typename U> friend class ArenaStlAllocator;
};

/// Simple arena-aware vector that doesn't have a destructor
/// Safe to return as span because both the vector and its data are
/// arena-allocated Uses arena's realloc() for efficient growth without copying
/// when possible
template <typename T> struct ArenaVector {
  explicit ArenaVector(ArenaAllocator *arena)
      : arena_(arena), data_(nullptr), size_(0), capacity_(0) {}

  void push_back(const T &item) {
    if (size_ >= capacity_) {
      grow();
    }
    data_[size_++] = item;
  }

  T *data() { return data_; }
  const T *data() const { return data_; }
  size_t size() const { return size_; }
  bool empty() const { return size_ == 0; }

  T &operator[](size_t index) { return data_[index]; }
  const T &operator[](size_t index) const { return data_[index]; }

  void clear() { size_ = 0; }

  // Implicit conversion to std::span
  operator std::span<T>() { return std::span<T>(data_, size_); }
  operator std::span<const T>() const {
    return std::span<const T>(data_, size_);
  }

  // Iterator support for range-based for loops
  T *begin() { return data_; }
  const T *begin() const { return data_; }
  T *end() { return data_ + size_; }
  const T *end() const { return data_ + size_; }

  // No destructor - arena cleanup handles memory

private:
  void grow() {
    size_t new_capacity = capacity_ == 0 ? 8 : capacity_ * 2;

    // arena.realloc() handles nullptr like standard realloc() - acts like
    // malloc() This avoids copying when growing in-place is possible
    data_ = arena_->realloc(data_, capacity_, new_capacity);
    capacity_ = new_capacity;
  }

  ArenaAllocator *arena_;
  T *data_;
  size_t size_;
  size_t capacity_;
};