# ION Memory Manager
# The "Slab Allocator" for Zero-Copy IO

import ../ring

proc console_write(p: pointer, len: csize_t) {.importc, cdecl.}
proc dbg(s: string) =
  console_write(unsafeAddr s[0], csize_t(s.len))
  var nl = "\n"
  console_write(unsafeAddr nl[0], csize_t(1))

const
  SLAB_SIZE* = 2048  # Max Packet Size (Ethernet Frame + Headroom)
  POOL_COUNT* = 1024 # Number of packets in the pool (2MB total RAM)
  POOL_ALIGN* = 4096 # VirtIO/Page Alignment

type
  # The Physical Token representing a packet
  IonPacket* = object
    data*: ptr UncheckedArray[byte] # Virtual Address of payload
    phys*: uint64                   # Physical Address (For VirtIO/DMA)
    len*: uint16                    # Actual data length
    id*: uint16                     # The Token ID (Slab Index)

                                    # The Monolithic Memory Pool
  PacketPool = object
    # We use 'align' to ensure the buffer starts on a page boundary for VirtIO
    buffer {.align: POOL_ALIGN.}: array[SLAB_SIZE * POOL_COUNT, byte]
    free_ring: RingBuffer[uint16, POOL_COUNT] # Stores IDs of free slabs
    base_phys: uint64

var global_pool: PacketPool

proc ion_pool_init*() {.exportc.} =
  ## Initialize the DMA Pool with REAL Physical Address
  dbg("[ION] Initializing Pool...")

  # 1. Get the VIRTUAL address of the static buffer
  let virt_addr = cast[uint64](addr global_pool.buffer[0])

  # 2. Translate to PHYSICAL (Identity Mapped for Phase 7)
  global_pool.base_phys = virt_addr

  dbg("[ION] Ring Init...")
  global_pool.free_ring.init()

  # Fill the free ring with all indices [0..1023]
  dbg("[ION] Filling Slabs...")
  var count = 0
  for i in 0 ..< POOL_COUNT:
    if global_pool.free_ring.push(uint16(i)):
      inc count

  dbg("[ION] Pool Ready.")

proc ion_alloc*(): IonPacket {.exportc.} =
  ## O(1) Allocation. Returns an empty packet struct.
  ## If OOM, returns packet with data = nil
  var pkt: IonPacket

  let (ok, idx) = global_pool.free_ring.pop()
  if not ok:
    dbg("[ION] ALLOC FAILED (Empty Ring?)")
    pkt.data = nil
    return pkt

  pkt.id = idx
  pkt.len = 0

  # Calculate Virtual Address
  let offset = int(idx) * SLAB_SIZE
  pkt.data = cast[ptr UncheckedArray[byte]](addr global_pool.buffer[offset])

  # Calculate Physical Address (1:1 map for Phase 7)
  pkt.phys = global_pool.base_phys + uint64(offset)

  return pkt

proc ion_free*(pkt: IonPacket) {.exportc.} =
  ## O(1) Free. Returns the token to the ring.
  if pkt.data == nil: return

  discard global_pool.free_ring.push(pkt.id)

# Helper for C/Zig Interop (Pure Pointers)
# Return physical address of a allocated block, put ID in out_id
proc ion_alloc_raw*(out_id: ptr uint16): uint64 {.exportc.} =
  let pkt = ion_alloc()
  if pkt.data == nil: return 0
  out_id[] = pkt.id
  return pkt.phys

proc ion_free_raw*(id: uint16) {.exportc.} =
  var pkt: IonPacket
  pkt.id = id
  # We don't reconstruct data/phys for free, just push ID
  # But for safety we might check bounds? Ring handles it.
  pkt.data = cast[ptr UncheckedArray[byte]](1) # Dummy non-nil
  ion_free(pkt)

proc ion_get_virt*(id: uint16): ptr byte {.exportc.} =
  let offset = int(id) * SLAB_SIZE
  return addr global_pool.buffer[offset]

proc ion_get_phys*(id: uint16): uint64 {.exportc.} =
  let offset = int(id) * SLAB_SIZE
  return global_pool.base_phys + uint64(offset)

# =========================================================
# The Global TX Ring (Multiplexing)
# =========================================================

var global_tx_ring*: RingBuffer[IonPacket, 256]

proc ion_tx_init*() {.exportc.} =
  global_tx_ring.init()

proc ion_tx_push*(pkt: IonPacket): bool {.exportc.} =
  global_tx_ring.push(pkt)

proc ion_tx_pop*(out_id: ptr uint16, out_len: ptr uint16): bool {.exportc.} =
  if global_tx_ring.isEmpty:
    return false

  let (ok, pkt) = global_tx_ring.pop()
  if not ok: return false

  out_id[] = pkt.id
  out_len[] = pkt.len
  return true