from nativesockets import Domain, SockType, Protocol
from net import IpAddress, Port, `$`, `==`, parseIpAddress
from sequtils import filter
import asyncdispatch, asyncnet, os, strformat, strutils
import asyncutils
import message
import options
import tables
import tcp_syni
import tcp_nutss
import udp

type
  Punchd = ref object
    unixSocket: AsyncSocket
    punchers: Table[string, Puncher]
    attempts: seq[Attempt]

  Sigint = object of CatchableError

const PunchdSocket = "/tmp/punchd.socket"

proc handleSigint() {.noconv.} =
  raise newException(Sigint, "received SIGINT")

proc sendToClient(unixSock: AsyncSocket, msg: string,
                  cmsgs: seq[ControlMessage] = @[]) {.async.} =
  if not unixSock.isClosed():
    let unixFd = unixSock.getFd.AsyncFD
    await unixFd.asyncSendMsg(msg, cmsgs)

proc findAttemptsByLocalAddr(punchd: Punchd, protocol: Protocol,
                             srcIp: IpAddress, srcPort: Port): seq[Attempt] =
  proc matchesLocalAddr(a: Attempt): bool =
    a.protocol == protocol and a.srcIp == srcIp and a.srcPort == srcPort
  punchd.attempts.filter(matchesLocalAddr)

proc acceptConnections(punchd: Punchd, ip: IpAddress, port: Port,
                       protocol: Protocol) {.async.} =
  var sockType: SockType
  case protocol:
  of IPPROTO_TCP:
    sockType = SOCK_STREAM
  of IPPROTO_UDP:
    sockType = SOCK_DGRAM
  else:
    assert(false, "can only accept TCP or UDP connections")
  let sock = newAsyncSocket(sockType = sockType, protocol = protocol)
  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 query = Attempt(protocol: protocol, srcIp: ip, srcPort: port,
                          dstIp: peerIp, dstPorts: @[peerPort])
      let i = punchd.attempts.find(query)
      if i == -1:
        echo "Accepted connection, but no attempt found. Discarding."
        peer.close()
        continue
      else:
        let acceptFuture = punchd.attempts[i].acceptFuture.get()
        acceptFuture.complete(peer)
    let localAddrMatches = punchd.findAttemptsByLocalAddr(protocol, ip, port)
    if localAddrMatches.len() <= 1:
      break
  sock.close()

proc addAttempt(punchd: Punchd, attempt: Attempt) =
  let localAddrMatches = punchd.findAttemptsByLocalAddr(attempt.protocol,
                                                        attempt.srcIp,
                                                        attempt.srcPort)
  punchd.attempts.add(attempt)
  if localAddrMatches.len() == 0:
    if attempt.acceptFuture.isSome():
      asyncCheck punchd.acceptConnections(attempt.srcIp, attempt.srcPort,
                                          attempt.protocol)
  elif localAddrMatches.contains(attempt):
    raise newException(PunchHoleError,
                       "hole punching for given parameters already active")

proc removeAttempt(punchd: Punchd, attempt: Attempt) =
  punchd.attempts.del(punchd.attempts.find(attempt))

proc handleRequest(punchd: Punchd, line: string,
                   unixSock: AsyncSocket) {.async.} =
  var id: string
  var sock: AsyncSocket
  var attempt: Attempt
  try:
    let args = line.parseArgs(4)
    id = args[1]
    let puncher = punchd.punchers[args[2]]

    case args[0]:
    of "initiate":
      attempt = puncher.parseInitiateRequest(args[3])
      punchd.addAttempt(attempt)
      proc progress(extraArgs: string) {.async.} =
        let msg = &"progress|{id}|{args[2]}|{args[3]}|{extraArgs}\n"
        await sendToClient(unixSock, msg)
      sock = await puncher.initiate(attempt, progress)
      punchd.removeAttempt(attempt)

    of "respond":
      attempt = puncher.parseRespondRequest(args[3])
      punchd.addAttempt(attempt)
      sock = await puncher.respond(attempt)
      punchd.removeAttempt(attempt)

    else:
      raise newException(ValueError, "invalid request")

    await sendToClient(unixSock, &"ok|{id}\n", @[fromFd(sock.getFd.AsyncFD)])
    sock.close()

  except PunchHoleError as e:
    punchd.removeAttempt(attempt)
    await sendToClient(unixSock, &"error|{id}|{e.msg}\n")
  except KeyError, ValueError:
    unixSock.close

proc handleRequests(punchd: Punchd, userSock: AsyncSocket) {.async.} =
  while true:
    if userSock.isClosed:
      break
    let line = await userSock.recvLine(maxLength = 400)
    if line.len == 0:
      userSock.close()
      break
    asyncCheck punchd.handleRequest(line, userSock)

proc handleUsers(punchd: Punchd) {.async.} =
  while true:
    let user = await punchd.unixSocket.accept()
    asyncCheck punchd.handleRequests(user)

proc main() =
  setControlCHook(handleSigint)
  removeFile(PunchdSocket)
  let unixSocket = newAsyncSocket(AF_UNIX, SOCK_STREAM, IPPROTO_IP)
  unixSocket.bindUnix(PunchdSocket)
  unixSocket.listen()
  setFilePermissions(PunchdSocket,
                     {fpUserRead, fpUserWrite, fpGroupRead, fpGroupWrite,
                      fpOthersRead, fpOthersWrite})
  let punchd = Punchd(unixSocket: unixSocket)
  punchd.punchers["tcp-syni"] = initTcpSyniPuncher()
  punchd.punchers["tcp-nutss"] = initTcpNutssPuncher()
  punchd.punchers["udp"] = initUdpPuncher()
  asyncCheck handleUsers(punchd)
  try:
    runForever()

  except Sigint:
    while punchd.attempts.len() != 0:
      waitFor punchd.attempts.pop().cleanup()
    punchd.unixSocket.close()
    removeFile(PunchdSocket)

when isMainModule:
  main()