{.used.}

import bearssl, chronos, options
import ../libp2p
import ../libp2p/[protocols/connectivity/relay/relay,
                  protocols/connectivity/relay/messages,
                  protocols/connectivity/relay/utils,
                  protocols/connectivity/relay/client]
import ./helpers
import std/times
import stew/byteutils

proc createSwitch(r: Relay): Switch =
  result = SwitchBuilder.new()
    .withRng(newRng())
    .withAddresses(@[ MultiAddress.init("/ip4/0.0.0.0/tcp/0").tryGet() ])
    .withTcpTransport()
    .withMplex()
    .withNoise()
    .withCircuitRelay(r)
    .build()

suite "Circuit Relay V2":

  suite "Reservation":
    asyncTeardown:
      await allFutures(src1.stop(), src2.stop(), rel.stop())
      checkTrackers()
    var
      ttl {.threadvar.}: int
      ldur {.threadvar.}: uint32
      ldata {.threadvar.}: uint64
      cl1 {.threadvar.}: RelayClient
      cl2 {.threadvar.}: RelayClient
      rv2 {.threadvar.}: Relay
      src1 {.threadvar.}: Switch
      src2 {.threadvar.}: Switch
      rel {.threadvar.}: Switch
      rsvp {.threadvar.}: Rsvp
      range {.threadvar.}: HSlice[times.DateTime, times.DateTime]

    asyncSetup:
      ttl = 3
      ldur = 60
      ldata = 2048
      cl1 = RelayClient.new()
      cl2 = RelayClient.new()
      rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                      limitDuration=ldur,
                      limitData=ldata,
                      maxCircuit=1)
      src1 = createSwitch(cl1)
      src2 = createSwitch(cl2)
      rel = createSwitch(rv2)

      await src1.start()
      await src2.start()
      await rel.start()
      await src1.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await src2.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      rsvp = await cl1.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata

    asyncTest "Too many reservations":
      let conn = await cl2.switch.dial(rel.peerInfo.peerId, rel.peerInfo.addrs, RelayV2HopCodec)
      let pb = encode(HopMessage(msgType: HopMessageType.Reserve))
      await conn.writeLp(pb.buffer)
      let msg = HopMessage.decode(await conn.readLp(RelayMsgSize)).get()
      check:
        msg.msgType == HopMessageType.Status
        msg.status == some(StatusV2.ReservationRefused)

    asyncTest "Too many reservations + Reconnect":
      expect(ReservationError):
        discard await cl2.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await rel.disconnect(src1.peerInfo.peerId)
      rsvp = await cl2.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata

    asyncTest "Reservation ttl expires":
      await sleepAsync(chronos.timer.seconds(ttl + 1))
      rsvp = await cl1.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata

    asyncTest "Reservation over relay":
      let
        rv2add = Relay.new()
        addrs = @[ MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                     $rel.peerInfo.peerId & "/p2p-circuit").get() ]
      rv2add.setup(src2)
      await rv2add.start()
      src2.mount(rv2add)
      rv2.maxCircuit.inc()

      rsvp = await cl2.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      range = now().utc + (ttl-3).seconds..now().utc + (ttl+3).seconds
      check:
        rsvp.expire.int64.fromUnix.utc in range
        rsvp.limitDuration == ldur
        rsvp.limitData == ldata
      expect(ReservationError):
        discard await cl1.reserve(src2.peerInfo.peerId, addrs)

  suite "Connection":
    asyncTeardown:
      checkTrackers()
    var
      customProtoCodec {.threadvar.}: string
      proto {.threadvar.}: LPProtocol
      ttl {.threadvar.}: int
      ldur {.threadvar.}: uint32
      ldata {.threadvar.}: uint64
      srcCl {.threadvar.}: RelayClient
      dstCl {.threadvar.}: RelayClient
      rv2 {.threadvar.}: Relay
      src {.threadvar.}: Switch
      dst {.threadvar.}: Switch
      rel {.threadvar.}: Switch
      rsvp {.threadvar.}: Rsvp
      conn {.threadvar.}: Connection

    asyncSetup:
      customProtoCodec = "/test"
      proto = new LPProtocol
      proto.codec = customProtoCodec
      ttl = 60
      ldur = 120
      ldata = 16384
      srcCl = RelayClient.new()
      dstCl = RelayClient.new()
      src = createSwitch(srcCl)
      dst = createSwitch(dstCl)
      rel = newStandardSwitch()

    asyncTest "Connection succeed":
      proto.handler = proc(conn: Connection, proto: string) {.async.} =
        check: "test1" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("test2")
        check: "test3" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("test4")
        await conn.close()
      rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                      limitDuration=ldur,
                      limitData=ldata)
      rv2.setup(rel)
      rel.mount(rv2)
      dst.mount(proto)

      await rel.start()
      await src.start()
      await dst.start()

      let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                $rel.peerInfo.peerId & "/p2p-circuit").get()

      await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

      rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)

      conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
      await conn.writeLp("test1")
      check: "test2" == string.fromBytes(await conn.readLp(1024))
      await conn.writeLp("test3")
      check: "test4" == string.fromBytes(await conn.readLp(1024))
      await allFutures(conn.close())
      await allFutures(src.stop(), dst.stop(), rel.stop())

    asyncTest "Connection duration exceeded":
      ldur = 3
      proto.handler = proc(conn: Connection, proto: string) {.async.} =
        check "wanna sleep?" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("yeah!")
        check "go!" == string.fromBytes(await conn.readLp(1024))
        await sleepAsync(chronos.timer.seconds(ldur + 1))
        await conn.writeLp("that was a cool power nap")
        await conn.close()
      rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                      limitDuration=ldur,
                      limitData=ldata)
      rv2.setup(rel)
      rel.mount(rv2)
      dst.mount(proto)

      await rel.start()
      await src.start()
      await dst.start()

      let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                $rel.peerInfo.peerId & "/p2p-circuit").get()

      await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

      rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
      await conn.writeLp("wanna sleep?")
      check: "yeah!" == string.fromBytes(await conn.readLp(1024))
      await conn.writeLp("go!")
      expect(LPStreamEOFError):
        discard await conn.readLp(1024)
      await allFutures(conn.close())
      await allFutures(src.stop(), dst.stop(), rel.stop())

    asyncTest "Connection data exceeded":
      ldata = 1000
      proto.handler = proc(conn: Connection, proto: string) {.async.} =
        check "count me the better story you know" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("do you expect a lorem ipsum or...?")
        check "surprise me!" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("""Call me Ishmael. Some years ago--never mind how long
precisely--having little or no money in my purse, and nothing
particular to interest me on shore, I thought I would sail about a
little and see the watery part of the world. It is a way I have of
driving off the spleen and regulating the circulation. Whenever I
find myself growing grim about the mouth; whenever it is a damp,
drizzly November in my soul; whenever I find myself involuntarily
pausing before coffin warehouses, and bringing up the rear of every
funeral I meet; and especially whenever my hypos get such an upper
hand of me, that it requires a strong moral principle to prevent me
from deliberately stepping into the street, and methodically knocking
people's hats off--then, I account it high time to get to sea as soon
as I can. This is my substitute for pistol and ball. With a
philosophical flourish Cato throws himself upon his sword; I quietly
take to the ship.""")
      rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                      limitDuration=ldur,
                      limitData=ldata)
      rv2.setup(rel)
      rel.mount(rv2)
      dst.mount(proto)

      await rel.start()
      await src.start()
      await dst.start()

      let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                $rel.peerInfo.peerId & "/p2p-circuit").get()

      await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

      rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
      await conn.writeLp("count me the better story you know")
      check: "do you expect a lorem ipsum or...?" == string.fromBytes(await conn.readLp(1024))
      await conn.writeLp("surprise me!")
      expect(LPStreamEOFError):
        discard await conn.readLp(1024)
      await allFutures(conn.close())
      await allFutures(src.stop(), dst.stop(), rel.stop())

    asyncTest "Reservation ttl expire during connection":
      ttl = 3
      proto.handler = proc(conn: Connection, proto: string) {.async.} =
        check: "test1" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("test2")
        check: "test3" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("test4")
        await conn.close()
      rv2 = Relay.new(reservationTTL=initDuration(seconds=ttl),
                      limitDuration=ldur,
                      limitData=ldata)
      rv2.setup(rel)
      rel.mount(rv2)
      dst.mount(proto)

      await rel.start()
      await src.start()
      await dst.start()

      let addrs = MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                $rel.peerInfo.peerId & "/p2p-circuit").get()

      await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await dst.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)

      rsvp = await dstCl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
      await conn.writeLp("test1")
      check: "test2" == string.fromBytes(await conn.readLp(1024))
      await conn.writeLp("test3")
      check: "test4" == string.fromBytes(await conn.readLp(1024))
      await src.disconnect(rel.peerInfo.peerId)
      await sleepAsync(chronos.timer.seconds(ttl + 1))

      expect(DialFailedError):
        check: conn.atEof()
        await conn.close()
        await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
        conn = await src.dial(dst.peerInfo.peerId, @[ addrs ], customProtoCodec)
      await allFutures(conn.close())
      await allFutures(src.stop(), dst.stop(), rel.stop())

    asyncTest "Connection over relay":
      # src => rel => rel2 => dst
      # rel2 reserve rel
      # dst reserve rel2
      # src try to connect with dst
      proto.handler = proc(conn: Connection, proto: string) {.async.} =
        raise newException(CatchableError, "Should not be here")
      let
        rel2Cl = RelayClient.new(canHop = true)
        rel2 = createSwitch(rel2Cl)
        rv2 = Relay.new()
      rv2.setup(rel)
      rel.mount(rv2)
      dst.mount(proto)
      await rel.start()
      await rel2.start()
      await src.start()
      await dst.start()

      let
        addrs = @[ MultiAddress.init($rel.peerInfo.addrs[0] & "/p2p/" &
                                     $rel.peerInfo.peerId & "/p2p-circuit/p2p/" &
                                     $rel2.peerInfo.peerId & "/p2p/" &
                                     $rel2.peerInfo.peerId & "/p2p-circuit").get() ]

      await src.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await rel2.connect(rel.peerInfo.peerId, rel.peerInfo.addrs)
      await dst.connect(rel2.peerInfo.peerId, rel2.peerInfo.addrs)

      rsvp = await rel2Cl.reserve(rel.peerInfo.peerId, rel.peerInfo.addrs)
      let rsvp2 = await dstCl.reserve(rel2.peerInfo.peerId, rel2.peerInfo.addrs)

      expect(DialFailedError):
        conn = await src.dial(dst.peerInfo.peerId, addrs, customProtoCodec)
      await allFutures(conn.close())
      await allFutures(src.stop(), dst.stop(), rel.stop(), rel2.stop())

    asyncTest "Connection using ClientRelay":
      var
        protoABC = new LPProtocol
        protoBCA = new LPProtocol
        protoCAB = new LPProtocol
      protoABC.codec = "/abctest"
      protoABC.handler = proc(conn: Connection, proto: string) {.async.} =
        check: "testABC1" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("testABC2")
        check: "testABC3" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("testABC4")
        await conn.close()
      protoBCA.codec = "/bcatest"
      protoBCA.handler = proc(conn: Connection, proto: string) {.async.} =
        check: "testBCA1" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("testBCA2")
        check: "testBCA3" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("testBCA4")
        await conn.close()
      protoCAB.codec = "/cabtest"
      protoCAB.handler = proc(conn: Connection, proto: string) {.async.} =
        check: "testCAB1" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("testCAB2")
        check: "testCAB3" == string.fromBytes(await conn.readLp(1024))
        await conn.writeLp("testCAB4")
        await conn.close()

      let
        clientA = RelayClient.new(canHop = true)
        clientB = RelayClient.new(canHop = true)
        clientC = RelayClient.new(canHop = true)
        switchA = createSwitch(clientA)
        switchB = createSwitch(clientB)
        switchC = createSwitch(clientC)

      switchA.mount(protoBCA)
      switchB.mount(protoCAB)
      switchC.mount(protoABC)

      await switchA.start()
      await switchB.start()
      await switchC.start()

      let
        addrsABC = MultiAddress.init($switchB.peerInfo.addrs[0] & "/p2p/" &
                                     $switchB.peerInfo.peerId & "/p2p-circuit").get()
        addrsBCA = MultiAddress.init($switchC.peerInfo.addrs[0] & "/p2p/" &
                                     $switchC.peerInfo.peerId & "/p2p-circuit").get()
        addrsCAB = MultiAddress.init($switchA.peerInfo.addrs[0] & "/p2p/" &
                                     $switchA.peerInfo.peerId & "/p2p-circuit").get()

      await switchA.connect(switchB.peerInfo.peerId, switchB.peerInfo.addrs)
      await switchB.connect(switchC.peerInfo.peerId, switchC.peerInfo.addrs)
      await switchC.connect(switchA.peerInfo.peerId, switchA.peerInfo.addrs)
      let rsvpABC = await clientA.reserve(switchC.peerInfo.peerId, switchC.peerInfo.addrs)
      let rsvpBCA = await clientB.reserve(switchA.peerInfo.peerId, switchA.peerInfo.addrs)
      let rsvpCAB = await clientC.reserve(switchB.peerInfo.peerId, switchB.peerInfo.addrs)
      let connABC = await switchA.dial(switchC.peerInfo.peerId, @[ addrsABC ], "/abctest")
      let connBCA = await switchB.dial(switchA.peerInfo.peerId, @[ addrsBCA ], "/bcatest")
      let connCAB = await switchC.dial(switchB.peerInfo.peerId, @[ addrsCAB ], "/cabtest")

      await connABC.writeLp("testABC1")
      await connBCA.writeLp("testBCA1")
      await connCAB.writeLp("testCAB1")
      check:
        "testABC2" == string.fromBytes(await connABC.readLp(1024))
        "testBCA2" == string.fromBytes(await connBCA.readLp(1024))
        "testCAB2" == string.fromBytes(await connCAB.readLp(1024))
      await connABC.writeLp("testABC3")
      await connBCA.writeLp("testBCA3")
      await connCAB.writeLp("testCAB3")
      check:
        "testABC4" == string.fromBytes(await connABC.readLp(1024))
        "testBCA4" == string.fromBytes(await connBCA.readLp(1024))
        "testCAB4" == string.fromBytes(await connCAB.readLp(1024))
      await allFutures(connABC.close(), connBCA.close(), connCAB.close())
      await allFutures(switchA.stop(), switchB.stop(), switchC.stop())