# SPDX-License-Identifier: LSL-1.0
# Copyright (c) 2026 Markus Maiwald
# Stewardship: Self Sovereign Society Foundation
#
# This file is part of the Nexus Sovereign Core.
# See legal/LICENSE_SOVEREIGN.md for license terms.

## Rumpk Layer 1: ION Slab Allocator

# 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

  SYSTABLE_BASE = 0x83000000'u64
  USER_SLAB_OFFSET = 0x10000'u64   # Offset within SYSTABLE
  USER_SLAB_BASE* = SYSTABLE_BASE + USER_SLAB_OFFSET  # 0x83010000
  USER_SLAB_COUNT = 512           # 512 packets to cover RX Ring (256) + TX
  USER_PKT_SIZE = 2048            # 2KB per packet
  USER_BITMAP_ADDR = SYSTABLE_BASE + 0x100

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_tx_ring*: RingBuffer[IonPacket, 256]
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

  # Tracing for Phase 37
  proc kprint_hex(v: uint64) {.importc, cdecl.}
  dbg("[ION] Pool Base Phys: ")
  kprint_hex(global_pool.base_phys)
  dbg("")

  dbg("[ION] Ring Init...")
  global_pool.free_ring.init()
  global_tx_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

  if (pkt.id and 0x8000) != 0:
    # User Slab - Clear shared bitmap
    let slotIdx = pkt.id and 0x7FFF
    if slotIdx >= USER_SLAB_COUNT: return
    let bitmap = cast[ptr array[16, byte]](USER_BITMAP_ADDR)
    let byteIdx = int(slotIdx) div 8
    let bitIdx = int(slotIdx) mod 8
    let mask = byte(1 shl bitIdx)
    bitmap[byteIdx] = bitmap[byteIdx] and (not mask)
    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, cdecl.} =
  let pkt = ion_alloc()
  if pkt.data == nil: return 0
  out_id[] = pkt.id
  return pkt.phys

proc ion_free_raw*(id: uint16) {.exportc, cdecl.} =
  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.} =
  if (id and 0x8000) != 0:
    let idx = id and 0x7FFF
    let offset = int(idx) * SLAB_SIZE
    return cast[ptr byte](USER_SLAB_BASE + uint64(offset))
  let offset = int(id) * SLAB_SIZE
  return addr global_pool.buffer[offset]

proc ion_get_phys*(id: uint16): uint64 {.exportc.} =
  if (id and 0x8000) != 0:
    let idx = id and 0x7FFF
    let offset = int(idx) * SLAB_SIZE
    return USER_SLAB_BASE + uint64(offset)
  let offset = int(id) * SLAB_SIZE
  return global_pool.base_phys + uint64(offset)

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


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

proc ion_tx_push*(pkt: IonPacket): bool {.exportc.} =
  if global_tx_ring.push(pkt):
     # dbg("[ION TX] Pushed")
     return true
  dbg("[ION TX] PUSH FAILED (Global Ring Full)")
  return false

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
  dbg("[ION TX] Popped Packet for VirtIO")
  return true

# =========================================================
# User-Visible Slab Allocator (Shared Memory)
# =========================================================
# NOTE: This allocator provides buffers in the SYSTABLE shared region
# (0x83010000+) which is mapped into both kernel and userland page tables.
# Used for network packet egress from userland.

# NOTE: Constants moved to top


# var user_slab_bitmap: array[USER_SLAB_COUNT, bool] # REMOVED: Use Shared Bitmap

proc ion_user_slab_init*() {.exportc.} =
  ## Initialize shared user slab bitmap (all free)
  let bitmap = cast[ptr array[64, byte]](USER_BITMAP_ADDR)
  for i in 0 ..< 64:
    bitmap[i] = 0

proc ion_alloc_shared*(out_id: ptr uint16): uint64 {.exportc, cdecl.} =
  ## Allocate a buffer from the user-visible slab (Kernel Side, Shared Bitmap)
  let bitmap = cast[ptr array[64, byte]](USER_BITMAP_ADDR)
  
  for byteIdx in 0 ..< 64:
    if bitmap[byteIdx] != 0xFF:
      for bitIdx in 0 ..< 8:
        let mask = byte(1 shl bitIdx)
        if (bitmap[byteIdx] and mask) == 0:
          # Found free
          bitmap[byteIdx] = bitmap[byteIdx] or mask
          let idx = byteIdx * 8 + bitIdx
          if idx >= USER_SLAB_COUNT: return 0
          
          out_id[] = uint16(idx) or 0x8000
          return USER_SLAB_BASE + uint64(idx) * USER_PKT_SIZE
  return 0