import asyncdispatch, asyncnet, strformat
from net import IpAddress, Port, `$`, `==`, parseIpAddress
from random import randomize, rand
from sequtils import any
import ip_packet
import message
import port_prediction
import puncher
import raw_socket
import utils

export Puncher, Initiator, PunchProgressCb, PunchHoleError, cleanup, initiate

type
  TNIAttempt = ref object of Attempt
    future: Future[AsyncSocket]

  TcpNutssInitiator* = ref object of Initiator

  Request = object
    srcIp: IpAddress
    srcPorts: seq[Port]
    dstIp: IpAddress
    dstPorts: seq[Port]

method cleanup*(puncher: TcpNutssInitiator) {.async.} =
  discard

proc initTcpNutssInitiator*(): TcpNutssInitiator =
  randomize()
  TcpNutssInitiator()

proc injectSynPackets(attempt: TNIAttempt) {.async.} =
  let injectFd = setupTcpInjectingSocket()
  for dstPort in attempt.dstPorts:
    let synOut = IpPacket(protocol: tcp, ipAddrSrc: attempt.srcIp,
                          ipAddrDst: attempt.dstIp, ipTTL: 2,
                          tcpPortSrc: attempt.srcPort, tcpPortDst: dstPort,
                          tcpSeqNumber: rand(uint32), tcpAckNumber: 0,
                          tcpFlags: {SYN}, tcpWindowSize: 1452 * 10)
    echo &"[{synOut.ipAddrSrc}:{synOut.tcpPortSrc} -> {synOut.ipAddrDst}:{synOut.tcpPortDst}, SEQ {synOut.tcpSeqNumber}] injecting outgoing SYN"
    await injectFd.injectTcpPacket(synOut)

proc accept(puncher: TcpNutssInitiator, ip: IpAddress, port: Port) {.async.} =
  let sock = newAsyncSocket()
  #sock.setSockOpt(OptReuseAddr, true)
  sock.bindAddr(port, $ip)
  sock.listen()
  while true:
    let acceptFuture = sock.accept()
    await acceptFuture or sleepAsync(Timeout)
    if acceptFuture.finished():
      let peer = acceptFuture.read()
      let (peerAddr, peerPort) = peer.getPeerAddr()
      let peerIp = parseIpAddress(peerAddr)
      let i = puncher.findAttempt(ip, port, peerIp, @[peerPort])
      if i == -1:
        echo "Accepted connection, but no attempt found. Discarding."
        peer.close()
        continue
      else:
        let attempt = TNIAttempt(puncher.attempts[i])
        attempt.future.complete(peer)
    let attempts = puncher.findAttemptsByLocalAddr(ip, port)
    if attempts.len() <= 1:
      break
  sock.close()

method initiate*(puncher: TcpNutssInitiator, args: string,
                 progressCb: PunchProgressCb): Future[AsyncSocket] {.async.} =
  let req = parseMessage[Request](args)
  let localIp = getPrimaryIPAddr(req.dstIp)
  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, req.srcPorts[0])
  if existingAttempts.len() == 0:
    echo &"initiating connection to {req.dstIp}:{req.dstPorts[0].int}"
    asyncCheck puncher.accept(localIp, req.srcPorts[0])
  else:
    for a in existingAttempts:
      if a.dstIp == req.dstIp and
         a.dstPorts.any(proc (p: Port): bool = p in req.dstPorts):
        raise newException(PunchHoleError, "hole punching for given parameters already active")
  try:
    let attempt = TNIAttempt(srcIp: localIp, srcPort: req.srcPorts[0],
                             dstIp: req.dstIp,
                             dstPorts: predictPortRange(req.dstPorts),
                             future: newFuture[AsyncSocket]("initiate"))
    puncher.attempts.add(attempt)
    await attempt.injectSynPackets()
    await progressCb("")
    await attempt.future or sleepAsync(Timeout)
    puncher.attempts.del(puncher.attempts.find(attempt))
    if attempt.future.finished():
      result = attempt.future.read()
    else:
      raise newException(PunchHoleError, "timeout")
  except OSError as e:
    raise newException(PunchHoleError, e.msg)