#include "arena_allocator.hpp"
#include <cassert>
#include <iomanip>
#include <limits>
#include <vector>

ArenaAllocator::~ArenaAllocator() {
  while (current_block_) {
    Block *prev = current_block_->prev;
    std::free(current_block_);
    current_block_ = prev;
  }
}

ArenaAllocator::ArenaAllocator(ArenaAllocator &&other) noexcept
    : initial_block_size_(other.initial_block_size_),
      current_block_(other.current_block_) {
  other.current_block_ = nullptr;
}

ArenaAllocator &ArenaAllocator::operator=(ArenaAllocator &&other) noexcept {
  if (this != &other) {
    while (current_block_) {
      Block *prev = current_block_->prev;
      std::free(current_block_);
      current_block_ = prev;
    }

    initial_block_size_ = other.initial_block_size_;
    current_block_ = other.current_block_;

    other.current_block_ = nullptr;
  }
  return *this;
}

void ArenaAllocator::reset() {
  if (!current_block_) {
    return;
  }

  // Find the first block by traversing backwards
  Block *first_block = current_block_;
  while (first_block && first_block->prev) {
    first_block = first_block->prev;
  }

  // Free all blocks after the first one
  Block *current = current_block_;
  while (current && current != first_block) {
    Block *prev = current->prev;
    std::free(current);
    current = prev;
  }

  // Update first block's counters to reflect only itself
  if (first_block) {
    first_block->total_size = first_block->size;
    first_block->total_used = 0;
  }

  current_block_ = first_block;
  current_block_->offset = 0;
}

void *ArenaAllocator::realloc_raw(void *ptr, uint32_t old_size,
                                  uint32_t new_size, uint32_t alignment) {
  if (ptr == nullptr) {
    return allocate_raw(new_size, alignment);
  }

  if (new_size == old_size) {
    return ptr;
  }

  // Assert that we have a current block if ptr is not null
  assert(current_block_ &&
         "realloc called with non-null ptr but no current block exists");

  // Assert that offset is large enough (should always be true for
  // valid callers)
  assert(current_block_->offset >= old_size &&
         "offset must be >= old_size for valid last allocation");

  // Check if this was the last allocation by comparing with expected location
  char *expected_last_alloc_start =
      current_block_->data() + current_block_->offset - old_size;

  if (ptr == expected_last_alloc_start) {
    // This is indeed the last allocation
    if (new_size > old_size) {
      // Growing - check if we have space
      size_t additional_space_needed = new_size - old_size;

      if (current_block_->offset + additional_space_needed <=
          current_block_->size) {
        // We can extend in place
        current_block_->offset += additional_space_needed;
        return ptr;
      }
    } else {
      // Shrinking - just update the offset
      size_t space_to_free = old_size - new_size;
      current_block_->offset -= space_to_free;
      return new_size == 0 ? nullptr : ptr;
    }
  }

  // Can't extend in place
  if (new_size == 0) {
    // For non-last allocations, we can't reclaim memory but still return
    // nullptr
    return nullptr;
  }

  if (new_size <= old_size) {
    // Shrinking - no need to allocate, just return the same pointer
    return ptr;
  }

  // Growing but can't extend in place - need to allocate new space and copy
  void *new_ptr = allocate_raw(new_size, alignment);
  if (new_ptr && ptr) {
    // Copy all the old data since we're growing
    std::memcpy(new_ptr, ptr, old_size);
  }

  return new_ptr;
}

void ArenaAllocator::debug_dump(std::ostream &out, bool show_memory_map,
                                bool show_content, size_t content_limit) const {
  out << "=== Arena Debug Dump ===" << std::endl;

  if (!current_block_) {
    out << "Arena is empty (no blocks allocated)" << std::endl;
    out << "Initial block size: " << initial_block_size_ << " bytes"
        << std::endl;
    return;
  }

  // Build list of blocks from current to first
  std::vector<Block *> blocks;
  Block *block = current_block_;
  while (block) {
    blocks.push_back(block);
    block = block->prev;
  }

  // Overall statistics
  size_t used = this->used_bytes();
  size_t total_alloc = this->total_allocated();
  double utilization = total_alloc > 0 ? (100.0 * used / total_alloc) : 0.0;

  out << "Total allocated: " << total_alloc << " bytes across " << num_blocks()
      << " blocks" << std::endl;
  out << "Currently used: " << used << " bytes (" << std::fixed
      << std::setprecision(1) << utilization << "% utilization)" << std::endl;
  out << "Available in current: " << available_in_current_block() << " bytes"
      << std::endl;
  out << std::endl;

  out << "Block Chain (newest to oldest):" << std::endl;

  // Display blocks in reverse order (current first)
  for (int i = 0; i < static_cast<int>(blocks.size()); ++i) {
    Block *b = blocks[i];

    // Calculate used bytes in this specific block
    size_t block_used = b->offset;

    double block_util = b->size > 0 ? (100.0 * block_used / b->size) : 0.0;

    out << "Block #" << (blocks.size() - i) << ": " << b->size << " bytes "
        << "[used: " << block_used << "/" << b->size << " = " << std::fixed
        << std::setprecision(1) << block_util << "%]";

    if (i == 0) {
      out << " <- current";
    }
    out << std::endl;

    // Show memory map if requested
    if (show_memory_map && b->size > 0) {
      const int map_width = 60;
      int used_chars = static_cast<int>((map_width * block_used) / b->size);
      used_chars = std::min(used_chars, map_width);

      out << "  [";
      for (int j = 0; j < map_width; ++j) {
        if (j < used_chars) {
          out << "#";
        } else {
          out << ".";
        }
      }
      out << "] (# = used, . = free)" << std::endl;
    }
  }

  out << std::endl;
  out << "Block addresses and relationships:" << std::endl;
  for (int i = 0; i < static_cast<int>(blocks.size()); ++i) {
    Block *b = blocks[i];
    out << "Block #" << (blocks.size() - i) << " @ " << static_cast<void *>(b)
        << " -> data @ " << static_cast<void *>(b->data());
    if (b->prev) {
      out << " (prev: " << static_cast<void *>(b->prev) << ")";
    } else {
      out << " (first block)";
    }
    out << std::endl;
  }

  // Show memory contents if requested
  if (show_content) {
    out << std::endl;
    out << "Memory Contents:" << std::endl;

    for (int i = 0; i < static_cast<int>(blocks.size()); ++i) {
      Block *b = blocks[i];
      int block_num = static_cast<int>(blocks.size()) - i;

      // Calculate used bytes in this specific block
      size_t block_used = b->offset;

      if (block_used == 0) {
        out << "Block #" << block_num << ": No content (empty)" << std::endl;
        continue;
      }

      size_t bytes_to_show = std::min(block_used, content_limit);
      out << "Block #" << block_num << " (first " << bytes_to_show << " of "
          << block_used << " used bytes):" << std::endl;

      const char *data = b->data();
      dump_memory_contents(out, data, bytes_to_show);

      if (bytes_to_show < block_used) {
        out << "  ... (" << (block_used - bytes_to_show) << " more bytes)"
            << std::endl;
      }
      out << std::endl;
    }
  }
}

void ArenaAllocator::add_block(size_t size) {
  Block *new_block = Block::create(size, current_block_);
  current_block_ = new_block;
}

size_t ArenaAllocator::calculate_next_block_size(size_t required_size) const {
  size_t doubled_size = (current_block_ ? current_block_->size : 0) * 2;
  doubled_size =
      std::min<size_t>(doubled_size, std::numeric_limits<uint32_t>::max());
  return std::max(required_size, doubled_size);
}

void ArenaAllocator::dump_memory_contents(std::ostream &out, const char *data,
                                          size_t size) {
  const int bytes_per_line = 16;

  for (int64_t offset = 0; offset < static_cast<int64_t>(size);
       offset += bytes_per_line) {
    // Print offset
    out << "  0x" << std::setfill('0') << std::setw(4) << std::hex << offset
        << ": ";

    int bytes_in_line =
        std::min(bytes_per_line, static_cast<int>(size - offset));

    // Print hex bytes
    for (int i = 0; i < bytes_per_line; ++i) {
      if (i < bytes_in_line) {
        unsigned char byte = static_cast<unsigned char>(data[offset + i]);
        out << std::setfill('0') << std::setw(2) << std::hex
            << static_cast<int>(byte);
      } else {
        out << "  "; // Padding for incomplete lines
      }

      // Add space every 4 bytes for readability
      if ((i + 1) % 4 == 0) {
        out << " ";
      }
    }

    // Print ASCII representation
    out << " |";
    for (int i = 0; i < bytes_in_line; ++i) {
      char c = data[offset + i];
      if (c >= 32 && c <= 126) { // Printable ASCII
        out << c;
      } else {
        out << '.'; // Non-printable characters
      }
    }
    out << "|" << std::dec << std::endl;
  }
}
