import asyncdispatch, asyncnet, net, port_prediction, strformat

from nativesockets import
  SockAddr,
  SockAddr_storage,
  SockLen,
  getSockOptInt,
  setSockOptInt
from sequtils import any, map

type
  Attempt* = object
    ## A hole punching attempt.
    socks*: seq[AsyncSocket]
    dstIp*: IpAddress
    dstPorts*: seq[Port]
    future*: Future[(AsyncSocket, Port)]

  Puncher* = ref object
    socks*: seq[AsyncSocket]
    natProps*: NatProperties
    attempts*: seq[Attempt]

  PunchHoleError* = object of ValueError

var IPPROTO_IP {.importc: "IPPROTO_IP", header: "<netinet/in.h>".}: cint
var IP_TTL {.importc: "IP_TTL", header: "<netinet/in.h>".}: cint

const Timeout = 3000
const InitiatorMaxOutPacketCount = 1000
const ResponderMaxOutPacketCount = 70
const MaxOutPacketCount = max(InitiatorMaxOutPacketCount,
                              ResponderMaxOutPacketCount)

proc srcPort(sock: AsyncSocket): Port =
  result = sock.getLocalAddr[1]

proc `==`(a, b: Attempt): bool =
  ## ``==`` for hole punching attempts.
  ##
  ## Two hole punching attempts are considered equal if their ``dstIp`` is
  ## equal and their ``dstPorts`` overlap.
  a.dstIp == b.dstIp and a.dstPorts.any(proc (p: Port): bool = p in b.dstPorts)

proc initPuncher*(sock: AsyncSocket, probedSrcPorts: seq[Port]): Puncher =
  # TODO: determine IP_TTL
  let (_, primarySrcPort) = sock.getLocalAddr()
  let natProps = getNatProperties(primarySrcPort, probedSrcPorts)
  result = Puncher(socks: @[sock], natProps: natProps)
  if result.natProps.natType == SymmetricRandom:
    # our NAT is of the evil symmetric type with random port allocation. We are
    # trying to help the other peer by allocating a lot of auxillary sockets
    # for punching more holes
    result.socks.setLen(MaxOutPacketCount)
    for i in 1 .. MaxOutPacketCount - 1:
      result.socks[i] = newAsyncSocket(sockType = SOCK_DGRAM,
                                       protocol = IPPROTO_UDP, buffered = false)
      result.socks[i].bindAddr(Port(0))

proc primarySrcPort*(puncher: Puncher): Port =
  puncher.socks[0].srcPort

proc punch(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
           peerProbedPorts: seq[Port], isInitiating: bool, msg: string):
          Future[Attempt] {.async.} =
  let punchFuture = newFuture[(AsyncSocket, Port)]("punch")
  let peerNatProps = getNatProperties(peerPort, peerProbedPorts)
  let maxOutPacketCount = if isInitiating:
    InitiatorMaxOutPacketCount
  else:
    ResponderMaxOutPacketCount
  var sockCount = 1
  if puncher.natProps.natType == SymmetricRandom:
    if peerNatProps.natType == SymmetricRandom:
      # If both peers are behind a SymmetricRandom NAT we give up.
      raise newException(PunchHoleError,
                         "both peers behind symmetric NAT with random port allocation")
    sockCount = maxOutPacketCount
  let predictedDstPorts = predictPortRange(peerNatProps, maxOutPacketCount.uint16)
  result = Attempt(dstIp: peerIp, dstPorts: predictedDstPorts,
                   future: punchFuture)
  if puncher.attempts.contains(result):
    raise newException(PunchHoleError,
                       "hole punching to given destination already active")
  puncher.attempts.add(result)
  let srcPorts = puncher.socks[0 .. sockCount - 1].map(srcPort)
  echo &"sending msg {msg} to {peerIp}, srcPorts: {srcPorts}, dstPorts: {result.dstPorts}"
  var peerAddr: Sockaddr_storage
  var peerSockLen: SockLen
  try:
    var defaultTTL: int
    for i in 0 .. sockCount - 1:
      let sock = puncher.socks[i]
      if isInitiating:
        defaultTTL = sock.getFd.getSockOptInt(IPPROTO_IP, IP_TTL)
        sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, 2)
      for dstPort in result.dstPorts:
        toSockAddr(result.dstIp, dstPort, peerAddr, peerSockLen)
        # TODO: replace asyncdispatch.sendTo with asyncnet.sendTo (Nim 1.4 required)
        await sendTo(sock.getFd.AsyncFD, msg.cstring, msg.len,
                     cast[ptr SockAddr](addr peerAddr), peerSockLen)
      if isInitiating:
        sock.getFd.setSockOptInt(IPPROTO_IP, IP_TTL, defaultTTL)
  except OSError as e:
    raise newException(PunchHoleError, e.msg)

proc initiate*(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
               peerProbedPorts: seq[Port]): Future[Attempt] =
  punch(puncher, peerIp, peerPort, peerProbedPorts, true, "SYN")
  
proc respond*(puncher: Puncher, peerIp: IpAddress, peerPort: Port,
              peerProbedPorts: seq[Port]): Future[Attempt] =
  punch(puncher, peerIp, peerPort, peerProbedPorts, false, "ACK")

proc finalize*(attempt: Attempt): Future[(AsyncSocket, Port)] {.async.} =
  await attempt.future or sleepAsync(Timeout)
  if attempt.future.finished:
    result = attempt.future.read()
  else:
    raise newException(PunchHoleError, "timeout")

proc handleMsg*(puncher: Puncher, msg: string, sock: AsyncSocket,
                peerIp: IpAddress, peerPort: Port) =
  ## Handles an incoming UDP message which may complete the Futures returned by
  ## ``initiate`` and ``respond``.
  if msg == "SYN":
    # We received a SYN packet. We ignore it because we expected it to be
    # filtered by our NAT.
    return
  let query = Attempt(dstIp: peerIp, dstPorts: @[peerPort])
  let i = puncher.attempts.find(query)
  if i != -1:
    if msg == "ACK":
      echo &"handling ACK message from {peerIp}:{peerPort}"
    else:
      echo &"handling QUIC message from {peerIp}:{peerPort}"
    puncher.attempts[i].future.complete((sock, peerPort))
    puncher.attempts.del(i)
  else:
    echo &"received unexpected packet from {peerIp}:{peerPort}"

proc handleMsg*(puncher: Puncher, msg: string, sock: AsyncSocket,
                peerAddr: SockAddr | Sockaddr_storage, peerSockLen: SockLen) =
  var peerIp: IpAddress
  var peerPort: Port
  fromSockAddr(peerAddr, peerSockLen, peerIp, peerPort)
  handleMsg(puncher, msg, sock, peerIp, peerPort)