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

{.used.}

{.push raises: [].}

import tables
import chronos, stew/[byteutils, endians2, shims/net]
import
  ../../libp2p/[
    stream/connection,
    protocols/connectivity/relay/utils,
    transports/tcptransport,
    transports/tortransport,
    upgrademngrs/upgrade,
    multiaddress,
    errors,
    builders,
  ]

type TorServerStub* = ref object of RootObj
  tcpTransport: TcpTransport
  addrTable: Table[string, string]

proc new*(T: typedesc[TorServerStub]): T {.public.} =
  T(
    tcpTransport: TcpTransport.new(flags = {ReuseAddr}, upgrade = Upgrade()),
    addrTable: initTable[string, string](),
  )

proc registerAddr*(self: TorServerStub, key: string, val: string) =
  self.addrTable[key] = val

proc start*(self: TorServerStub, address: TransportAddress) {.async.} =
  let ma = @[MultiAddress.init(address).tryGet()]

  await self.tcpTransport.start(ma)

  var msg = newSeq[byte](3)
  while self.tcpTransport.running:
    let connSrc = await self.tcpTransport.accept()
    await connSrc.readExactly(addr msg[0], 3)

    await connSrc.write(@[05'u8, 00])

    msg = newSeq[byte](4)
    await connSrc.readExactly(addr msg[0], 4)
    let atyp = msg[3]
    let address =
      case atyp
      of Socks5AddressType.IPv4.byte:
        let n = 4 + 2 # +2 bytes for the port
        msg = newSeq[byte](n)
        await connSrc.readExactly(addr msg[0], n)
        var ip: array[4, byte]
        for i, e in msg[0 ..^ 3]:
          ip[i] = e
        $(ipv4(ip)) & ":" & $(Port(fromBytesBE(uint16, msg[^2 ..^ 1])))
      of Socks5AddressType.IPv6.byte:
        let n = 16 + 2 # +2 bytes for the port
        msg = newSeq[byte](n) # +2 bytes for the port
        await connSrc.readExactly(addr msg[0], n)
        var ip: array[16, byte]
        for i, e in msg[0 ..^ 3]:
          ip[i] = e
        $(ipv6(ip)) & ":" & $(Port(fromBytesBE(uint16, msg[^2 ..^ 1])))
      of Socks5AddressType.FQDN.byte:
        await connSrc.readExactly(addr msg[0], 1)
        let n = int(uint8.fromBytes(msg[0 .. 0])) + 2 # +2 bytes for the port
        msg = newSeq[byte](n)
        await connSrc.readExactly(addr msg[0], n)
        string.fromBytes(msg[0 ..^ 3]) & ":" &
          $(Port(fromBytesBE(uint16, msg[^2 ..^ 1])))
      else:
        raise newException(LPError, "Address not supported")

    let tcpIpAddr = self.addrTable[$(address)]

    await connSrc.write(@[05'u8, 00, 00, 01, 00, 00, 00, 00, 00, 00])

    let connDst =
      await self.tcpTransport.dial("", MultiAddress.init(tcpIpAddr).tryGet())

    await bridge(connSrc, connDst)
    await allFutures(connSrc.close(), connDst.close())

proc stop*(self: TorServerStub) {.async.} =
  await self.tcpTransport.stop()