import algorithm
import net
import sequtils
import unittest

type
  NatType* = enum
    Unknown,
    Cone,
    SymmetricProgressive,
    SymmetricRandom

  NatProperties* = object
    case natType*: NatType
    of Unknown:
      guess*: seq[uint16]
    of Cone:
      prediction*: uint16
    of SymmetricProgressive:
      order*: SortOrder
      previousPort*: uint16
      minDistance*: uint16
      maxDistance*: uint16
    of SymmetricRandom:
      minPort*: uint16
      maxPort*: uint16

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 getNatProperties*(localPort: uint16, probedPorts: seq[uint16]):
                      NatProperties =
  if probedPorts.len == 0:
    # No probed ports, so our only guess can be that the NAT is a cone-type NAT
    # and the port allocation preserves the local Port.
    return NatProperties(natType: Unknown, guess: @[localPort])
  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 allocation).
    result = NatProperties(natType: Unknown, guess: @[probedPorts[0]])
    if probedPorts[0] > localPort:
      let offset = probedPorts[0] - localPort
      result.guess.add(probedPorts[0].addOffset(offset))
    elif probedPorts[0] < localPort:
      let offset = localPort - probedPorts[0]
      result.guess.add(probedPorts[0].subtractOffset(offset))
    return
  let deduplicatedPorts = probedPorts.deduplicate()
  if deduplicatedPorts.len() == 1:
    # It looks like the NAT is a cone-type NAT.
    return NatProperties(natType: Cone, prediction: deduplicatedPorts[0])
  let probedPortsSorted = probedPorts.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 probedPorts.isSorted(Ascending):
      # assume symmetric NAT with positive-progressive port allocation
      return NatProperties(natType: SymmetricProgressive,
                           order: Ascending, previousPort: maxPort,
                           minDistance: minDistance, maxDistance: maxDistance)
    if probedPorts.isSorted(Descending):
      # assume symmetric NAT with negative-progressive port allocation
      return NatProperties(natType: SymmetricProgressive,
                           order: Descending, previousPort: minPort,
                           minDistance: minDistance, maxDistance: maxDistance)
  # assume symmetric NAT with random port allocation
  return NatProperties(natType: SymmetricRandom, minPort: minPort,
                       maxPort: maxPort)

proc getNatProperties*(localPort: Port, probedPorts: seq[Port]): NatProperties =
  getNatProperties(localPort.uint16, probedPorts.map(toUInt16))

proc predictPortRange*(props: NatProperties, maxPortCount: uint16): seq[Port] =
  case props.natType
  of Unknown:
    result = props.guess.map(toPort)
  of Cone:
    result = @[Port(props.prediction)]
  of SymmetricProgressive:
    if props.order == Ascending:
      if props.minDistance == props.maxDistance:
        return @[Port(props.previousPort.addOffset(props.maxDistance))]
      else:
        let minPort = props.previousPort - props.maxDistance
        for i in countup(0'u16, props.maxDistance):
          result.add(Port(minPort.addOffset(i)))
    else:
      if props.minDistance == props.maxDistance:
        return @[Port(props.previousPort.subtractOffset(props.maxDistance))]
      else:
        let maxPort = props.previousPort + props.maxDistance
        for i in countup(0'u16, props.maxDistance):
          result.add(Port(maxPort.subtractOffset(i)))
  of SymmetricRandom:
    assert(maxPortCount > 0)
    let center = props.minPort + (props.maxPort - props.minPort) div 2
    let half = maxPortCount div 2
    let first = if (1024'u16 + half) < center:
      min(center - half, uint16.high - maxPortCount + 1)
    else:
      1024'u16
    result = newSeq[Port](maxPortCount)
    for i in 0'u16 .. maxPortCount - 1'u16:
      result[i] = Port(first + i)

suite "port prediction tests":
  setup:
    let maxPortCount = 1000'u16

  test "single port":
    let props = getNatProperties(1234'u16, @[])
    check(props.natType == Unknown)
    check(props.guess == @[1234'u16])
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1234)])

  test "single probe equal":
    let props = getNatProperties(1234'u16, @[1234'u16])
    check(props.natType == Unknown)
    check(props.guess == @[1234'u16])
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1234)])

  test "single probe positive-progressive":
    let props = getNatProperties(1234'u16, @[1236'u16])
    check(props.natType == Unknown)
    check(props.guess == @[1236'u16, 1238'u16])
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1236), Port(1238)])

  test "single probe negative-progressive":
    let props = getNatProperties(1234'u16, @[1232'u16])
    check(props.natType == Unknown)
    check(props.guess == @[1232'u16, 1230'u16])
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1232), Port(1230)])

  test "all equal":
    let props = getNatProperties(1234'u16, @[1234'u16, 1234'u16])
    check(props.natType == Cone)
    check(props.prediction == 1234'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1234)])

  test "positive-progressive, offset 1":
    let props = getNatProperties(1234'u16, @[2034'u16, 2035'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Ascending)
    check(props.previousPort == 2035'u16)
    check(props.minDistance == 1'u16)
    check(props.maxDistance == 1'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(2036)])

  test "positive-progressive, offset 9":
    let props = getNatProperties(1234'u16, @[2034'u16, 2043'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Ascending)
    check(props.previousPort == 2043'u16)
    check(props.minDistance == 9'u16)
    check(props.maxDistance == 9'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(2052)])

  test "negative-progressive, offset 1":
    let props = getNatProperties(1234'u16, @[1100'u16, 1099'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Descending)
    check(props.previousPort == 1099'u16)
    check(props.minDistance == 1'u16)
    check(props.maxDistance == 1'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1098)])

  test "negative-progressive, offset 9":
    let props = getNatProperties(1234'u16, @[1100'u16, 1091'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Descending)
    check(props.previousPort == 1091'u16)
    check(props.minDistance == 9'u16)
    check(props.maxDistance == 9'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1082)])

  test "positive-progressive, 3 probed ports, low offset":
    let props = getNatProperties(1234'u16, @[2000'u16, 2000'u16, 2002'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Ascending)
    check(props.previousPort == 2002)
    check(props.minDistance == 0'u16)
    check(props.maxDistance == 2'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(2000), Port(2001), Port(2002)])

  test "negative-progressive, 3 probed ports, low offset":
    let props = getNatProperties(1234'u16, @[2002'u16, 2000'u16, 2000'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Descending)
    check(props.previousPort == 2000)
    check(props.minDistance == 0'u16)
    check(props.maxDistance == 2'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(2002), Port(2001), Port(2000)])

  test "high port, positive-progressive, offset 1":
    let props = getNatProperties(1234'u16, @[65534'u16, 65535'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Ascending)
    check(props.previousPort == 65535'u16)
    check(props.minDistance == 1'u16)
    check(props.maxDistance == 1'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1024)])

  test "high port, positive-progressive, offset 9":
    let props = getNatProperties(1234'u16, @[65520'u16, 65529'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Ascending)
    check(props.previousPort == 65529'u16)
    check(props.minDistance == 9'u16)
    check(props.maxDistance == 9'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(1026)])

  test "low port, negative-progressive, offset 1":
    let props = getNatProperties(1234'u16, @[1025'u16, 1024'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Descending)
    check(props.previousPort == 1024'u16)
    check(props.minDistance == 1'u16)
    check(props.maxDistance == 1'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(65535)])

  test "low port, negative-progressive, offset 9":
    let props = getNatProperties(1234'u16, @[1039'u16, 1030'u16])
    check(props.natType == SymmetricProgressive)
    check(props.order == Descending)
    check(props.previousPort == 1030'u16)
    check(props.minDistance == 9'u16)
    check(props.maxDistance == 9'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted == @[Port(65533)])

  test "random port allocation":
    let props = getNatProperties(1234'u16, @[20000'u16, 24000'u16])
    check(props.natType == SymmetricRandom)
    check(props.minPort == 20000'u16)
    check(props.maxPort == 24000'u16)
    let predicted = predictPortRange(props, maxPortCount)
    let half = maxPortCount div 2'u16
    check(predicted == toSeq(countup(22000'u16 - half, 22000'u16 + half - 1)).map(toPort))

  test "random port allocation":
    let props = getNatProperties(1234'u16, @[1200'u16, 1600'u16])
    check(props.natType == SymmetricRandom)
    check(props.minPort == 1200'u16)
    check(props.maxPort == 1600'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted.len == maxPortCount.int)
    check(predicted == toSeq(countup(1024'u16, 1024'u16 + maxPortCount - 1)).map(toPort))

  test "random port allocation, high":
    let props = getNatProperties(1234'u16, @[65000'u16, 65400'u16])
    check(props.natType == SymmetricRandom)
    check(props.minPort == 65000'u16)
    check(props.maxPort == 65400'u16)
    let predicted = predictPortRange(props, maxPortCount)
    check(predicted.len == maxPortCount.int)
    check(predicted == toSeq(countup(uint16.high - maxPortCount + 1, uint16.high)).map(toPort))