# Nim-LibP2P
# Copyright (c) 2023-2024 Status Research & Development GmbH
# Licensed under either of
#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
# at your option.
# This file may not be copied, modified, or distributed except according to
# those terms.

{.push raises: [].}

import pkg/[chronos, chronicles, stew/byteutils]
import ../../stream/connection, ../../utility, ../../varint, ../../vbuffer, ../muxer

logScope:
  topics = "libp2p mplexcoder"

type
  MessageType* {.pure.} = enum
    New
    MsgIn
    MsgOut
    CloseIn
    CloseOut
    ResetIn
    ResetOut

  Msg* = tuple[id: uint64, msgType: MessageType, data: seq[byte]]

  InvalidMplexMsgType* = object of MuxerError

# https://github.com/libp2p/specs/tree/master/mplex#writing-to-a-stream
const MaxMsgSize* = 1 shl 20 # 1mb

proc newInvalidMplexMsgType*(): ref InvalidMplexMsgType =
  newException(InvalidMplexMsgType, "invalid message type")

proc readMsg*(
    conn: Connection
): Future[Msg] {.async: (raises: [CancelledError, LPStreamError, MuxerError]).} =
  let header = await conn.readVarint()
  trace "read header varint", varint = header, conn

  let data = await conn.readLp(MaxMsgSize)
  trace "read data", dataLen = data.len, data = shortLog(data), conn

  let msgType = header and 0x7
  if msgType.int > ord(MessageType.ResetOut):
    raise newInvalidMplexMsgType()

  return (header shr 3, MessageType(msgType), data)

proc writeMsg*(
    conn: Connection, id: uint64, msgType: MessageType, data: seq[byte] = @[]
): Future[void] {.async: (raises: [CancelledError, LPStreamError], raw: true).} =
  var
    left = data.len
    offset = 0
    buf = initVBuffer()

  # Split message into length-prefixed chunks
  while left > 0 or data.len == 0:
    let chunkSize =
      if left > MaxMsgSize:
        MaxMsgSize - 64
      else:
        left

    buf.writePBVarint(id shl 3 or ord(msgType).uint64)
    buf.writeSeq(data.toOpenArray(offset, offset + chunkSize - 1))
    left = left - chunkSize
    offset = offset + chunkSize

    if data.len == 0:
      break

  trace "writing mplex message",
    conn, id, msgType, data = data.len, encoded = buf.buffer.len

  # Write all chunks in a single write to avoid async races where a close
  # message gets written before some of the chunks
  conn.write(buf.buffer)

proc writeMsg*(
    conn: Connection, id: uint64, msgType: MessageType, data: string
): Future[void] {.async: (raises: [CancelledError, LPStreamError], raw: true).} =
  conn.writeMsg(id, msgType, data.toBytes())