-- TODO: Scrap these hacks and write a proper test suite.

pcall(require, 'luarocks.require')

local socket = require 'socket'

local geoip = require 'geoip'
local geoip_country = require 'geoip.country'

local geoip_country_filename = select(1, ...) or "./GeoIP.dat"
local geoip_country6_filename = select(3, ...) or "./GeoIPv6.dat"

-- Since lua 5.2 unpack was moved to table.unpack
local unpack = unpack or table.unpack

print("TESTING lua-geoip")
print("")
print("VERSION: ", assert(geoip._VERSION))
print("DESCRIPTION: ", assert(geoip._DESCRIPTION))
print("COPYRIGHT: ", assert(geoip._COPYRIGHT))
print("")
print("VERSION: ", assert(geoip_country._VERSION))
print("DESCRIPTION: ", assert(geoip_country._DESCRIPTION))
print("COPYRIGHT: ", assert(geoip_country._COPYRIGHT))
print("")

-- Check that required files exist
-- See README on info on how to get them
assert(io.open(geoip_country_filename, "r")):close()
assert(io.open(geoip_country6_filename, "r")):close()

do
  local id = assert(geoip.id_by_code('RU'))

  assert(geoip.code_by_id(id) == 'RU')
  assert(geoip.code3_by_id(id) == 'RUS')
  assert(geoip.name_by_id(id) == 'Russian Federation')

  -- Depends on libgeoip version o_O
  assert(geoip.continent_by_id(id) == 'EU' or geoip.continent_by_id(id) == 'AS')

  assert(geoip.region_name_by_code('RU', '77') == "Tver'") -- WTF? MSK?
  assert(geoip.time_zone_by_country_and_region('RU', '77') == 'Europe/Moscow')
end

do
  assert(geoip_country.open("./BADFILENAME") == nil)

  assert(pcall(geoip_country.open, nil, geoip.INDEX_CACHE) == false)

  --assert(geoip_country.open(nil, 2 ^ 10) == nil) -- TODO: This should fail
  --assert(geoip_country.open(nil, nil, -1) == nil) -- TODO: This should fail
end

do
  local flags =
  {
    geoip.STANDARD;
    geoip.MEMORY_CACHE;
    geoip.CHECK_CACHE;
    geoip.INDEX_CACHE;
    geoip.MMAP_CACHE;
  }

  for _, flag in ipairs(flags) do
    if flag ~= geoip.INDEX_CACHE then
      assert(geoip_country.open(nil, flag)):close()
      assert(geoip_country.open(geoip_country_filename, flag)):close()
    end
  end
end

do
  local geodb = assert(
      geoip_country.open(geoip_country_filename)
    )
  geodb:close()
  geodb:close()
end

do
  local check_country = function(db, method, arg)
    local id = assert(db[method](db, arg, "id"))
    assert(type(id) == "number")

    local expected =
    {
      id = id;
      code = assert(geoip.code_by_id(id));
      code3 = assert(geoip.code3_by_id(id));
      name = assert(geoip.name_by_id(id));
      continent = assert(geoip.continent_by_id(id));
    }

    local all = assert(db[method](db, arg))

    local keys = { }
    for k, v in pairs(expected) do
      assert(all[k] == expected[k])
      assert(db[method](db, arg, k) == expected[k])
      keys[#keys + 1] = k
    end

    local r = { db[method](db, arg, unpack(keys)) }
    assert(#r == #keys)
    for i = 1, #keys do
      assert(r[i] == expected[keys[i]])
    end
  end

  local apack = function(...)
    return select("#", ...), { ... }
  end

  local geodb_country = assert(geoip_country.open(geoip_country_filename))

  local checkers =
  {
    [geodb_country] = check_country;
  }

  for _, geodb in ipairs { geodb_country } do
    local checker = checkers[geodb]

--    checker(geodb, "query_by_name", "google-public-dns-a.google.com")
    checker(geodb, "query_by_addr", "8.8.8.8")
    checker(geodb, "query_by_ipnum", 134744072) -- 8.8.8.8
  end

  geodb_country:close()
end

-- Country IPv6 Edition
do
  local geodb_country6 = assert(geoip_country.open(geoip_country6_filename, geoip.MEMORY_CACHE, geoip.COUNTRY_V6))

  res = geodb_country6:query_by_addr6("2a01:e0c:1::1")
  assert(res.code == "FR")

  res = geodb_country6:query_by_addr6("2a03:2880:f127:83:face:b00c:0:25de")
  assert(res.code == "US")
  geodb_country6:close()
end

-- TODO: Test two different DBs open in parallel work properly

local profiles =
{
  {
    name = "country";
    module = geoip_country;
    file = geoip_country_filename;
    field = "id";
  };
}

for i = 1, #profiles do
  local p = profiles[i]

  local geodb = assert(p.module.open(p.file))

  do
    print(p.name, "profiling ipnum queries")

    local num_queries = 1e5

    local cases = { }
    for i = 1, num_queries do
      cases[i] = math.random(0x7FFFFFFF)
    end

    local time_start = socket.gettime()
    for i = 1, num_queries do
      if i % 1e4 == 0 then
        print("#", i, "of", num_queries)
      end
      local result, err = geodb:query_by_ipnum(cases[i], p.field)
      if not result and err ~= "not found" then
        error(err)
      end
    end

    print(
        p.name,
        num_queries / (socket.gettime() - time_start),
        "ipnum queries per second"
      )
    print()
  end

  do
    print(p.name, "profiling addr queries") -- slow due to dns resolution

    local num_queries = 1e5

    local cases = { }
    for i = 1, num_queries do
      cases[i] = ('%d.%d.%d.%d'):format(
          math.random(255),
          math.random(255),
          math.random(255),
          math.random(255)
        )
    end


    local time_start = socket.gettime()
    for i = 1, num_queries do
      if i % 1e4 == 0 then
        print("#", i, "of", num_queries)
      end
      local result, err = geodb:query_by_name(cases[i], p.field)
      if not result and err ~= "not found" then
        error(err)
      end
    end

    print(
        p.name,
        num_queries / (socket.gettime() - time_start),
        "addr queries per second"
      )
    print()
  end

--[[ Disable test: access to network
  do
    print(p.name, "profiling name queries")

    local num_queries = 500 -- slow due to dns resolution

    local time_start = socket.gettime()
    for i = 1, num_queries do
      if i % 50 == 0 then
        print("#", i, "of", num_queries)
      end
      assert(geodb:query_by_name("ya.ru", p.field))
    end

    print(
        p.name,
        num_queries / (socket.gettime() - time_start),
        "name queries per second"
      )
    print()
  end
]]--

  geodb:close()
end

print("")
print("OK")