# Nim-LibP2P
# Copyright (c) 2023 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 macros
import stew/[objects, results]
import ../../../peerinfo,
       ../../../signed_envelope

# Circuit Relay V1 Message

type
  RelayType* {.pure.} = enum
    Hop = 1
    Stop = 2
    Status = 3
    CanHop = 4

  StatusV1* {.pure.} = enum
    Success = 100
    HopSrcAddrTooLong = 220
    HopDstAddrTooLong = 221
    HopSrcMultiaddrInvalid = 250
    HopDstMultiaddrInvalid = 251
    HopNoConnToDst = 260
    HopCantDialDst = 261
    HopCantOpenDstStream = 262
    HopCantSpeakRelay = 270
    HopCantRelayToSelf = 280
    StopSrcAddrTooLong = 320
    StopDstAddrTooLong = 321
    StopSrcMultiaddrInvalid = 350
    StopDstMultiaddrInvalid = 351
    StopRelayRefused = 390
    MalformedMessage = 400

  RelayPeer* = object
    peerId*: PeerId
    addrs*: seq[MultiAddress]

  RelayMessage* = object
    msgType*: Opt[RelayType]
    srcPeer*: Opt[RelayPeer]
    dstPeer*: Opt[RelayPeer]
    status*: Opt[StatusV1]

proc encode*(msg: RelayMessage): ProtoBuffer =
  result = initProtoBuffer()

  msg.msgType.withValue(typ):
    result.write(1, typ.ord.uint)
  msg.srcPeer.withValue(srcPeer):
    var peer = initProtoBuffer()
    peer.write(1, srcPeer.peerId)
    for ma in srcPeer.addrs:
      peer.write(2, ma.data.buffer)
    peer.finish()
    result.write(2, peer.buffer)
  msg.dstPeer.withValue(dstPeer):
    var peer = initProtoBuffer()
    peer.write(1, dstPeer.peerId)
    for ma in dstPeer.addrs:
      peer.write(2, ma.data.buffer)
    peer.finish()
    result.write(3, peer.buffer)
  msg.status.withValue(status):
    result.write(4, status.ord.uint)

  result.finish()

proc decode*(_: typedesc[RelayMessage], buf: seq[byte]): Opt[RelayMessage] =
  var
    rMsg: RelayMessage
    msgTypeOrd: uint32
    src: RelayPeer
    dst: RelayPeer
    statusOrd: uint32
    pbSrc: ProtoBuffer
    pbDst: ProtoBuffer

  let pb = initProtoBuffer(buf)

  if ? pb.getField(1, msgTypeOrd).toOpt():
    if msgTypeOrd.int notin RelayType:
      return Opt.none(RelayMessage)
    rMsg.msgType = Opt.some(RelayType(msgTypeOrd))

  if ? pb.getField(2, pbSrc).toOpt():
    discard ? pbSrc.getField(1, src.peerId).toOpt()
    discard ? pbSrc.getRepeatedField(2, src.addrs).toOpt()
    rMsg.srcPeer = Opt.some(src)

  if ? pb.getField(3, pbDst).toOpt():
    discard ? pbDst.getField(1, dst.peerId).toOpt()
    discard ? pbDst.getRepeatedField(2, dst.addrs).toOpt()
    rMsg.dstPeer = Opt.some(dst)

  if ? pb.getField(4, statusOrd).toOpt():
    var status: StatusV1
    if not checkedEnumAssign(status, statusOrd):
      return Opt.none(RelayMessage)
    rMsg.status = Opt.some(status)
  Opt.some(rMsg)

# Voucher

type
  Voucher* = object
    relayPeerId*: PeerId     # peer ID of the relay
    reservingPeerId*: PeerId # peer ID of the reserving peer
    expiration*: uint64      # UNIX UTC expiration time for the reservation

proc decode*(_: typedesc[Voucher], buf: seq[byte]): Result[Voucher, ProtoError] =
  let pb = initProtoBuffer(buf)
  var v = Voucher()

  ? pb.getRequiredField(1, v.relayPeerId)
  ? pb.getRequiredField(2, v.reservingPeerId)
  ? pb.getRequiredField(3, v.expiration)

  ok(v)

proc encode*(v: Voucher): seq[byte] =
  var pb = initProtoBuffer()

  pb.write(1, v.relayPeerId)
  pb.write(2, v.reservingPeerId)
  pb.write(3, v.expiration)

  pb.finish()
  pb.buffer

proc init*(T: typedesc[Voucher],
          relayPeerId: PeerId,
          reservingPeerId: PeerId,
          expiration: uint64): T =
  T(
    relayPeerId = relayPeerId,
    reservingPeerId = reservingPeerId,
    expiration: expiration
  )

type SignedVoucher* = SignedPayload[Voucher]

proc payloadDomain*(_: typedesc[Voucher]): string = "libp2p-relay-rsvp"
proc payloadType*(_: typedesc[Voucher]): seq[byte] = @[ (byte)0x03, (byte)0x02 ]

proc checkValid*(spr: SignedVoucher): Result[void, EnvelopeError] =
  if not spr.data.relayPeerId.match(spr.envelope.publicKey):
    err(EnvelopeInvalidSignature)
  else:
    ok()

# Circuit Relay V2 Hop Message

type
  Peer* = object
    peerId*: PeerId
    addrs*: seq[MultiAddress]
  Reservation* = object
    expire*: uint64              # required, Unix expiration time (UTC)
    addrs*: seq[MultiAddress]    # relay address for reserving peer
    svoucher*: Opt[seq[byte]]    # optional, reservation voucher
  Limit* = object
    duration*: uint32            # seconds
    data*: uint64                # bytes

  StatusV2* = enum
    Ok = 100
    ReservationRefused = 200
    ResourceLimitExceeded = 201
    PermissionDenied = 202
    ConnectionFailed = 203
    NoReservation = 204
    MalformedMessage = 400
    UnexpectedMessage = 401
  HopMessageType* {.pure.} = enum
    Reserve = 0
    Connect = 1
    Status = 2
  HopMessage* = object
    msgType*: HopMessageType
    peer*: Opt[Peer]
    reservation*: Opt[Reservation]
    limit*: Limit
    status*: Opt[StatusV2]

proc encode*(msg: HopMessage): ProtoBuffer =
  var pb = initProtoBuffer()

  pb.write(1, msg.msgType.ord.uint)
  msg.peer.withValue(peer):
    var ppb = initProtoBuffer()
    ppb.write(1, peer.peerId)
    for ma in peer.addrs:
      ppb.write(2, ma.data.buffer)
    ppb.finish()
    pb.write(2, ppb.buffer)
  msg.reservation.withValue(rsrv):
    var rpb = initProtoBuffer()
    rpb.write(1, rsrv.expire)
    for ma in rsrv.addrs:
      rpb.write(2, ma.data.buffer)
    rsrv.svoucher.withValue(vouch):
      rpb.write(3, vouch)
    rpb.finish()
    pb.write(3, rpb.buffer)
  if msg.limit.duration > 0 or msg.limit.data > 0:
    var lpb = initProtoBuffer()
    if msg.limit.duration > 0: lpb.write(1, msg.limit.duration)
    if msg.limit.data > 0: lpb.write(2, msg.limit.data)
    lpb.finish()
    pb.write(4, lpb.buffer)
  msg.status.withValue(status):
    pb.write(5, status.ord.uint)

  pb.finish()
  pb

proc decode*(_: typedesc[HopMessage], buf: seq[byte]): Opt[HopMessage] =
  var msg: HopMessage
  let pb = initProtoBuffer(buf)

  var msgTypeOrd: uint32
  ? pb.getRequiredField(1, msgTypeOrd).toOpt()
  if not checkedEnumAssign(msg.msgType, msgTypeOrd):
    return Opt.none(HopMessage)

  var pbPeer: ProtoBuffer
  if ? pb.getField(2, pbPeer).toOpt():
    var peer: Peer
    ? pbPeer.getRequiredField(1, peer.peerId).toOpt()
    discard ? pbPeer.getRepeatedField(2, peer.addrs).toOpt()
    msg.peer = Opt.some(peer)

  var pbReservation: ProtoBuffer
  if ? pb.getField(3, pbReservation).toOpt():
    var
      svoucher: seq[byte]
      reservation: Reservation
    if ? pbReservation.getField(3, svoucher).toOpt():
      reservation.svoucher = Opt.some(svoucher)
    ? pbReservation.getRequiredField(1, reservation.expire).toOpt()
    discard ? pbReservation.getRepeatedField(2, reservation.addrs).toOpt()
    msg.reservation = Opt.some(reservation)

  var pbLimit: ProtoBuffer
  if ? pb.getField(4, pbLimit).toOpt():
    discard ? pbLimit.getField(1, msg.limit.duration).toOpt()
    discard ? pbLimit.getField(2, msg.limit.data).toOpt()

  var statusOrd: uint32
  if ? pb.getField(5, statusOrd).toOpt():
    var status: StatusV2
    if not checkedEnumAssign(status, statusOrd):
      return Opt.none(HopMessage)
    msg.status = Opt.some(status)
  Opt.some(msg)

# Circuit Relay V2 Stop Message

type
  StopMessageType* {.pure.} = enum
    Connect = 0
    Status = 1
  StopMessage* = object
    msgType*: StopMessageType
    peer*: Opt[Peer]
    limit*: Limit
    status*: Opt[StatusV2]


proc encode*(msg: StopMessage): ProtoBuffer =
  var pb = initProtoBuffer()

  pb.write(1, msg.msgType.ord.uint)
  msg.peer.withValue(peer):
    var ppb = initProtoBuffer()
    ppb.write(1, peer.peerId)
    for ma in peer.addrs:
      ppb.write(2, ma.data.buffer)
    ppb.finish()
    pb.write(2, ppb.buffer)
  if msg.limit.duration > 0 or msg.limit.data > 0:
    var lpb = initProtoBuffer()
    if msg.limit.duration > 0: lpb.write(1, msg.limit.duration)
    if msg.limit.data > 0: lpb.write(2, msg.limit.data)
    lpb.finish()
    pb.write(3, lpb.buffer)
  msg.status.withValue(status):
    pb.write(4, status.ord.uint)

  pb.finish()
  pb

proc decode*(_: typedesc[StopMessage], buf: seq[byte]): Opt[StopMessage] =
  var msg: StopMessage

  let pb = initProtoBuffer(buf)

  var msgTypeOrd: uint32
  ? pb.getRequiredField(1, msgTypeOrd).toOpt()
  if msgTypeOrd.int notin StopMessageType:
    return Opt.none(StopMessage)
  msg.msgType = StopMessageType(msgTypeOrd)


  var pbPeer: ProtoBuffer
  if ? pb.getField(2, pbPeer).toOpt():
    var peer: Peer
    ? pbPeer.getRequiredField(1, peer.peerId).toOpt()
    discard ? pbPeer.getRepeatedField(2, peer.addrs).toOpt()
    msg.peer = Opt.some(peer)

  var pbLimit: ProtoBuffer
  if ? pb.getField(3, pbLimit).toOpt():
    discard ? pbLimit.getField(1, msg.limit.duration).toOpt()
    discard ? pbLimit.getField(2, msg.limit.data).toOpt()

  var statusOrd: uint32
  if ? pb.getField(4, statusOrd).toOpt():
    var status: StatusV2
    if not checkedEnumAssign(status, statusOrd):
      return Opt.none(StopMessage)
    msg.status = Opt.some(status)
  Opt.some(msg)