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

export PunchHoleError

type
  Attempt = object
    srcIp: IpAddress
    srcPort: Port
    dstIp: IpAddress
    dstPorts: seq[Port]
    future: Future[AsyncSocket]

  TcpNutssInitiator* = Puncher[Attempt]

  PunchProgressCb* = proc() {.async.}

proc cleanup*(puncher: TcpNutssInitiator) {.async.} =
  while puncher.attempts.len() != 0:
    await puncher.attempts.pop().deleteFirewallRules() # FIXME: needed?

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

proc injectSynPackets(attempt: Attempt) {.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 = puncher.attempts[i]
        attempt.future.complete(peer)
    let attempts = puncher.findAttemptsByLocalAddr(ip, port)
    if attempts.len() <= 1:
      break
  sock.close()

proc initiate*(puncher: TcpNutssInitiator, srcPort: Port, dstIp: IpAddress,
               dstPorts: seq[Port], progressCb: PunchProgressCb):
               Future[AsyncSocket] {.async.} =
  let localIp = getPrimaryIPAddr(dstIp)
  let existingAttempts = puncher.findAttemptsByLocalAddr(localIp, srcPort)
  if existingAttempts.len() == 0:
    echo &"initiating connection to {dstIp}:{dstPorts[0].int}"
    asyncCheck puncher.accept(localIp, srcPort)
  else:
    for a in existingAttempts:
      if a.dstIp == dstIp and
         a.dstPorts.any(proc (p: Port): bool = p in dstPorts):
        raise newException(PunchHoleError, "hole punching for given parameters already active")
  try:
    let attempt = Attempt(srcIp: localIp, srcPort: srcPort, dstIp: dstIp,
                          dstPorts: predictPortRange(dstPorts),
                          future: newFuture[AsyncSocket]("initiate"))
    puncher.attempts.add(attempt)
    #await attempt.addFirewallRules()
    await attempt.injectSynPackets()
    await progressCb()
    await attempt.future or sleepAsync(Timeout)
    #await attempt.deleteFirewallRules()
    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)