// arena.h – Arena allocator and string utilities (nullprogram style)

#include <stddef.h>
#include <stdint.h>
#include <stdarg.h>
#include <string.h>

// helpers for printing/defns
#define pstr(cstr) (int)((cstr).len), (cstr).data
#define kstr(str) str_from_buf(str, KSTR_LEN(str))
#define skstr(str) {.data=(byte*)str, .len=KSTR_LEN(str)}

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-function"

// -----------------------------------------------------------------------------
// Arena – simple bump allocator (allocates from the end backwards)
// -----------------------------------------------------------------------------
//
// The arena uses two pointers: 'beg' points to the start of the memory block,
// 'end' points to the current allocation frontier. Memory is allocated by
// moving 'end' *downwards* (towards 'beg'), so the beginning of the block
// stays free. This is key for the concatenation trick below.
//
//   +-------------------+--------------------+
//   |                   |                    |
//   beg                end                   (original end)
//   (free space)        (already allocated)
//
// Resetting the arena simply sets 'end' back to the original top.

typedef struct arena {
  uint8_t *beg;   // start of the memory block (fixed)
  uint8_t *end;   // current allocation frontier (moves downwards)
  uint8_t *lim;   // original end of the block (for reset)
} arena;

// Initialize an arena from an existing buffer.
static void arena_init(arena *a, void *buffer, size_t size) {
  a->beg = (uint8_t *)buffer;
  a->end = a->beg + size;
  a->lim = a->end;
}

// Reset the arena – discards all allocations.
static void arena_reset(arena *a) {
  a->end = a->lim;
}

// Core allocation function – returns zero‑initialised memory.
//   objsize : size of one object
//   align   : required alignment (must be a power of two)
//   count   : number of objects to allocate
// Returns NULL if out of memory.
static void *arena_alloc(arena *a, size_t objsize, size_t align, size_t count) {
  // Alignment adjustment: move 'end' down to the next aligned address.
  size_t pad = (uintptr_t)a->end & (align - 1);
  if (pad) pad = align - pad;

  // Total bytes needed, checking for overflow.
  size_t total = objsize * count;
  if (objsize != 0 && total / objsize != count) return NULL; // overflow

  // Check if we have enough space.
  if (total + pad > (size_t)(a->end - a->beg)) return NULL;

  // Bump the pointer downwards.
  a->end -= pad + total;
  return memset(a->end, 0, total);
}

// Convenience macros – similar to C++'s new.
#define new(a, t, n)   ((t *)arena_alloc((a), sizeof(t), _Alignof(t), (n)))
#define new_array(a, t, n) new(a, t, n)

// -----------------------------------------------------------------------------
// String view – non‑owning, pointer + length
// (a string with len=0 is counted as an empty string, regardless of the value of data)
// -----------------------------------------------------------------------------
typedef struct str {
  uint8_t *data;
  ptrdiff_t len;
} str;

// Create a str from a null‑terminated C string.
static str str_from_cstr(const char *s) {
  str r;
  r.data = (uint8_t *)s;
  r.len = (ptrdiff_t)strlen(s);
  return r;
}

// Create a str from a byte buffer.
static str str_from_buf(const void *buf, size_t len) {
  str r;
  r.data = (uint8_t *)buf;
  r.len = (ptrdiff_t)len;
  return r;
}

// destructively update dst in-place and change its len (must be writable)
// (normally used when dst/src will always have the same length)
static void str_update(str *dst, str src) {
  assert(dst->len >= src.len);
  memcpy(dst->data, src.data, src.len);
  dst->len = src.len;
}

// copying versions
static str str_new_from_buf(arena *a, const void *buf, size_t len) {
  if (!buf || len == 0) return (str){0};
  uint8_t *copy = new(a, uint8_t, len);
  if (!copy) return (str){0};
  memcpy(copy, buf, len);
  return (str){ .data = copy, .len = (ptrdiff_t)len };
}

static str str_new_from_cstr(arena *a, const char *s) {
  if (!s) return (str){0};
  return str_new_from_buf(a, s, strlen(s));
}

static str str_dup(arena *a, str s) {
  if (0==s.len) return (str){0};
  return str_new_from_buf(a, s.data, s.len);
}

// Convert a str to a null‑terminated C string (allocates in the arena!).
static char *str_to_cstr(arena *a, str s) {
  char *cstr = new(a, char, s.len + 1);
  if (!cstr) return NULL;
  memcpy(cstr, s.data, (size_t)s.len);
  cstr[s.len] = 0;
  return cstr;
}

// Compare two strings for equality.
static int str_eq(str a, str b) {
  return a.len == b.len && (a.data == b.data || memcmp(a.data, b.data, (size_t)a.len) == 0);
}

// -----------------------------------------------------------------------------
// Almighty concatenation – builds strings at the front of the arena
// -----------------------------------------------------------------------------
//
// The idea: allocate the *first* part of the string using arena_front_alloc().
// This reserves space at the very beginning of the arena block. Later, we can
// extend the string in place (without copying) as long as nothing else has been
// allocated after it. The arena_front_alloc() function allocates from the front
// (increasing 'beg'), so the string grows to the right.

// Allocate a block from the *front* of the arena. This is separate from the
// normal arena_alloc() which allocates from the end. For normal objects you
// should use arena_alloc() / new(). Use arena_front_alloc() only when you
// plan to later extend the allocation.
static void *arena_front_alloc(arena *a, size_t objsize, size_t align, size_t count) {
  uintptr_t addr = (uintptr_t)(a->beg);
  size_t pad = addr & (align - 1);
  if (pad) pad = align - pad;

  size_t total = objsize * count;
  if (objsize != 0 && total / objsize != count) return NULL;

  if (a->beg + pad + total > a->end) return NULL; // would collide with main allocator

  void *ptr = a->beg + pad;
  a->beg += pad + total;
  return memset(ptr, 0, total);
}

#define new_front(a, t, n) ((t *)arena_front_alloc((a), sizeof(t), _Alignof(t), (n)))

// Append a suffix to a string that was allocated with new_front().
// The base string must be the most recent front allocation, otherwise
// we cannot extend it in place and we fall back to copying.
// Returns an extended str (which may be a new copy if in‑place failed).
static str str_concat(arena *a, str base, const void *suffix, size_t suffix_len) {
  // Check if 'base' is exactly the last thing allocated from the front.
  // This is true if base.data + base.len == a->beg.
  if (base.data + base.len == a->beg) {
    // Extend in place
    size_t total = (size_t)(base.len + suffix_len);
    if (a->beg + total > a->end) {
      // Not enough room – fall back to copy
      goto fallback;
    }
    memcpy(a->beg, suffix, suffix_len);
    a->beg += suffix_len;
    return (str){ .data = base.data, .len = (ptrdiff_t)total };
  }
fallback:
  {
    // Cannot extend in place – allocate a new copy from the front.
    str newstr;
    newstr.len = base.len + (ptrdiff_t)suffix_len;
    newstr.data = new_front(a, uint8_t, (size_t)newstr.len);
    if (!newstr.data) return (str){0};
    memcpy(newstr.data, base.data, (size_t)base.len);
    memcpy(newstr.data + base.len, suffix, suffix_len);
    return newstr;
  }
}

// Convenience: concatenate a null‑terminated C string.
static str str_cat_cstr(arena *a, str base, const char *suffix) {
  return str_concat(a, base, suffix, strlen(suffix));
}

// -----------------------------------------------------------------------------
// Formatted string (printf style) allocated from the front.
// -----------------------------------------------------------------------------
static str str_printf(arena *a, const char *fmt, ...) {
  va_list ap;
  va_start(ap, fmt);
  // First, determine the required length.
  int len = vsnprintf(NULL, 0, fmt, ap);
  va_end(ap);
  if (len < 0) return (str){0};

  str result;
  result.len = len;
  result.data = new_front(a, uint8_t, (size_t)len);
  if (!result.data) return (str){0};

  va_start(ap, fmt);
  vsnprintf((char *)result.data, (size_t)len + 1, fmt, ap);
  va_end(ap);
  return result;
}

// Append formatted text to a string that was allocated at the front.
// Returns the new string (may be the same as `s` if extended in place).
static str str_printf_append(arena *a, str s, const char *fmt, ...) {
    va_list ap;
    va_start(ap, fmt);
    int needed = vsnprintf(NULL, 0, fmt, ap);
    va_end(ap);
    if (needed < 0) return (str){0};

    // Case 1: empty string – just allocate fresh from the front
    if (s.len == 0 || s.data == NULL) {
      str result;
      result.len = needed;
      result.data = new_front(a, uint8_t, (size_t)needed);
      if (!result.data) return (str){0};

      va_start(ap, fmt);
      vsnprintf((char *)result.data, (size_t)needed + 1, fmt, ap);
      va_end(ap);
      return result;
    }

    // Case 2: can we extend in place?
    if (s.data + s.len == a->beg) {
        // Enough room?
        if (a->beg + (size_t)needed <= a->end) {
            va_start(ap, fmt);
            vsnprintf((char *)a->beg, (size_t)needed + 1, fmt, ap);
            va_end(ap);
            a->beg += needed;
            return (str){ .data = s.data, .len = s.len + needed };
        }
    }

    // TODO, we should have some way to abort this request and close the connection instead (maybe an error state in arena checked on end request?)
    die("ERROR str_printf_append failed to append: needed: %d s={%p,%ld} A->beg:%p\n", needed, s.data, s.len, a->beg);
}

// -----------------------------------------------------------------------------
// Scratch arena pattern – pass by value to automatically discard temporaries
// -----------------------------------------------------------------------------
//
// When you need temporary allocations inside a function, take a *copy* of an
// arena. Allocations made on that copy do not affect the caller, and they all
// disappear when the function returns (the copy goes out of scope).
//
// Example:
//   void do_something(arena scratch) {
//       char *tmp = new(&scratch, char, 256);
//       // ... use tmp ...
//   } // scratch memory is implicitly freed

#pragma GCC diagnostic pop