import algorithm
import net
import sequtils
import unittest

const RandomPortCount = 1000

proc min(a, b: uint16): uint16 =
  min(a.int32, b.int32).uint16

proc toUint16(p: Port): uint16 = uint16(p)

proc toPort(u: uint16): Port = Port(u)

proc addOffset(port: uint16, offset: uint16, minValue = 1024'u16,
               maxValue = uint16.high): uint16 =
  assert(port >= minValue)
  assert(port <= maxValue)
  let distanceToMaxValue = maxValue - port
  if distanceToMaxValue < offset:
    return minValue + offset - distanceToMaxValue - 1
  return port + offset

proc subtractOffset(port: uint16, offset: uint16, minValue = 1024'u16,
                    maxValue = uint16.high): uint16 =
  assert(port >= minValue)
  assert(port <= maxValue)
  let distanceToMinValue = port - minValue
  if distanceToMinValue < offset:
    return maxValue - offset + distanceToMinValue + 1
  return port - offset

proc predictPortRange*(localPort: Port, probedPorts: seq[Port]): seq[Port] =
  if probedPorts.len == 0:
    # No probed ports, so our only guess can be that the NAT is a cone-type NAT
    # and the port mapping preserves the local Port.
    return @[localPort]
  let localPortUint = localPort.uint16
  let probedPortsUint = probedPorts.map(toUint16)
  if probedPorts.len == 1:
    # Only one server was used for probing, so we cannot know if the NAT is
    # symmetric or not. We are trying the probed port (assuming cone-type NAT)
    # and the next port in a progressive sequence if applicable (assuming
    # symmetric NAT with progressive port mapping).
    result.add(probedPorts[0])
    if probedPortsUint[0] > localPortUint:
      let offset = probedPortsUint[0] - localPortUint
      result.add(Port(probedPortsUint[0].addOffset(offset)))
    elif probedPortsUint[0] < localPortUint:
      let offset = localPortUint - probedPortsUint[0]
      result.add(Port(probedPortsUint[0].subtractOffset(offset)))
    return
  let deduplicatedPorts = probedPortsUint.deduplicate()
  if deduplicatedPorts.len() == 1:
    # It looks like the NAT is a cone-type NAT.
    return deduplicatedPorts.map(toPort)
  let probedPortsSorted = probedPortsUint.sorted()
  let minPort = probedPortsSorted[probedPortsSorted.minIndex()]
  let maxPort = probedPortsSorted[probedPortsSorted.maxIndex()]
  var minDistance = uint16.high()
  var maxDistance = uint16.low()
  for i in 1 .. probedPortsSorted.len() - 1:
    # FIXME: use rotated distance
    let distance = probedPortsSorted[i] - probedPortsSorted[i - 1]
    minDistance = min(minDistance, distance)
    maxDistance = max(maxDistance, distance)
  if maxDistance < 10:
    if probedPortsUint.isSorted(Ascending):
      # assume symmetric NAT with positive-progressive port mapping
      if minDistance == maxDistance:
        return @[Port(maxPort.addOffset(maxDistance))]
      else:
        for i in countup(0'u16, maxDistance):
          result.add(Port(minPort.addOffset(i)))
        return
    if probedPortsUint.isSorted(Descending):
      # assume symmetric NAT with negative-progressive port mapping
      if minDistance == maxDistance:
        return @[Port(minPort.subtractOffset(maxDistance))]
      else:
        for i in countup(0'u16, maxDistance):
          result.add(Port(maxPort.subtractOffset(i)))
        return
  # assume symmetric NAT with random port mapping
  let portRange = maxPort - minPort
  let first = if portRange > RandomPortCount:
    minPort
  else:
    let notCovered = RandomPortCount - portRange
    max(minPort - notCovered shr 1, 1024)
  let last = first + RandomPortCount
  for i in first .. last:
    result.add(Port(i))

suite "port prediction tests":
  test "single port":
    let predicted = predictPortRange(Port(1234), @[])
    check(predicted == @[Port(1234)])

  test "single probe equal":
    let predicted = predictPortRange(Port(1234), @[Port(1234)])
    check(predicted == @[Port(1234)])

  test "single probe positive-progressive":
    let predicted = predictPortRange(Port(1234), @[Port(1236)])
    check(predicted == @[Port(1236), Port(1238)])

  test "single probe negative-progressive":
    let predicted = predictPortRange(Port(1234), @[Port(1232)])
    check(predicted == @[Port(1232), Port(1230)])

  test "all equal":
    let predicted = predictPortRange(Port(1234), @[Port(1234), Port(1234)])
    check(predicted == @[Port(1234)])

  test "positive-progressive, offset 1":
    let predicted = predictPortRange(Port(1234), @[Port(2034), Port(2035)])
    check(predicted == @[Port(2036)])

  test "positive-progressive, offset 9":
    let predicted = predictPortRange(Port(1234), @[Port(2034), Port(2043)])
    check(predicted == @[Port(2052)])

  test "negative-progressive, offset 1":
    let predicted = predictPortRange(Port(1234), @[Port(1100), Port(1099)])
    check(predicted == @[Port(1098)])

  test "negative-progressive, offset 9":
    let predicted = predictPortRange(Port(1234), @[Port(1100), Port(1091)])
    check(predicted == @[Port(1082)])

  test "positive-progressive, 3 probed ports, low offset":
    let predicted = predictPortRange(Port(1234), @[Port(2000), Port(2000), Port(2002)])
    check(predicted == @[Port(2000), Port(2001), Port(2002)])

  test "negative-progressive, 3 probed ports, low offset":
    let predicted = predictPortRange(Port(1234), @[Port(2002), Port(2000), Port(2000)])
    check(predicted == @[Port(2002), Port(2001), Port(2000)])

  test "high port, positive-progressive, offset 1":
    let predicted = predictPortRange(Port(1234), @[Port(65534), Port(65535)])
    check(predicted == @[Port(1024)])

  test "high port, positive-progressive, offset 9":
    let predicted = predictPortRange(Port(1234), @[Port(65520), Port(65529)])
    check(predicted == @[Port(1026)])

  test "low port, negative-progressive, offset 1":
    let predicted = predictPortRange(Port(1234), @[Port(1025), Port(1024)])
    check(predicted == @[Port(65535)])

  test "low port, negative-progressive, offset 9":
    let predicted = predictPortRange(Port(1234), @[Port(1039), Port(1030)])
    check(predicted == @[Port(65533)])

  test "random mapping, distance > RandomPortCount":
    let predicted = predictPortRange(Port(1234), @[Port(3546), Port(7624)])
    check(predicted == toSeq(countup(3546'u16, 3546'u16 + RandomPortCount)).map(toPort))

  test "random mapping, distance < RandomPortCount":
    let centerPort = 30000'u16
    let minPort = centerPort - RandomPortCount.uint16 shr 1 + 1
    let maxPort = centerPort + RandomPortCount.uint16 shr 1 - 1
    let predicted = predictPortRange(Port(centerPort), @[Port(minPort), Port(maxPort)])
    check(predicted == toSeq(countup(minPort - 1, maxPort + 1)).map(toPort))