import asyncdispatch, asyncnet, os, strformat, strutils
from nativesockets import Domain, SockType, Protocol
from net import IpAddress, Port, `$`
import asyncutils
import message
import tcp_syni_initiator
import tcp_syni_responder
import tcp_nutss_initiator
import tcp_nutss_responder

from strutils import format, join
from nativesockets import setSockOptInt

type
  Punchd = ref object
    unixSocket: AsyncSocket
    tcpSyniInitiator: TcpSyniInitiator
    tcpSyniResponder: TcpSyniResponder
    tcpNutssInitiator: TcpNutssInitiator
    tcpNutssResponder: TcpNutssResponder

  Sigint = object of CatchableError

  # Requests
  InitiateTcpSyni = object
    srcIp: IpAddress
    srcPorts: seq[Port]
    dstIp: IpAddress
    dstPorts: seq[Port]

  RespondTcpSyni = object
    dstIp: IpAddress
    dstPorts: seq[Port]
    srcIp: IpAddress
    srcPorts: seq[Port]
    seqNums: seq[uint32]

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

  RespondTcpNutss = object
    dstIp: IpAddress
    dstPorts: seq[Port]
    srcIp: IpAddress
    srcPorts: seq[Port]
    extraData: string

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 handleRequest(punchd: Punchd, line: string,
                   unixSock: AsyncSocket) {.async.} =
  var id: string
  var sock: AsyncSocket
  try:
    let args = line.parseArgs(4)
    id = args[1]

    case args[0]:
    of "initiate":
      case args[2]:
      of "tcp-syni":
        let req = parseMessage[InitiateTcpSyni](args[3])
        proc progress(seqNumbers: seq[uint32]) {.async.} =
          echo "progress! seqNumbers: ", seqNumbers
          let content = @["tcp-syni", $req.srcIp, req.srcPorts.join(","),
                          $req.dstIp, req.dstPorts.join(","),
                          seqNumbers.join(",")].join("|")
          await sendToClient(unixSock, &"progress|{id}|{content}\n")
        sock = await punchd.tcpSyniInitiator.initiate(req.srcPorts[0], req.dstIp,
                                                      req.dstPorts, progress)

      of "tcp-nutss":
        let req = parseMessage[InitiateTcpNutss](args[3])
        proc progress() {.async.} =
          echo "progress!"
          let content = @["tcp-nutss", $req.srcIp, req.srcPorts.join(","),
                          $req.dstIp, req.dstPorts.join(","), ""].join("|")
          await sendToClient(unixSock, &"progress|{id}|{content}\n")
        sock = await punchd.tcpNutssInitiator.initiate(req.srcPorts[0], req.dstIp,
                                                       req.dstPorts, progress)

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


    of "respond":
      case args[2]:
      of "tcp-syni":
        let req = parseMessage[RespondTcpSyni](args[3])
        sock = await punchd.tcpSyniResponder.respond(req.srcPorts[0], req.dstIp,
                                                     req.dstPorts, req.seqNums)

      of "tcp-nutss":
        let req = parseMessage[RespondTcpNutss](args[3])
        sock = await punchd.tcpNutssResponder.respond(req.srcPorts[0],
                                                      req.dstIp, req.dstPorts)

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

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

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

  except PunchHoleError as e:
    await sendToClient(unixSock, &"error|{id}|{e.msg}\n")
  except 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,
                      tcpSyniInitiator: initTcpSyniInitiator(),
                      tcpSyniResponder: initTcpSyniResponder(),
                      tcpNutssInitiator: initTcpNutssInitiator(),
                      tcpNutssResponder: initTcpNutssResponder())
  asyncCheck handleUsers(punchd)
  try:
    runForever()

  except Sigint:
    waitFor punchd.tcpSyniInitiator.cleanup()
    waitFor punchd.tcpSyniResponder.cleanup()
    waitFor punchd.tcpNutssInitiator.cleanup()
    waitFor punchd.tcpNutssResponder.cleanup()
    punchd.unixSocket.close()
    removeFile(PunchdSocket)

when isMainModule:
  main()