# 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*: RingBuffer # Bytes waiting for read() - TODO

var socket_table*: array[1024, ptr NexusSock]

# 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()
  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
  sys_check_timeouts()

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 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

  # 1. Setup LwIP Callback
  tcp_arg(sock.pcb, sock)
  # tcp_err(sock.pcb, on_tcp_error) # TODO
  # tcp_recv(sock.pcb, on_tcp_recv) # 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