# libs/membrane/net_glue.nim

import ../../core/ion/memory
import ion_client

proc console_write*(buf: cstring, len: csize_t) {.importc, cdecl.}
# Define LwIP Raw API types (Internal/External)
# We need to bridge to the C headers of LwIP
type
  TcpPcb* {.importc: "struct tcp_pcb", header: "lwip/tcp.h", pure.} = object
  IpAddr* {.importc: "ip_addr_t", header: "lwip/ip_addr.h", pure.} = object
  Pbuf* {.importc: "struct pbuf", header: "lwip/pbuf.h", pure.} = object
    next*: ptr Pbuf
    payload*: pointer
    tot_len*: uint16
    len*: uint16
  Netif* {.importc: "struct netif", header: "lwip/netif.h", pure.} = object
    input*: proc(p: ptr Pbuf, ni: ptr Netif): ErrT {.cdecl.}
    linkoutput*: proc(ni: ptr Netif, p: ptr Pbuf): ErrT {.cdecl.}
    mtu*: uint16
    flags*: uint8
    hwaddr_len*: uint8
    hwaddr*: array[6, byte]
  ErrT* = int8

const
  ERR_OK* = 0
  PBUF_RAW* = 0
  PBUF_POOL* = 3
  ERR_MEM* = -1
  ERR_TIMEOUT* = -3
  ERR_ABRT* = -4
  ERR_RST* = -5
  ERR_CLSD* = -6
  ERR_CONN* = -7
  ERR_VAL* = -8
  ERR_ARG* = -9
  ERR_USE* = -10
  ERR_IF* = -11
  ERR_ISCONN* = -12
  ERR_INPROGRESS* = -13

# External LwIP Procs
proc tcp_new*(): ptr TcpPcb {.importc: "tcp_new", header: "lwip/tcp.h".}
proc tcp_arg*(pcb: ptr TcpPcb; arg: pointer) {.importc: "tcp_arg",
    header: "lwip/tcp.h".}
proc tcp_connect*(pcb: ptr TcpPcb; ip: ptr IpAddr; port: uint16;
    cb: pointer): ErrT {.importc: "tcp_connect", header: "lwip/tcp.h".}
proc sys_check_timeouts*() {.importc: "sys_check_timeouts",
    header: "lwip/timeouts.h".}

proc pbuf_alloc*(layer, length, pType: int): ptr Pbuf {.importc: "pbuf_alloc",
    header: "lwip/pbuf.h".}
proc pbuf_free*(p: ptr Pbuf): uint8 {.importc: "pbuf_free",
    header: "lwip/pbuf.h".}
proc pbuf_take*(p: ptr Pbuf; dataptr: pointer;
    len: uint16): ErrT {.importc: "pbuf_take", header: "lwip/pbuf.h".}

var membrane_netif*: Netif

proc pbuf_copy_partial*(p: ptr Pbuf; dataptr: pointer; len,
    offset: uint16): uint16 {.
    importc: "pbuf_copy_partial", header: "lwip/pbuf.h".}

# The "Lungs" Logic
proc membrane_output*(ni: ptr Netif; p: ptr Pbuf): ErrT {.cdecl.} =
  ## Called by LwIP to send a packet
  var pkt: IonPacket
  if not ion_user_alloc(addr pkt): return -1

  # Copy pbuf chain to fixed slab
  let tot_len = p.tot_len
  discard pbuf_copy_partial(p, pkt.data, tot_len, 0)
  pkt.len = tot_len

  if ion_user_tx(pkt):
    return ERR_OK
  else:
    ion_user_free(pkt)
    return -1

type
  SocketState* = enum
    CLOSED, LISTEN, SYN_SENT, ESTABLISHED

  # The Shadow Socket
  NexusSock* = object
    fd*: int
    state*: SocketState
    pcb*: ptr TcpPcb           # The LwIP Object
    rx_buf*: array[8192, byte] # 8KB RX Buffer
    rx_head*: int
    rx_tail*: int
    rx_len*: int

var socket_table*: array[1024, ptr NexusSock]

# LwIP Callbacks
proc on_tcp_recv_cb(arg: pointer; pcb: ptr TcpPcb; p: ptr Pbuf;
    err: ErrT): ErrT {.cdecl.} =
  let sock = cast[ptr NexusSock](arg)
  if p == nil:
    # Connection closed
    sock.state = CLOSED
    return ERR_OK

  # Copy pbuf data to circular buffer
  let tot_len = p.tot_len
  var offset: uint16 = 0

  # Check for overflow
  if sock.rx_len + int(tot_len) > 8192:
    # For now, discard or handle backpressure?
    # TODO: real backpressure would be NOT calling tcp_recved until consumed
    discard pbuf_free(p)
    return ERR_OK

  while offset < tot_len:
    let space = 8192 - sock.rx_tail
    let chunk = min(int(tot_len - offset), space)
    discard pbuf_copy_partial(p, addr sock.rx_buf[sock.rx_tail], uint16(chunk), offset)
    sock.rx_tail = (sock.rx_tail + chunk) mod 8192
    sock.rx_len += chunk
    offset += uint16(chunk)

  discard pbuf_free(p)
  return ERR_OK

proc tcp_recved*(pcb: ptr TcpPcb; len: uint16) {.importc: "tcp_recved",
    header: "lwip/tcp.h".}

proc glue_read*(sock: ptr NexusSock; buf: pointer; len: int): int =
  if sock.rx_len == 0:
    if sock.state == CLOSED: return 0 # EOF
    return -1 # EAGAIN

  let to_read = min(len, sock.rx_len)
  var read_so_far = 0

  while read_so_far < to_read:
    let available = 8192 - sock.rx_head
    let chunk = min(to_read - read_so_far, available)
    copyMem(cast[pointer](cast[uint](buf) + uint(read_so_far)),
        addr sock.rx_buf[sock.rx_head], chunk)
    sock.rx_head = (sock.rx_head + chunk) mod 8192
    sock.rx_len -= chunk
    read_so_far += chunk

  # Notify LwIP we consumed data to open window
  if sock.pcb != nil:
    tcp_recved(sock.pcb, uint16(read_so_far))

  return read_so_far

# LwIP Callbacks
proc on_connected_cb(arg: pointer; pcb: ptr TcpPcb; err: ErrT): ErrT {.cdecl.} =
  let sock = cast[ptr NexusSock](arg)
  if err == ERR_OK:
    sock.state = ESTABLISHED
  return ERR_OK

proc lwip_init() {.importc: "lwip_init", header: "lwip/init.h".}

proc netif_add(netif: ptr Netif; ipaddr, netmask, gw: ptr IpAddr; state: pointer;
               init, input: pointer): ptr Netif {.importc: "netif_add",
                   header: "lwip/netif.h".}
proc netif_set_up(netif: ptr Netif) {.importc: "netif_set_up",
    header: "lwip/netif.h".}
proc netif_set_link_up(netif: ptr Netif) {.importc: "netif_set_link_up",
    header: "lwip/netif.h".}
proc netif_set_default(netif: ptr Netif) {.importc: "netif_set_default",
    header: "lwip/netif.h".}
proc ethernet_input(p: ptr Pbuf, ni: ptr Netif): ErrT {.importc: "ethernet_input",
    header: "netif/ethernet.h".}

proc dummy_netif_init(ni: ptr Netif): ErrT {.cdecl.} =
  console_write(cstring("[Glue] Netif Init Called\n"), 25)
  return 0

proc membrane_init*() {.exportc, cdecl.} =
  when not defined(is_membrane):
    ion_pool_init()
  ion_user_init()

  # EMERGENCY PHASE 34.3: Address Verify (Userland Side)
  let sys = cast[ptr SysTable](SYS_TABLE_ADDR)
  if sys != nil and sys.fn_vfs_write != nil:
    var msg = "[Membrane] Input Ring Ptr @ 0x"
    discard sys.fn_vfs_write(1, unsafeAddr msg[0], uint64(msg.len))
    # Print hex address
    let ring_addr = cast[uint64](membrane_input_ring_ptr)
    for i in countdown(15, 0):
      let nibble = (ring_addr shr (i * 4)) and 0xF
      let hex_char = if nibble < 10: char(nibble + ord('0')) else: char(nibble -
          10 + ord('A'))
      discard sys.fn_vfs_write(1, unsafeAddr hex_char, 1)
    let newline = "\n"
    discard sys.fn_vfs_write(1, unsafeAddr newline[0], 1)

  lwip_init()

  # Set up Virtual Interface for Subject (10.0.2.16)
  var ip, mask, gw: IpAddr
  # Identity mapped packing (A.B.C.D -> uint32)
  proc pack(a, b, c, d: uint8): uint32 =
    (uint32(a) shl 0) or (uint32(b) shl 8) or (uint32(c) shl 16) or (uint32(d) shl 24)

  cast[ptr uint32](addr ip)[] = pack(10, 0, 2, 16)
  cast[ptr uint32](addr mask)[] = pack(255, 255, 255, 0)
  cast[ptr uint32](addr gw)[] = pack(10, 0, 2, 2)

  # Initialize netif struct
  membrane_netif.mtu = 1500
  membrane_netif.linkoutput = membrane_output
  # Flags: UP, ETHARP, LINK_UP, ETHERNET (0x01 | 0x08 | 0x10 | 0x40 = 0x59)
  membrane_netif.flags = 0x59

  # Set MAC (52:54:00:12:34:57)
  membrane_netif.hwaddr_len = 6
  membrane_netif.hwaddr[0] = 0x52
  membrane_netif.hwaddr[1] = 0x54
  membrane_netif.hwaddr[2] = 0x00
  membrane_netif.hwaddr[3] = 0x12
  membrane_netif.hwaddr[4] = 0x34
  membrane_netif.hwaddr[5] = 0x57

  discard netif_add(addr membrane_netif, addr ip, addr mask, addr gw,
                    nil, cast[pointer](dummy_netif_init), cast[pointer](ethernet_input))
  netif_set_default(addr membrane_netif)
  netif_set_up(addr membrane_netif)
  netif_set_link_up(addr membrane_netif)

proc pump_membrane_stack*() {.exportc, cdecl.} =
  # 1. Drain ION RX Ring -> LwIP input
  var pkt: IonPacket
  while ion_user_rx(addr pkt):
    # Wrap in Pbuf
    let pb = pbuf_alloc(PBUF_RAW, int(pkt.len), PBUF_POOL)
    if pb != nil:
      discard pbuf_take(pb, pkt.data, pkt.len)
      # Feed to Stack
      if membrane_netif.input(pb, addr membrane_netif) != ERR_OK:
        discard pbuf_free(pb)

    # Return slab to pool
    ion_user_free(pkt)

  # 2. Check Timers
  # console_write(cstring("P"), 1)
  sys_check_timeouts()
  # Phase 33: Explicit yield if we aren't calling sys_read
  let sys = cast[ptr SysTable](SYS_TABLE_ADDR)
  if sys.fn_yield != nil:
    sys.fn_yield()

proc tcp_write*(pcb: ptr TcpPcb; dataptr: pointer; len: uint16;
    apiflags: uint8): ErrT {.
    importc: "tcp_write", header: "lwip/tcp.h".}
proc tcp_output*(pcb: ptr TcpPcb): ErrT {.importc: "tcp_output",
    header: "lwip/tcp.h".}

proc glue_write*(sock: ptr NexusSock; buf: pointer; len: int): int =
  if sock.pcb == nil or sock.state != ESTABLISHED: return -1

  # Flags: TCP_WRITE_FLAG_COPY = 0x01
  let err = tcp_write(sock.pcb, buf, uint16(len), 0x01)
  if err == ERR_OK:
    discard tcp_output(sock.pcb)
    return len
  return -1

proc tcp_recv*(pcb: ptr TcpPcb; cb: pointer) {.importc: "tcp_recv",
    header: "lwip/tcp.h".}

proc glue_connect*(sock: ptr NexusSock; ip: ptr IpAddr; port: uint16): int =
  if sock.pcb == nil:
    sock.pcb = tcp_new()
    if sock.pcb == nil: return -1

  # Reset RX state
  sock.rx_head = 0
  sock.rx_tail = 0
  sock.rx_len = 0

  # 1. Setup LwIP Callback
  tcp_arg(sock.pcb, sock)
  tcp_recv(sock.pcb, on_tcp_recv_cb)
  # tcp_err(sock.pcb, on_tcp_error) # Todo
  # tcp_sent(sock.pcb, on_tcp_sent) # Todo

  # 2. Start Handshake
  let err = tcp_connect(sock.pcb, ip, port, on_connected_cb)

  if err == ERR_OK:
    sock.state = SYN_SENT
    return 0
  return -1

# The Yield Mechanism (Cooperative Multitasking)
proc fiber_yield*() {.exportc, cdecl.} =
  ## Yield control back to the Kernel's networking fiber.
  ## This allows VirtIO polling and packet processing to occur.
  when defined(is_membrane):
    # Use the Kernel-provided yield function pointer
    type YieldFunc = proc() {.cdecl.}
    let yield_ptr = cast[YieldFunc](0x83000FF0'u64)
    if yield_ptr != nil:
      yield_ptr()