import asyncdispatch, asyncnet, os, sequtils, strformat, strutils, tables
from nativesockets import htonl
from net import IpAddress, Port, getPrimaryIPAddr, parseIpAddress, `$`
import ../../asyncutils
import ../../message
import ../../network_interface

type
  Client = object
    sock: AsyncSocket
    ip: IpAddress
    ports: array[3, Port]

  # Requests
  Register = object
    peerId: string
    ip: IpAddress
    ports: array[3, Port]

  GetPeerinfo = object
    peerId: string

  NotifyPeer = object
    sender: string
    recipient: string
    data: string

proc isPrivateIp(ip: IpAddress): bool =
  const ranges: array[5, tuple[first: IpAddress, last: IpAddress]] =
    [(parseIpAddress("10.0.0.0"), parseIpAddress("10.255.255.255")),
     (parseIpAddress("172.16.0.0"), parseIpAddress("172.31.255.255")),
     (parseIpAddress("192.168.0.0"), parseIpAddress("192.168.255.255")),
     (parseIpAddress("169.254.0.0"), parseIpAddress("169.254.255.255")),
     (parseIpAddress("127.0.0.0"), parseIpAddress("127.255.255.255"))]
  for r in ranges:
    let ipScalar = htonl(cast[uint32](ip.address_v4))
    if ipScalar > htonl(cast[uint32](r.first.address_v4)) and
       ipScalar < htonl(cast[uint32](r.last.address_v4)):
      return true
  return false

proc isInNetwork(ip: IpAddress, iface: NetworkInterface): bool =
  let ipScalar = htonl(cast[uint32](ip.address_v4))
  let ifaceIpScalar = htonl(cast[uint32](iface.ipAddress.address_v4))
  let netmaskScalar = htonl(cast[uint32](iface.netMask.address_v4))
  (ipScalar and netmaskScalar) == (ifaceIpScalar and netmaskScalar)

proc probePublicIp(): Future[IpAddress] {.async.} =
  let output = await asyncExecCmd("ping -R -c 1 -s 1 -n 193.0.14.129")
  let ipLines = output.splitLines()
    .filter(proc(l: string): bool = l.startsWith("\t") or l.startsWith("RR:"))
    .map(proc(l: string): string = l.strip(true, false, {'R', ':', '\t', ' '}))
  for line in ipLines:
    let ipAddr = parseIpAddress(line)
    if not isPrivateIp(ipAddr):
      return ipAddr
  block:
    raise newException(OSError, "cannot probe public IP address")

proc removeClient(clients: TableRef[string, Client], peerId: string) =
  if peerId.len > 0: clients.del(peerId)

proc sendEndpoint(client: AsyncSocket, requestId: string) {.async.} =
  let (address, port) = client.getPeerAddr()
  var ipAddr = parseIpAddress(address)
  if ipAddr.isPrivateIp() and
     ipAddr.isInNetwork(fromIpAddress(getPrimaryIPAddr(ipAddr))):
    ipAddr = await probePublicIp()
  await client.send(&"ok|{requestId}|{ipAddr}|{port.int}\n")
  client.close

proc processClient(client: AsyncSocket,
                   clients: TableRef[string, Client]) {.async.} =
  var id = ""
  var peerId = ""
  while true:
    var line = await client.recvLine(maxLength = 400)
    line = line.strip(leading = false, trailing = true, chars = {'\r', '\n'})
    if line.len == 0:
      removeClient(clients, peerId)
      break
    try:
      let args = line.parseArgs(3, 1)
      id = args[1]
      case args[0]:
      of "register":
        let req = parseMessage[Register](args[2])
        echo "register: ", req
        peerId = req.peerId
        clients[peerId] = Client(sock: client, ip: req.ip, ports: req.ports)
        asyncCheck client.send(&"ok|{id}\n")
      of "get-endpoint":
        echo "get-endpoint"
        asyncCheck client.sendEndpoint(id)
        removeClient(clients, peerId)
        break
      of "get-peerinfo":
        let req = parseMessage[GetPeerinfo](args[2])
        echo "get-info: ", req
        let peer = clients[req.peerId]
        let peerPorts = peer.ports.join(",")
        asyncCheck client.send(&"ok|{id}|{peer.ip}|{peerPorts}\n")
      of "notify-peer":
        let req = parseMessage[NotifyPeer](args[2])
        echo "notify-peer: ", req
        let recipient = clients[req.recipient]
        asyncCheck recipient.sock.send(&"notify-peer|{req.sender}|{req.recipient}|{req.data}\n")
        asyncCheck client.send(&"ok|{id}\n")
      else:
        echo "invalid request"
        client.close()
        removeClient(clients, peerId)
        break

    except KeyError:
      asyncCheck client.send(&"error|{id}|peer not registered\n")

    except ValueError:
      echo "invalid message"
      client.close
      removeClient(clients, peerId)
      break

proc serve(port: Port) {.async.} =
  var clients = newTable[string, Client]()
  var server = newAsyncSocket()
  server.setSockOpt(OptReuseAddr, true)
  server.bindAddr(port)
  server.listen()

  while true:
    let client = await server.accept()
    asyncCheck processClient(client, clients)

proc main() =
  if paramCount() != 1:
    echo(fmt"usage: {paramStr(0)} PORT")
    quit(1)

  try:
    let portNumber = paramStr(1).parseUInt
    if portNumber > uint16.high:
      raise newException(ValueError, "port out of range")
    let port = Port(portNumber)
    asyncCheck serve(port)
    runForever()

  except ValueError as e:
    echo e.msg

when isMainModule:
  main()