--------------------------------------------------------------------------------
--- Small table utilities
-- @module lua-nucleo.table-utils
-- This file is a part of lua-nucleo library
-- @copyright lua-nucleo authors (see file `COPYRIGHT` for the license)
--------------------------------------------------------------------------------

local setmetatable, error, pairs, ipairs, tostring, select, type, assert
    = setmetatable, error, pairs, ipairs, tostring, select, type, assert

local rawget = rawget

local table_insert, table_remove = table.insert, table.remove

local math_min, math_max = math.min, math.max

--------------------------------------------------------------------------------

local arguments,
      optional_arguments,
      method_arguments
      = import 'lua-nucleo/args.lua'
      {
        'arguments',
        'optional_arguments',
        'method_arguments'
      }

local is_number,
      is_table
      = import 'lua-nucleo/type.lua'
      {
        'is_number',
        'is_table'
      }

local assert_is_table
      = import 'lua-nucleo/typeassert.lua'
      {
        'assert_is_table'
      }

--------------------------------------------------------------------------------

-- Warning: it is possible to corrupt this with rawset and debug.setmetatable.
local empty_table = setmetatable(
    { },
    {
      __newindex = function(t, k, v)
        error("attempted to change the empty table", 2)
      end;

      __metatable = "empty_table";
    }
  )

local function toverride_many(t, s, ...)
  if s then
    for k, v in pairs(s) do
      t[k] = v
    end
    -- Recursion is usually faster than calling select()
    return toverride_many(t, ...)
  end

  return t
end

local function tappend_many(t, s, ...)
  if s then
    for k, v in pairs(s) do
      if t[k] == nil then
        t[k] = v
      else
        error("attempted to override table key `" .. tostring(k) .. "'", 2)
      end
    end

    -- Recursion is usually faster than calling select()
    return tappend_many(t, ...)
  end

  return t
end

local function tijoin_many(t, s, ...)
  if s then
    -- Note: can't use ipairs() since we want to support tijoin_many(t, t)
    for i = 1, #s do
      t[#t + 1] = s[i]
    end

    -- Recursion is usually faster than calling select()
    return tijoin_many(t, ...)
  end

  return t
end

-- Keys are ordered in undetermined order
local tkeys = function(t)
  local r = { }

  for k, v in pairs(t) do
    r[#r + 1] = k
  end

  return r
end

-- Values are ordered in undetermined order
local tvalues = function(t)
  local r = { }

  for k, v in pairs(t) do
    r[#r + 1] = v
  end

  return r
end

-- Keys and values are ordered in undetermined order
local tkeysvalues = function(t)
  local keys, values = { }, { }

  for k, v in pairs(t) do
    keys[#keys + 1] = k
    values[#values + 1] = v
  end

  return keys, values
end

-- If table contains multiple keys with the same value,
-- only one key is stored in the result, picked in undetermined way.
local tflip = function(t)
  local r = { }

  for k, v in pairs(t) do
    r[v] = k
  end

  return r
end

-- If table contains multiple keys with the same value,
-- only one key is stored in the result, picked in undetermined way.
local tflip_inplace = function(t)
  for k, v in pairs(t) do
    t[v] = k
  end

  return t
end

-- If table contains multiple keys with the same value,
-- only the last such key (highest one) is stored in the result.
local tiflip = function(t)
  local r = { }

  for i = 1, #t do
    r[t[i]] = i
  end

  return r
end

local tset = function(t)
  local r = { }

  for k, v in pairs(t) do
    r[v] = true
  end

  return r
end

local tiset = function(t)
  local r = { }

  for i = 1, #t do
    r[t[i]] = true
  end

  return r
end

local function tiinsert_args(t, a, ...)
  if a ~= nil then
    t[#t + 1] = a
    -- Recursion is usually faster than calling select() in a loop.
    return tiinsert_args(t, ...)
  end

  return t
end

local timap_inplace = function(fn, t, ...)
  for i = 1, #t do
    t[i] = fn(t[i], ...)
  end

  return t
end

local timap = function(fn, t, ...)
  local r = { }
  for i = 1, #t do
    r[i] = fn(t[i], ...)
  end
  return r
end

local timap_sliding = function(fn, t, ...)
  local r = {}

  for i = 1, #t do
    tiinsert_args(r, fn(t[i], ...))
  end

  return r
end

local tiwalk = function(fn, t, ...)
  for i = 1, #t do
    fn(t[i], ...)
  end
end

local tiwalker = function(fn)
  return function(t)
    for i = 1, #t do
      fn(t[i])
    end
  end
end

local twalk_pairs = function(fn, t)
  for k, v in pairs(t) do
    fn(k, v)
  end
end

local tequals = function(lhs, rhs)
  for k, v in pairs(lhs) do
    if v ~= rhs[k] then
      return false
    end
  end

  for k, v in pairs(rhs) do
    if lhs[k] == nil then
      return false
    end
  end

  return true
end

local tiunique = function(t)
  return tkeys(tiflip(t))
end

-- Deprecated, use tgenerate_1d_linear instead
local tgenerate_n = function(n, generator, ...)
  local r = { }
  for i = 1, n do
    r[i] = generator(...)
  end
  return r
end

local tgenerate_1d_linear = function(n, fn, ...)
  local r = { }
  for i = 1, n do
    r[#r + 1] = fn(i, ...)
  end
  return r
end

local tgenerate_2d_linear = function(w, h, fn, ...)
  local r = { }
  for y = 1, h do
    for x = 1, w do
      r[#r + 1] = fn(x, y, ...)
    end
  end
  return r
end

local taccumulate = function(t, init)
  local sum = init or 0
  for k, v in pairs(t) do
    sum = sum + v
  end
  return sum
end

local tnormalize, tnormalize_inplace
do
  local impl = function(t, r, sum)
    sum = sum or taccumulate(t)

    for k, v in pairs(t) do
      r[k] = v / sum
    end

    return r
  end

  tnormalize = function(t, sum)
    return impl(t, { }, sum)
  end

  tnormalize_inplace = function(t, sum)
    return impl(t, t, sum)
  end
end

local tclone
do
  local function impl(t, visited)
    local t_type = type(t)
    if t_type ~= "table" then
      return t
    end

    assert(not visited[t], "recursion detected")
    visited[t] = true

    local r = { }
    for k, v in pairs(t) do
      r[impl(k, visited)] = impl(v, visited)
    end

    visited[t] = nil

    return r
  end

  tclone = function(t)
    return impl(t, { })
  end
end

-- Slow
local tcount_elements = function(t)
  local n = 0
  for _ in pairs(t) do
    n = n + 1
  end
  return n
end

local tremap_to_array = function(fn, t)
  local r = { }
  for k, v in pairs(t) do
    r[#r + 1] = fn(k, v)
  end
  return r
end

local tmap_values = function(fn, t, ...)
  local r = { }
  for k, v in pairs(t) do
    r[k] = fn(v, ...)
  end
  return r
end

--------------------------------------------------------------------------------

local torderedset = function(t)
  local r = { }

  for i = 1, #t do
    local v = t[i]

    -- Have to add this limitation to avoid size ambiquity.
    -- If you need ordered set of numbers, use separate storage
    -- for set and array parts (write make_ordered_set then).
    assert(type(v) ~= "number", "can't insert number into ordered set")

    r[v] = i
    r[i] = v
  end

  return r
end

-- Returns false if item already exists
-- Returns true otherwise
local torderedset_insert = function(t, v)
  -- See torderedset() for motivation
  assert(type(v) ~= "number", "can't insert number into ordered set")

  if not t[v] then
    local i = #t + 1
    t[v] = i
    t[i] = v

    return true
  end

  return false
end

-- Returns false if item didn't existed
-- Returns true otherwise
-- Note this operation is really slow
local torderedset_remove = function(t, v)
  -- See torderedset() for motivation
  assert(type(v) ~= "number", "can't remove number from ordered set")

  local pos = t[v]
  if pos then
    t[v] = nil
    -- TODO: Do table.remove manually then to do all in a single loop.
    table_remove(t, pos)
    for i = pos, #t do
      t[t[i]] = i -- Update changed numbers
    end
  end

  return false
end

--------------------------------------------------------------------------------

-- Handles subtables (is "deep").
-- Does not support recursive defaults tables
-- WARNING: Uses tclone()! Do not use on tables with metatables!
local twithdefaults
do
  twithdefaults = function(t, defaults)
    for k, d in pairs(defaults) do
      local v = t[k]
      if v == nil then
        if type(d) == "table" then
          d = tclone(d)
        end
        t[k] = d
      elseif type(v) == "table" and type(d) == "table" then
        twithdefaults(v, d)
      end
    end

    return t
  end
end

--------------------------------------------------------------------------------

local tifilter = function(pred, t, ...)
  local r = { }
  for i = 1, #t do
    local v = t[i]
    if pred(v, ...) then
      r[#r + 1] = v
    end
  end
  return r
end

--------------------------------------------------------------------------------

local tsetof = function(value, t)
  local r = { }

  for k, v in pairs(t) do
    r[v] = value
  end

  return r
end

--------------------------------------------------------------------------------

local tset_many = function(...)
  local r = { }

  for i = 1, select("#", ...) do
    for k, v in pairs((select(i, ...))) do
      r[v] = true
    end
  end

  return r
end

-- TODO: Pick a better name?
local tidentityset = function(t)
  local r = { }

  for k, v in pairs(t) do
    r[v] = v
  end

  return r
end

--------------------------------------------------------------------------------

local timapofrecords = function(t, key)
  local r = { }

  for i = 1, #t do
    local v = t[i]
    r[assert(v[key], "missing record key field")] = v
  end

  return r
end

local tivalues = function(t)
  local r = { }

  for i = 1, #t do
    r[#r + 1] = t[i]
  end

  return r
end

--------------------------------------------------------------------------------

-- NOTE: Optimized to be fast at simple value indexing.
--       Slower on initialization and on table value fetching.
-- WARNING: This does not protect userdata.
local treadonly, treadonly_ex
do
  local newindex = function()
    error("attempted to change read-only table")
  end

  treadonly = function(value, callbacks, tostring_fn, disable_nil)
    callbacks = callbacks or empty_table
    if disable_nil == nil then
      disable_nil = true
    end

    arguments(
        "table", value,
        "table", callbacks
      )

    optional_arguments(
        "function", tostring_fn,
        "boolean", disable_nil -- TODO: ?! Not exactly optional
      )

    local mt =
    {
      __metatable = "treadonly"; -- protect metatable

      __index = function(t, k)
        local v = rawget(value, k)
        if is_table(v) then
          -- TODO: Optimize
          v = treadonly(v, callbacks, tostring_fn, disable_nil)
        end
        if v == nil then -- TODO: Try to use metatables
          -- Note: __index does not support multiple return values in 5.1,
          --       so we can not do call right here.
          local fn = callbacks[k]
          if fn then
            return function(...) return fn(value, ...) end
          end
          if disable_nil then
            error(
                "attempted to read inexistant value at key " .. tostring(k),
                2
              )
          end
        end
        return v
      end;

      __newindex = newindex;
    }

    if tostring_fn then
      mt.__tostring = function() return tostring_fn(value) end
    end

    return setmetatable({ }, mt)
  end

  -- Changes to second return value are guaranteed to affect first one
  treadonly_ex = function(value, ...)
    local protected = treadonly(value, ...)
    return protected, value
  end
end

local tmap_kv = function(fn, t)
  local r = { }
  for k, v in pairs(t) do
    k, v = fn(k, v)
    r[k] = v
  end
  return r
end

local tmapofrecordgroups = function(t, key_name)
  local r = { }
  for k, v in pairs(t) do
    local v = t[k]
    local key = assert(v[key_name], "missing required key")
    local g = r[key]
    if not g then
      g = { }
      r[key] = g
    end
    g[#g + 1] = v
  end

  return r
end

local timapofrecordgroups = function(t, key_name)
  local r = { }
  for i = 1, #t do
    local v = t[i]
    local key = assert(v[key_name], "missing required key")
    local g = r[key]
    if not g then
      g = { }
      r[key] = g
    end
    g[#g + 1] = v
  end

  return r
end

local tilistofrecordfields = function(t, k)
  local r = { }
  for i = 1, #t do
    local v = t[i][k]
    assert(v ~= nil, "missing required key")
    r[#r + 1] = v
  end
  return r
end

local tipermute_inplace = function(t, n, count, random)
  n = n or #t
  count = count or n
  random = random or math.random

  for i = 1, count do
    local j = random(i, n)
    t[i], t[j] = t[j], t[i]
  end

  return t
end

local tkvtorecordlist = function(t, key_name, value_name)
  local result = { }
  for k, v in pairs(t) do
    result[#result + 1] = { [key_name] = k, [value_name] = v }
  end
  return result
end

local function tgetpath(t, k, nextk, ...)
  if k == nil then
    return nil
  end

  local v = t[k]
  if not is_table(v) or nextk == nil then
    return v
  end

  return tgetpath(v, nextk, ...)
end

--  tsetpath(tabl, "a", "b", "c", d)
--  tabl.a.b.c[d] = val
local tsetpath
do
  local function impl(nargs, dest, key, ...)

    if nargs == 0 then
      return dest
    end

    if key == nil then
      error("tsetpath: nil can't be a table key")
    end

    dest[key] = assert_is_table(
        dest[key] or { },
        "key `" .. tostring(key)
     .. "' already exists and its value is not a table"
      )

    return impl(nargs - 1, dest[key], ...)
  end

  tsetpath = function(dest, ...)
    local nargs = select("#", ...)
    if nargs == 0 then
      return dest
    end

    return impl(nargs, dest, ...)
  end
end

local tsetpathvalue
do
  local function impl(nargs, value, dest, key, ...)
    assert(nargs > 0)

    if key == nil then
      error("tsetpathvalue: nil can't be a table key")
    end

    if nargs == 1 then
      dest[key] = value
      return dest
    end

    dest[key] = assert_is_table(
        dest[key] or { },
        "key `" .. tostring(key)
     .. "' already exists and its value is not a table"
      )

    return impl(nargs - 1, value, dest[key], ...)
  end

  tsetpathvalue = function(value, dest, ...)
    local nargs = select("#", ...)
    if nargs == 0 then
      return dest
    end

    return impl(nargs, value, dest, ...)
  end
end

-- TODO: rename to tislice
local tslice = function(t, start_i, end_i)
  local r = { }

  start_i = math_max(start_i, 1)
  end_i = math_min(end_i, #t)
  for i = start_i, end_i do
    r[i - start_i + 1] = t[i]
  end

  return r
end

local tarraylisttohashlist = function(t, ...)
  local r = { }
  local nargs = select("#", ...)

  for i = 1, #t do
    local item = { }
    for j = 1, nargs do
      local hash = select(j, ...)
      if hash ~= nil then -- ignore nil from arguments
        item[hash] = t[i][j]
      end
    end
    r[#r + 1] = item
  end

  return r
end

local tarraytohash = function(t, ...)
  local r = { }
  local nargs = select("#", ...)

  for i = 1, nargs do
    local hash = select(i, ...)
    if hash ~= nil then -- ignore nil from arguments
      r[hash] = t[i]
    end
  end

  return r
end

local tisempty = function(t)
  return next(t) == nil
end

local tifindvalue_nonrecursive = function(t, v)
  for i = 1, #t do
    if t[i] == v then
      return true
    end
  end
  return false
end

local tkvlist2kvpairs = function(t)
  local r = { }
  for i = 1, #t, 2 do
    local k, v = t[i], t[i+1]
    if k ~= nil then
      r[k] = v
    end
  end
  return r
end

local tfilterkeylist = function(t, f, strict)
  strict = strict or false
  local r = { }

  for i = 1, #f do
    local k = f[i]
    if t[k] ~= nil then
      r[k] = t[k]
    elseif strict then
      return nil, "Field `" .. tostring(k) .. "' is absent"
    end
  end
  return r
end

local tkvmap_unpack = function(fn, t, ...)
  local r = { }
  for k, v in pairs(t) do
    k, v = fn(k, ...), fn(v, ...)

    if k ~= nil and v ~= nil then
      r[#r + 1] = k
      r[#r + 1] = v
    end
  end
  return unpack(r)
end

local tkvlist_to_hash = function(t)
  local r = { }
  for i = 1, #t, 2 do
    r[t[i]] = t[i + 1]
  end
  return r
end

local tmerge_many = function(...)
  return toverride_many({ }, ...)
end

-- Returns true is a table is an array
-- Returns false otherwise
-- Note the empty table is treated as an array
local tisarray = function(t)
  for k, _ in pairs(t) do
    if
      -- Array keys should be numbers...
      not is_number(k)
      -- ...greater than 1...
      or k < 1
      -- ...in a continuous sequence...
      or (k > 1 and t[k - 1] == nil)
      -- ...of integers...
      or k % 1 ~= 0
      -- ...avoiding floating point overflow
      or k == k - 1
    then
      return false
    end
  end
  return true
end

--------------------------------------------------------------------------------

return
{
  empty_table = empty_table;
  toverride_many = toverride_many;
  tappend_many = tappend_many;
  tijoin_many = tijoin_many;
  tkeys = tkeys;
  tvalues = tvalues;
  tkeysvalues = tkeysvalues;
  tflip = tflip;
  tflip_inplace = tflip_inplace;
  tiflip = tiflip;
  tset = tset;
  tiset = tiset;
  tisarray = tisarray;
  tiinsert_args = tiinsert_args;
  timap_inplace = timap_inplace;
  timap = timap;
  timap_sliding = timap_sliding;
  tiwalk = tiwalk;
  tiwalker = tiwalker;
  tequals = tequals;
  tiunique = tiunique;
  tgenerate_n = tgenerate_n; -- deprecated
  tgenerate_1d_linear = tgenerate_1d_linear;
  tgenerate_2d_linear = tgenerate_2d_linear;
  taccumulate = taccumulate;
  tnormalize = tnormalize;
  tnormalize_inplace = tnormalize_inplace;
  tclone = tclone;
  tcount_elements = tcount_elements;
  tremap_to_array = tremap_to_array;
  twalk_pairs = twalk_pairs;
  tmap_values = tmap_values;
  torderedset = torderedset;
  torderedset_insert = torderedset_insert;
  torderedset_remove = torderedset_remove;
  twithdefaults = twithdefaults;
  tifilter = tifilter;
  tsetof = tsetof;
  tset_many = tset_many;
  tidentityset = tidentityset;
  timapofrecords = timapofrecords;
  tivalues = tivalues;
  treadonly = treadonly;
  treadonly_ex = treadonly_ex;
  tmap_kv = tmap_kv;
  tmapofrecordgroups = tmapofrecordgroups;
  timapofrecordgroups = timapofrecordgroups;
  tilistofrecordfields = tilistofrecordfields;
  tipermute_inplace = tipermute_inplace;
  tkvtorecordlist = tkvtorecordlist;
  tgetpath = tgetpath;
  tsetpath = tsetpath;
  tsetpathvalue = tsetpathvalue;
  tslice = tslice;
  tarraylisttohashlist = tarraylisttohashlist;
  tarraytohash = tarraytohash;
  tkvlist2kvpairs = tkvlist2kvpairs;
  tfilterkeylist = tfilterkeylist;
  tisempty = tisempty;
  tifindvalue_nonrecursive = tifindvalue_nonrecursive;
  tkvmap_unpack = tkvmap_unpack;
  tkvlist_to_hash = tkvlist_to_hash;
  tmerge_many = tmerge_many;
}