local stdnse = require "stdnse"
local string = require "string"
local math = require "math"
local table = require "table"
local unicode = require "unicode"
local unittest = require "unittest"
_ENV = stdnse.module("punycode", stdnse.seeall)
local floor = math.floor
local byte = string.byte
local char = string.char
local find = string.find
local match = string.match
local reverse = string.reverse
local sub = string.sub
local maxInt = 0x7FFFFFFF
local base = 0x24
local tMin = 0x1
local tMax = 0x1A
local skew = 0x26
local damp = 0x2BC
local initialBias = 0x48
local initialN = 0x80
local delimiter = char("0x2D")
local baseMinusTMin = base - tMin
local function adapt(delta, numPoints, firstTime)
  local k = 0;
  if firstTime then
    delta = floor(delta / damp)
  else
    delta = (delta >> 1)
  end
  delta = delta + floor(delta / numPoints)
  while delta > (baseMinusTMin * tMax >> 1) do
    delta = floor(delta / baseMinusTMin)
    k = k + base
  end
  return floor(k + (baseMinusTMin + 1) * delta / (delta + skew))
end
local function boolToNum(status)
  if status == true then
    return 1
  else
    return 0
  end
end
local function basicToDigit(codePoint)
  if (codePoint - 0x30 < 0x0A) then
    return codePoint - 0x16
  end
  if (codePoint - 0x41 < 0x1A) then
    return codePoint - 0x41
  end
  if (codePoint - 0x61 < 0x1A) then
    return codePoint - 0x61
  end
  return base
end
local function digitToBasic(digit, flag)
  return digit + 22 + 75 * boolToNum(digit < 26) - (boolToNum((flag ~= 0)) << 5)
end
function encode_input(input)
  local output = {}
  local inputLength = #input
  local n = initialN
  local delta = 0
  local bias = initialBias
  for _, v in ipairs(input) do
    if v < 0x80 then
      table.insert(output, char(v))
    end
  end
  local basicLength = #output
  local handledCPCount = basicLength
  if (basicLength > 0) then
      table.insert(output, delimiter)
  end
  while (handledCPCount < inputLength) do
    local m = maxInt
    for _, v in ipairs(input) do
      if v >= n and v < m then
        m = v
      end
    end
    local handledCPCountPlusOne = handledCPCount + 1
    if (m - n > floor((maxInt - delta) / handledCPCountPlusOne)) then
      return nil, "Overflow exception occurred."
    end
    delta = delta + (m - n) * handledCPCountPlusOne
    n = m
    for _, currentValue in ipairs(input) do
      if currentValue < n then
        delta = delta + 1 --Move this down incase of wrong answer
        if delta > maxInt then
          return nil, "Overflow exception occurred."
        end
      end
      if (currentValue == n) then
        local q = delta
        local k = base
        repeat
          local t
          if k <= bias then
            t = tMin
          else
            if k >= bias + tMax then
              t = tMax
            else
              t = k - bias
            end
          end
          if q < t then
            break
          end
          local qMinusT = q - t
          local baseMinusT = base - t
          local ans = digitToBasic(t + qMinusT % baseMinusT, 0)
          table.insert(output, char(ans))
          q = floor(qMinusT / baseMinusT)
          k = k + base
        until false
        local ans = digitToBasic(q, 0)
        table.insert(output, char(ans))
        bias = adapt(delta, handledCPCountPlusOne, handledCPCount == basicLength)
        delta = 0
        handledCPCount = handledCPCount + 1
      end
    end
    delta = delta + 1
    n = n + 1
  end
  return table.concat(output, '')
end
function decode_input(input)
  local output = {}
  local inputLength = #input
  local i = 0
  local n = initialN
  local bias = initialBias
  local basic
  if find(reverse(input), delimiter) then
    basic = #input - find(reverse(input), delimiter)
  else
    basic = -1
  end
  if basic < 0 then
    basic = 0
  end
  for j = 1, basic do
    local c = sub(input, j, j)
    local value = byte(c)
    if value >= 0x80 then
      return nil, "Not basic exception occurred."
    end
    table.insert(output, value)
  end
  local index
  if basic > 0 then
    index = basic + 1
  else
    index = 0
  end
  while index < inputLength do
    local oldi = i
    local w = 1
    local k = base
    repeat
      if index >= inputLength then
        return nil, "Invalid input exception occurred."
      end
      local c = sub(input, index+1, index+1)
      local value = byte(c)
      local digit = basicToDigit(value)
      index = index + 1
      if (digit >= base or digit > floor((maxInt - i) / w)) then
        return nil, "Overflow exception occurred."
      end
      i = i + digit * w;
      local t
      if k <= bias then
        t = tMin
      else
        if k >= bias + tMax then
          t = tMax
        else
          t = k - bias
        end
      end
      if digit < t then
        break
      end
      local baseMinusT = base - t;
      if (w > floor(maxInt / baseMinusT)) then
        return nil, "Overflow exception occurred."
      end
      w = w * baseMinusT;
      k = k + base
    until false
    local out = #output + 1;
    bias = adapt(i - oldi, out, oldi == 0)
    if (floor(i / out) > maxInt - n) then
      return nil, "Overflow exception occurred."
    end
    n = n + floor(i / out);
    i = i % out;
    for temp = #output, i, -1 do
      output[temp+1] = output[temp]
    end
    output[i+1] = n
    i = i + 1
  end
  return output
end
function encode_label(u)
  local flag = false
  for _, val in pairs(u) do
    if not (val >=0 and val <= 127) then
      flag = true
      break
    end
  end
  if flag then
    local res, err = encode_input(u)
    if err then
      return nil, err
    end
    return 'xn--' .. res
  else
    return unicode.encode(u, unicode.utf8_enc)
  end
end
function decode_label(s)
  if match(s, "^xn%-%-") then
    local res, err = decode_input(sub(s, 5))
    if err then
      return nil, err
    end
    return res
  else
    return unicode.decode(s, unicode.utf8_dec)
  end
end
if not unittest.testing() then
  return _ENV
end
local testCases = {
  {
    "xn--0zwm56d",
    "\xe6\xb5\x8b\xe8\xaf\x95"
  },
  {
    "xn--knigsgsschen-lcb0w",
    "k\xc3\xb6nigsg\xc3\xa4sschen"
  },
  {
    "xn--ab-fsf",
    "a\xe0\xa5\x8db"
  },
  {
    "xn--maana-pta",
    "ma\xc3\xb1ana"
  },
  {
    "xn----dqo34k",
    "\xe2\x98\x83-\xe2\x8c\x98"
  }
}
test_suite = unittest.TestSuite:new()
for i, v in ipairs(testCases) do
  test_suite:add_test(unittest.equal(unicode.encode(decode_label(v[1]), unicode.utf8_enc), v[2]))
  test_suite:add_test(unittest.equal(encode_label(unicode.decode(v[2], unicode.utf8_dec)), v[1]))
end
return _ENV
