firebird

defmodule Firebird.Memory do
  @moduledoc """
  Structured memory management for WASM instances.

  Raw `Firebird.read_memory/3` and `Firebird.write_memory/3` work at the byte
  level. This module provides a higher-level allocator that tracks offsets,
  encodes/decodes typed data, and supports arena-style bulk deallocation.

  ## Why?

  Real Go/Rust WASM modules pass strings and arrays through linear memory.
  Without an allocator you must manually calculate offsets, encode data in the
  correct byte order, and avoid overlapping writes. `Firebird.Memory` handles
  all of that.

  ## Quick Start

      {:ok, wasm} = Firebird.load("string_utils.wasm")
      alloc = Firebird.Memory.new(wasm)

      # Write a string, get back {allocator, offset, byte_length}
      {alloc, ptr, len} = Firebird.Memory.write_string(alloc, "hello world")

      # Pass ptr+len to WASM function that operates on the string
      {:ok, [result_ptr, result_len]} = Firebird.call(wasm, :to_upper, [ptr, len])

      # Read the result back
      {:ok, upper} = Firebird.Memory.read_string(alloc, result_ptr, result_len)
      # => "HELLO WORLD"

      # Free everything at once (arena-style)
      alloc = Firebird.Memory.free_all(alloc)

  ## Typed Helpers

      # i32 arrays
      {alloc, ptr, count} = Firebird.Memory.write_i32_array(alloc, [10, 20, 30])
      {:ok, [10, 20, 30]} = Firebird.Memory.read_i32_array(alloc, ptr, count)

      # f64 arrays
      {alloc, ptr, count} = Firebird.Memory.write_f64_array(alloc, [3.14, 2.718])
      {:ok, [3.14, 2.718]} = Firebird.Memory.read_f64_array(alloc, ptr, count)

  ## Struct Layouts (Go/Rust/C Interop)

  Pass structured data to WASM functions using C-compatible struct layouts:

      layout = Firebird.Memory.define_layout([
        {:x, :f64},
        {:y, :f64},
        {:id, :i32}
      ])

      {alloc, ptr} = Firebird.Memory.write_struct(alloc, layout, %{x: 1.5, y: 2.5, id: 42})
      {:ok, [distance]} = Firebird.call(wasm, :distance, [ptr])

      # Arrays of structs
      points = [%{x: 1.0, y: 2.0, id: 1}, %{x: 3.0, y: 4.0, id: 2}]
      {alloc, ptr, count} = Firebird.Memory.write_struct_array(alloc, layout, points)

  Supported field types: `:i8`, `:u8`, `:i16`, `:u16`, `:i32`, `:u32`,
  `:i64`, `:u64`, `:f32`, `:f64`. Fields use natural alignment matching
  C/Rust/TinyGo defaults.

  ## Alignment

  All allocations are aligned to 8 bytes by default (suitable for f64 and
  most WASM ABIs). You can change this with the `:alignment` option to `new/2`.

  ## Auto-Growing

  If an allocation would exceed the current linear memory size, the allocator
  automatically calls `Firebird.grow_memory/2` to extend it. This is
  transparent — you never need to grow memory manually.
  """

  @default_alignment 8
  # Offset where the bump allocator starts. Reserving the first 1024 bytes
  # avoids clobbering low addresses that some WASM toolchains use for the
  # stack pointer, data segment, or null-guard pages.
  @default_base_offset 1024

  defstruct [
    :instance,
    :alignment,
    :base_offset,
    :cursor,
    :known_memory_size,
    :memory_ref,
    :store_ref,
    allocations: []
  ]

  @typedoc "An allocation record: `{offset, byte_size}`."
  @type allocation :: {non_neg_integer(), non_neg_integer()}

  @type t :: %__MODULE__{
          instance: pid(),
          alignment: pos_integer(),
          base_offset: non_neg_integer(),
          cursor: non_neg_integer(),
          known_memory_size: non_neg_integer() | nil,
          memory_ref: term() | nil,
          store_ref: term() | nil,
          allocations: [allocation()]
        }

  # ── Constructor ──────────────────────────────────────────────────────────

  @doc """
  Create a new allocator bound to a WASM instance.

  ## Options

  - `:alignment` — byte alignment for allocations (default: 8)
  - `:base_offset` — starting offset in linear memory (default: 1024)

  ## Examples

      {:ok, wasm} = Firebird.load("math.wasm")
      alloc = Firebird.Memory.new(wasm)
  """
  @spec new(pid(), keyword()) :: t()
  def new(instance, opts \\ []) do
    alignment = Keyword.get(opts, :alignment, @default_alignment)
    base = Keyword.get(opts, :base_offset, @default_base_offset)
    aligned_base = align_up(base, alignment)

    # Cache the memory and store references upfront. These are stable for
    # the lifetime of the WASM instance and never change. Without caching,
    # every read/write goes through Firebird.read_memory/write_memory →
    # Runtime.memory(pid) which makes 2 GenServer round-trips (Wasmex.memory
    # + Wasmex.store) just to obtain these same references. For batch
    # operations (e.g., 100 string writes), that's 200+ wasted GenServer
    # calls. Caching here eliminates them entirely — all subsequent memory
    # operations use direct NIF calls with zero GenServer overhead.
    {memory_ref, store_ref, known_size} =
      case Firebird.Runtime.memory(instance) do
        {:ok, %{memory: mem, store: store}} ->
          size = Wasmex.Memory.size(store, mem)
          {mem, store, size}

        _ ->
          {nil, nil, nil}
      end

    %__MODULE__{
      instance: instance,
      alignment: alignment,
      base_offset: aligned_base,
      cursor: aligned_base,
      known_memory_size: known_size,
      memory_ref: memory_ref,
      store_ref: store_ref,
      allocations: []
    }
  end

  # ── Raw allocation ───────────────────────────────────────────────────────

  @doc """
  Reserve `byte_size` bytes of linear memory and return the offset.

  Auto-grows memory if needed. The returned offset is aligned to the
  allocator's alignment setting.

  Returns `{updated_allocator, offset}`.

  ## Examples

      {alloc, ptr} = Firebird.Memory.alloc(alloc, 256)
  """
  @spec alloc(t(), non_neg_integer()) :: {t(), non_neg_integer()}
  def alloc(%__MODULE__{} = mem, byte_size) when is_integer(byte_size) and byte_size >= 0 do
    offset = mem.cursor
    new_cursor = align_up(offset + byte_size, mem.alignment)

    # Auto-grow if needed
    mem = ensure_capacity(mem, new_cursor)

    entry = {offset, byte_size}
    mem = %{mem | cursor: new_cursor, allocations: [entry | mem.allocations]}

    {mem, offset}
  end

  @doc """
  Write raw bytes at the current cursor position.

  Returns `{updated_allocator, offset, byte_size}`.
  """
  @spec write_bytes(t(), binary()) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_bytes(%__MODULE__{} = mem, data) when is_binary(data) do
    size = byte_size(data)
    {mem, offset} = alloc(mem, size)
    :ok = write_memory_direct(mem, offset, data)
    {mem, offset, size}
  end

  @doc """
  Read raw bytes from an offset in the WASM instance's memory.

  This is a thin convenience over `Firebird.read_memory/3` that uses
  the allocator's bound instance.
  """
  @spec read_bytes(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, binary()} | {:error, term()}
  def read_bytes(%__MODULE__{} = mem, offset, length) do
    read_memory_direct(mem, offset, length)
  end

  # ── String helpers ───────────────────────────────────────────────────────

  @doc """
  Write a UTF-8 string into linear memory.

  Returns `{updated_allocator, offset, byte_length}`. The byte length is
  what you pass to WASM functions that accept `(ptr, len)` pairs.

  ## Options

  - `:null_terminated` — append a NUL byte (default: false). Useful for
    C-style strings expected by some WASM modules.

  ## Examples

      {alloc, ptr, len} = Firebird.Memory.write_string(alloc, "hello")
      {:ok, [result]} = Firebird.call(wasm, :process_string, [ptr, len])
  """
  @spec write_string(t(), String.t(), keyword()) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_string(%__MODULE__{} = mem, string, opts \\ []) when is_binary(string) do
    null_term = Keyword.get(opts, :null_terminated, false)
    data = if null_term, do: string <> <<0>>, else: string
    # length without NUL terminator
    logical_len = byte_size(string)
    {mem, offset, _raw_size} = write_bytes(mem, data)
    {mem, offset, logical_len}
  end

  @doc """
  Read a UTF-8 string from linear memory.

  ## Examples

      {:ok, str} = Firebird.Memory.read_string(alloc, ptr, len)
  """
  @spec read_string(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, String.t()} | {:error, term()}
  def read_string(%__MODULE__{} = mem, offset, length) do
    case read_bytes(mem, offset, length) do
      {:ok, bytes} ->
        if String.valid?(bytes) do
          {:ok, bytes}
        else
          {:error, :invalid_utf8}
        end

      error ->
        error
    end
  end

  # ── i32 array helpers ────────────────────────────────────────────────────

  @doc """
  Write a list of 32-bit signed integers into linear memory (little-endian).

  Returns `{updated_allocator, offset, element_count}`.

  ## Examples

      {alloc, ptr, count} = Firebird.Memory.write_i32_array(alloc, [1, 2, 3])
      {:ok, [result]} = Firebird.call(wasm, :sum_array, [ptr, count])
  """
  @spec write_i32_array(t(), [integer()]) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_i32_array(%__MODULE__{} = mem, values) when is_list(values) do
    data = encode_array(values, fn v -> <<v::little-signed-32>> end)
    {mem, offset, _size} = write_bytes(mem, data)
    {mem, offset, length(values)}
  end

  @doc """
  Read a list of 32-bit signed integers from linear memory (little-endian).

  ## Examples

      {:ok, [1, 2, 3]} = Firebird.Memory.read_i32_array(alloc, ptr, 3)
  """
  @spec read_i32_array(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [integer()]} | {:error, term()}
  def read_i32_array(%__MODULE__{} = mem, offset, count) when is_integer(count) and count >= 0 do
    case read_bytes(mem, offset, count * 4) do
      {:ok, bytes} ->
        values = for <<v::little-signed-32 <- bytes>>, do: v
        {:ok, values}

      error ->
        error
    end
  end

  # ── u32 array helpers ────────────────────────────────────────────────────

  @doc """
  Write a list of 32-bit unsigned integers into linear memory (little-endian).

  Returns `{updated_allocator, offset, element_count}`.
  """
  @spec write_u32_array(t(), [non_neg_integer()]) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_u32_array(%__MODULE__{} = mem, values) when is_list(values) do
    data = encode_array(values, fn v -> <<v::little-unsigned-32>> end)
    {mem, offset, _size} = write_bytes(mem, data)
    {mem, offset, length(values)}
  end

  @doc """
  Read a list of 32-bit unsigned integers from linear memory (little-endian).
  """
  @spec read_u32_array(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [non_neg_integer()]} | {:error, term()}
  def read_u32_array(%__MODULE__{} = mem, offset, count) when is_integer(count) and count >= 0 do
    case read_bytes(mem, offset, count * 4) do
      {:ok, bytes} ->
        values = for <<v::little-unsigned-32 <- bytes>>, do: v
        {:ok, values}

      error ->
        error
    end
  end

  # ── i64 array helpers ────────────────────────────────────────────────────

  @doc """
  Write a list of 64-bit signed integers into linear memory (little-endian).

  Returns `{updated_allocator, offset, element_count}`.
  """
  @spec write_i64_array(t(), [integer()]) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_i64_array(%__MODULE__{} = mem, values) when is_list(values) do
    data = encode_array(values, fn v -> <<v::little-signed-64>> end)
    {mem, offset, _size} = write_bytes(mem, data)
    {mem, offset, length(values)}
  end

  @doc """
  Read a list of 64-bit signed integers from linear memory (little-endian).
  """
  @spec read_i64_array(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [integer()]} | {:error, term()}
  def read_i64_array(%__MODULE__{} = mem, offset, count) when is_integer(count) and count >= 0 do
    case read_bytes(mem, offset, count * 8) do
      {:ok, bytes} ->
        values = for <<v::little-signed-64 <- bytes>>, do: v
        {:ok, values}

      error ->
        error
    end
  end

  # ── f32 array helpers ────────────────────────────────────────────────────

  @doc """
  Write a list of 32-bit floats into linear memory (little-endian IEEE 754).

  Returns `{updated_allocator, offset, element_count}`.
  """
  @spec write_f32_array(t(), [float()]) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_f32_array(%__MODULE__{} = mem, values) when is_list(values) do
    data = encode_array(values, fn v -> <<v::little-float-32>> end)
    {mem, offset, _size} = write_bytes(mem, data)
    {mem, offset, length(values)}
  end

  @doc """
  Read a list of 32-bit floats from linear memory (little-endian IEEE 754).
  """
  @spec read_f32_array(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [float()]} | {:error, term()}
  def read_f32_array(%__MODULE__{} = mem, offset, count) when is_integer(count) and count >= 0 do
    case read_bytes(mem, offset, count * 4) do
      {:ok, bytes} ->
        values = for <<v::little-float-32 <- bytes>>, do: v
        {:ok, values}

      error ->
        error
    end
  end

  # ── f64 array helpers ────────────────────────────────────────────────────

  @doc """
  Write a list of 64-bit floats into linear memory (little-endian IEEE 754).

  Returns `{updated_allocator, offset, element_count}`.

  ## Examples

      {alloc, ptr, count} = Firebird.Memory.write_f64_array(alloc, [3.14, 2.718])
  """
  @spec write_f64_array(t(), [float()]) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_f64_array(%__MODULE__{} = mem, values) when is_list(values) do
    data = encode_array(values, fn v -> <<v::little-float-64>> end)
    {mem, offset, _size} = write_bytes(mem, data)
    {mem, offset, length(values)}
  end

  @doc """
  Read a list of 64-bit floats from linear memory (little-endian IEEE 754).

  ## Examples

      {:ok, [3.14, 2.718]} = Firebird.Memory.read_f64_array(alloc, ptr, 2)
  """
  @spec read_f64_array(t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [float()]} | {:error, term()}
  def read_f64_array(%__MODULE__{} = mem, offset, count) when is_integer(count) and count >= 0 do
    case read_bytes(mem, offset, count * 8) do
      {:ok, bytes} ->
        values = for <<v::little-float-64 <- bytes>>, do: v
        {:ok, values}

      error ->
        error
    end
  end

  # ── Scalar helpers ───────────────────────────────────────────────────────

  @doc """
  Write a single i32 value and return its offset.

  Returns `{updated_allocator, offset}`.
  """
  @spec write_i32(t(), integer()) :: {t(), non_neg_integer()}
  def write_i32(%__MODULE__{} = mem, value) when is_integer(value) do
    {mem, offset, _size} = write_bytes(mem, <<value::little-signed-32>>)
    {mem, offset}
  end

  @doc "Read a single i32 from the given offset."
  @spec read_i32(t(), non_neg_integer()) :: {:ok, integer()} | {:error, term()}
  def read_i32(%__MODULE__{} = mem, offset) do
    case read_bytes(mem, offset, 4) do
      {:ok, <<v::little-signed-32>>} -> {:ok, v}
      error -> error
    end
  end

  @doc """
  Write a single i64 value and return its offset.

  Returns `{updated_allocator, offset}`.
  """
  @spec write_i64(t(), integer()) :: {t(), non_neg_integer()}
  def write_i64(%__MODULE__{} = mem, value) when is_integer(value) do
    {mem, offset, _size} = write_bytes(mem, <<value::little-signed-64>>)
    {mem, offset}
  end

  @doc "Read a single i64 from the given offset."
  @spec read_i64(t(), non_neg_integer()) :: {:ok, integer()} | {:error, term()}
  def read_i64(%__MODULE__{} = mem, offset) do
    case read_bytes(mem, offset, 8) do
      {:ok, <<v::little-signed-64>>} -> {:ok, v}
      error -> error
    end
  end

  @doc """
  Write a single f64 value and return its offset.

  Returns `{updated_allocator, offset}`.
  """
  @spec write_f64(t(), float()) :: {t(), non_neg_integer()}
  def write_f64(%__MODULE__{} = mem, value) when is_number(value) do
    float_val = value * 1.0
    {mem, offset, _size} = write_bytes(mem, <<float_val::little-float-64>>)
    {mem, offset}
  end

  @doc "Read a single f64 from the given offset."
  @spec read_f64(t(), non_neg_integer()) :: {:ok, float()} | {:error, term()}
  def read_f64(%__MODULE__{} = mem, offset) do
    case read_bytes(mem, offset, 8) do
      {:ok, <<v::little-float-64>>} -> {:ok, v}
      error -> error
    end
  end

  # ── Arena-style deallocation ─────────────────────────────────────────────

  @doc """
  Reset the allocator to its initial state (arena-style free-all).

  All previous offsets become invalid. The underlying WASM memory is NOT
  zeroed — it simply becomes available for reuse. This is the fastest way
  to reclaim memory between batches of work.

  ## Examples

      alloc = Firebird.Memory.free_all(alloc)
  """
  @spec free_all(t()) :: t()
  def free_all(%__MODULE__{} = mem) do
    %{mem | cursor: mem.base_offset, allocations: []}
  end

  @doc """
  Reset and zero the allocated region (secure free).

  Like `free_all/1` but also writes zeros over the previously used region.
  Use this when the memory may have contained sensitive data.
  """
  @spec free_all_secure(t()) :: t()
  def free_all_secure(%__MODULE__{} = mem) do
    used = mem.cursor - mem.base_offset

    if used > 0 do
      zeros = :binary.copy(<<0>>, used)
      write_memory_direct(mem, mem.base_offset, zeros)
    end

    free_all(mem)
  end

  # ── Introspection ────────────────────────────────────────────────────────

  @doc """
  Return the total number of bytes currently allocated (including alignment padding).
  """
  @spec allocated_bytes(t()) :: non_neg_integer()
  def allocated_bytes(%__MODULE__{} = mem) do
    mem.cursor - mem.base_offset
  end

  @doc """
  Return the number of individual allocations tracked.
  """
  @spec allocation_count(t()) :: non_neg_integer()
  def allocation_count(%__MODULE__{} = mem) do
    length(mem.allocations)
  end

  @doc """
  Return the list of allocations as `[{offset, byte_size}]` (most recent first).
  """
  @spec allocations(t()) :: [allocation()]
  def allocations(%__MODULE__{} = mem), do: mem.allocations

  @doc """
  Return a summary map of the allocator's state.

  ## Examples

      info = Firebird.Memory.info(alloc)
      # => %{
      #   cursor: 1080,
      #   base_offset: 1024,
      #   allocated_bytes: 56,
      #   allocation_count: 3,
      #   alignment: 8,
      #   wasm_memory_bytes: 65536
      # }
  """
  @spec info(t()) :: map()
  def info(%__MODULE__{} = mem) do
    # Use cached size when available to avoid NIF calls; fall back to live query.
    wasm_mem =
      mem.known_memory_size ||
        case Firebird.memory_size(mem.instance) do
          {:ok, size} -> size
          _ -> nil
        end

    %{
      cursor: mem.cursor,
      base_offset: mem.base_offset,
      allocated_bytes: allocated_bytes(mem),
      allocation_count: allocation_count(mem),
      alignment: mem.alignment,
      wasm_memory_bytes: wasm_mem
    }
  end

  # ── Convenience: scoped allocation ───────────────────────────────────────

  @doc """
  Execute a function with temporary allocations, then free everything.

  The allocator is reset after the function returns (arena pattern).
  Returns `{result, reset_allocator}`.

  ## Examples

      {result, alloc} = Firebird.Memory.with_arena(alloc, fn alloc ->
        {alloc, ptr, len} = Firebird.Memory.write_string(alloc, "temp data")
        {:ok, [hash]} = Firebird.call(wasm, :hash, [ptr, len])
        {hash, alloc}
      end)
      # alloc is reset, but `result` (the hash) is preserved
  """
  @spec with_arena(t(), (t() -> {term(), t()})) :: {term(), t()}
  def with_arena(%__MODULE__{} = mem, fun) when is_function(fun, 1) do
    saved_cursor = mem.cursor
    saved_allocs = mem.allocations

    {result, mem} = fun.(mem)

    reset_mem = %{mem | cursor: saved_cursor, allocations: saved_allocs}
    {result, reset_mem}
  end

  # ── Struct layout helpers ──────────────────────────────────────────────

  @typedoc """
  A struct layout describes a sequence of typed fields for packing/unpacking
  binary data to/from WASM linear memory.

  Each field is `{name, type}` where type is one of:
  `:i8`, `:u8`, `:i16`, `:u16`, `:i32`, `:u32`, `:i64`, `:u64`, `:f32`, `:f64`

  Fields are laid out sequentially with natural alignment (each field aligned
  to its own size). This matches the default struct packing used by C, Rust,
  and TinyGo WASM targets.
  """
  @type field_type :: :i8 | :u8 | :i16 | :u16 | :i32 | :u32 | :i64 | :u64 | :f32 | :f64
  @type struct_layout :: [{atom(), field_type()}]

  @doc """
  Define a struct layout and compute its total byte size and field offsets.

  Returns a map with `:fields`, `:size`, and `:offsets` (a map of field name
  to byte offset within the struct). Fields use natural alignment (each field
  aligned to its own size), matching C/Rust/TinyGo defaults.

  ## Examples

      layout = Firebird.Memory.define_layout([
        {:x, :f64},
        {:y, :f64},
        {:id, :i32},
        {:flags, :u8}
      ])
      # => %{
      #   fields: [{:x, :f64}, {:y, :f64}, {:id, :i32}, {:flags, :u8}],
      #   offsets: %{x: 0, y: 8, id: 16, flags: 20},
      #   size: 24  # padded to max alignment (8)
      # }
  """
  @spec define_layout(struct_layout()) :: map()
  def define_layout(fields) when is_list(fields) do
    {offsets_map, cursor, max_align} =
      Enum.reduce(fields, {%{}, 0, 1}, fn {name, type}, {offsets, cur, max_a} ->
        size = type_size(type)
        alignment = size
        aligned_cur = align_up(cur, alignment)
        {Map.put(offsets, name, aligned_cur), aligned_cur + size, max(max_a, alignment)}
      end)

    # Total struct size padded to max alignment (for arrays of structs)
    total_size = align_up(cursor, max_align)

    %{fields: fields, offsets: offsets_map, size: total_size}
  end

  @doc """
  Write a struct (map of field values) into linear memory using a layout.

  Returns `{updated_allocator, offset}` where offset points to the start
  of the struct in WASM memory.

  ## Examples

      layout = Firebird.Memory.define_layout([{:x, :f64}, {:y, :f64}, {:id, :i32}])
      point = %{x: 1.5, y: 2.5, id: 42}
      {alloc, ptr} = Firebird.Memory.write_struct(alloc, layout, point)

      # Pass ptr to WASM function
      {:ok, [distance]} = Firebird.call(wasm, :distance_from_origin, [ptr])
  """
  @spec write_struct(t(), map(), map()) :: {t(), non_neg_integer()}
  def write_struct(%__MODULE__{} = mem, layout, values) when is_map(values) do
    %{size: struct_size} = layout
    {mem, base} = alloc(mem, struct_size)
    binary = encode_struct_binary(layout, values)
    :ok = write_memory_direct(mem, base, binary)
    {mem, base}
  end

  @doc """
  Read a struct from linear memory using a layout.

  Returns `{:ok, map}` with field names as keys.

  ## Examples

      layout = Firebird.Memory.define_layout([{:x, :f64}, {:y, :f64}, {:id, :i32}])
      {:ok, point} = Firebird.Memory.read_struct(alloc, layout, ptr)
      # => {:ok, %{x: 1.5, y: 2.5, id: 42}}
  """
  @spec read_struct(t(), map(), non_neg_integer()) :: {:ok, map()} | {:error, term()}
  def read_struct(%__MODULE__{} = mem, layout, offset) do
    %{fields: fields, offsets: offsets, size: struct_size} = layout

    # Read the entire struct in one NIF call, then decode fields from the binary.
    case read_bytes(mem, offset, struct_size) do
      {:ok, buf} ->
        result =
          Enum.reduce(fields, %{}, fn {name, type}, acc ->
            field_offset = Map.fetch!(offsets, name)
            field_size = type_size(type)

            <<_::binary-size(field_offset), field_bytes::binary-size(field_size), _::binary>> =
              buf

            Map.put(acc, name, decode_field(type, field_bytes))
          end)

        {:ok, result}

      {:error, _} = err ->
        err
    end
  end

  @doc """
  Write an array of structs into linear memory.

  Returns `{updated_allocator, offset, count}`.

  ## Examples

      layout = Firebird.Memory.define_layout([{:x, :f64}, {:y, :f64}])
      points = [%{x: 1.0, y: 2.0}, %{x: 3.0, y: 4.0}]
      {alloc, ptr, count} = Firebird.Memory.write_struct_array(alloc, layout, points)
  """
  @spec write_struct_array(t(), map(), [map()]) :: {t(), non_neg_integer(), non_neg_integer()}
  def write_struct_array(%__MODULE__{} = mem, layout, structs) when is_list(structs) do
    %{size: struct_size} = layout
    count = length(structs)
    total_bytes = struct_size * count
    {mem, base} = alloc(mem, total_bytes)

    # Build the entire array as a single binary, then write once.
    # This replaces count × fields_per_struct individual NIF calls with ONE.
    binary =
      structs
      |> Enum.map(&encode_struct_binary(layout, &1))
      |> IO.iodata_to_binary()

    :ok = write_memory_direct(mem, base, binary)
    {mem, base, count}
  end

  @doc """
  Read an array of structs from linear memory.

  Returns `{:ok, [map()]}`.

  ## Examples

      layout = Firebird.Memory.define_layout([{:x, :f64}, {:y, :f64}])
      {:ok, points} = Firebird.Memory.read_struct_array(alloc, layout, ptr, 2)
      # => {:ok, [%{x: 1.0, y: 2.0}, %{x: 3.0, y: 4.0}]}
  """
  @spec read_struct_array(t(), map(), non_neg_integer(), non_neg_integer()) ::
          {:ok, [map()]} | {:error, term()}
  def read_struct_array(%__MODULE__{} = mem, layout, offset, count)
      when is_integer(count) and count >= 0 do
    if count == 0 do
      {:ok, []}
    else
      %{fields: fields, offsets: offsets, size: struct_size} = layout
      total_bytes = struct_size * count

      # Read the entire array in one NIF call, then decode all structs from the binary.
      case read_bytes(mem, offset, total_bytes) do
        {:ok, buf} ->
          structs =
            for idx <- 0..(count - 1) do
              # Slice this struct's chunk first (O(1) sub-binary), then decode
              # fields from the small chunk instead of the full buffer.
              struct_start = idx * struct_size
              <<_::binary-size(struct_start), chunk::binary-size(struct_size), _::binary>> = buf

              Enum.reduce(fields, %{}, fn {name, type}, acc ->
                field_offset = Map.fetch!(offsets, name)
                field_size = type_size(type)

                <<_::binary-size(field_offset), field_bytes::binary-size(field_size), _::binary>> =
                  chunk

                Map.put(acc, name, decode_field(type, field_bytes))
              end)
            end

          {:ok, structs}

        {:error, _} = err ->
          err
      end
    end
  end

  # ── Struct binary encoding (batched — single NIF call) ────────────────

  # Build a complete struct as a single binary, with fields placed at their
  # correct offsets and alignment padding zeroed. This enables write_struct
  # and write_struct_array to cross the NIF boundary exactly once instead of
  # once per field (or once per field per struct for arrays).
  #
  # Uses an iodata builder that emits zero-padding between fields and
  # converts to a flat binary once at the end. This avoids the previous
  # approach of splicing into a pre-allocated binary (which copied the
  # entire buffer on every field insertion — O(fields × struct_size) bytes
  # copied total). The iodata approach copies O(struct_size) bytes total.
  @spec encode_struct_binary(map(), map()) :: binary()
  defp encode_struct_binary(%{fields: fields, offsets: offsets, size: struct_size}, values) do
    # Sort fields by offset so we can emit padding + field bytes sequentially.
    sorted =
      fields
      |> Enum.map(fn {name, type} ->
        {Map.fetch!(offsets, name), name, type}
      end)
      |> Enum.sort_by(&elem(&1, 0))

    # Walk sorted fields, emitting zero-padding gaps and encoded field bytes.
    {iodata, cursor} =
      Enum.reduce(sorted, {[], 0}, fn {field_offset, name, type}, {acc, cur} ->
        value = Map.fetch!(values, name)
        field_bytes = encode_field(type, value)
        gap = field_offset - cur
        padding = if gap > 0, do: :binary.copy(<<0>>, gap), else: <<>>
        new_cur = field_offset + byte_size(field_bytes)
        {[acc, padding, field_bytes], new_cur}
      end)

    # Pad to full struct size (for array-of-structs alignment)
    trailing = struct_size - cursor

    final =
      if trailing > 0,
        do: [iodata, :binary.copy(<<0>>, trailing)],
        else: iodata

    IO.iodata_to_binary(final)
  end

  # ── Struct field encoding/decoding ───────────────────────────────────────

  defp type_size(:i8), do: 1
  defp type_size(:u8), do: 1
  defp type_size(:i16), do: 2
  defp type_size(:u16), do: 2
  defp type_size(:i32), do: 4
  defp type_size(:u32), do: 4
  defp type_size(:i64), do: 8
  defp type_size(:u64), do: 8
  defp type_size(:f32), do: 4
  defp type_size(:f64), do: 8

  defp encode_field(:i8, v), do: <<v::little-signed-8>>
  defp encode_field(:u8, v), do: <<v::little-unsigned-8>>
  defp encode_field(:i16, v), do: <<v::little-signed-16>>
  defp encode_field(:u16, v), do: <<v::little-unsigned-16>>
  defp encode_field(:i32, v), do: <<v::little-signed-32>>
  defp encode_field(:u32, v), do: <<v::little-unsigned-32>>
  defp encode_field(:i64, v), do: <<v::little-signed-64>>
  defp encode_field(:u64, v), do: <<v::little-unsigned-64>>
  defp encode_field(:f32, v), do: <<v * 1.0::little-float-32>>
  defp encode_field(:f64, v), do: <<v * 1.0::little-float-64>>

  defp decode_field(:i8, <<v::little-signed-8>>), do: v
  defp decode_field(:u8, <<v::little-unsigned-8>>), do: v
  defp decode_field(:i16, <<v::little-signed-16>>), do: v
  defp decode_field(:u16, <<v::little-unsigned-16>>), do: v
  defp decode_field(:i32, <<v::little-signed-32>>), do: v
  defp decode_field(:u32, <<v::little-unsigned-32>>), do: v
  defp decode_field(:i64, <<v::little-signed-64>>), do: v
  defp decode_field(:u64, <<v::little-unsigned-64>>), do: v
  defp decode_field(:f32, <<v::little-float-32>>), do: v
  defp decode_field(:f64, <<v::little-float-64>>), do: v

  # ── Private ──────────────────────────────────────────────────────────────

  # Encode a list of values into a flat binary via iodata.
  #
  # The previous `for v <- values, into: <<>>, do: ...` pattern builds the
  # binary by repeated append — each iteration copies the accumulated bytes
  # so far, making it O(n²) for large arrays (e.g., 10k elements).
  #
  # This helper maps each value to a small binary, producing an iodata list,
  # then concatenates once with IO.iodata_to_binary/1 — O(n) total.
  @spec encode_array([term()], (term() -> binary())) :: binary()
  defp encode_array(values, encoder) do
    values |> Enum.map(encoder) |> IO.iodata_to_binary()
  end

  defp align_up(offset, alignment) do
    case rem(offset, alignment) do
      0 -> offset
      r -> offset + (alignment - r)
    end
  end

  # Direct memory read using cached refs — bypasses 2 GenServer round-trips
  # (Wasmex.memory + Wasmex.store) that Firebird.read_memory would make.
  # Falls back to the public API if refs aren't cached.
  defp read_memory_direct(%__MODULE__{memory_ref: mref, store_ref: sref}, offset, length)
       when mref != nil and sref != nil do
    {:ok, Wasmex.Memory.read_binary(sref, mref, offset, length)}
  end

  defp read_memory_direct(%__MODULE__{instance: instance}, offset, length) do
    Firebird.read_memory(instance, offset, length)
  end

  # Direct memory write using cached refs — same optimization as read.
  defp write_memory_direct(%__MODULE__{memory_ref: mref, store_ref: sref}, offset, data)
       when mref != nil and sref != nil do
    Wasmex.Memory.write_binary(sref, mref, offset, data)
    :ok
  end

  defp write_memory_direct(%__MODULE__{instance: instance}, offset, data) do
    Firebird.write_memory(instance, offset, data)
  end

  # Ensure the WASM linear memory is large enough for `needed_end` bytes.
  #
  # Uses the cached `known_memory_size` to avoid NIF round-trips when the
  # allocation fits within the already-known capacity. This eliminates 3 NIF
  # calls per allocation (memory + store + size) in the common case, which
  # adds up fast for batch operations (e.g., 100 string writes saves ~300
  # NIF calls).
  #
  # When growth IS needed, we query the actual size (in case another process
  # grew memory), grow, and update the cache.
  defp ensure_capacity(%__MODULE__{known_memory_size: known} = mem, needed_end)
       when is_integer(known) and needed_end <= known do
    # Fast path: allocation fits within cached capacity — zero NIF calls.
    mem
  end

  defp ensure_capacity(%__MODULE__{memory_ref: mref, store_ref: sref} = mem, needed_end)
       when mref != nil and sref != nil do
    # Use cached refs for direct NIF calls — no GenServer round-trips.
    current_size = Wasmex.Memory.size(sref, mref)

    if needed_end > current_size do
      deficit = needed_end - current_size
      pages = div(deficit + 65535, 65536)
      Wasmex.Memory.grow(sref, mref, pages)
      new_size = current_size + pages * 65536
      %{mem | known_memory_size: new_size}
    else
      %{mem | known_memory_size: current_size}
    end
  end

  defp ensure_capacity(%__MODULE__{} = mem, needed_end) do
    # Fallback: no cached refs — go through the public API (GenServer path).
    case Firebird.memory_size(mem.instance) do
      {:ok, current_size} when needed_end > current_size ->
        deficit = needed_end - current_size
        pages = div(deficit + 65535, 65536)
        Firebird.grow_memory(mem.instance, pages)
        new_size = current_size + pages * 65536
        %{mem | known_memory_size: new_size}

      {:ok, current_size} ->
        %{mem | known_memory_size: current_size}

      _ ->
        mem
    end
  end
end