local classlib = require("core.class")
local table_is_class = function(t)
  return rawget(t, "CLASS") == t
end
local pickle_tableref_add = function(P, t)
  do
    local tableref = P.tmp.table_to_tableref[t]
    assert(tableref == nil, "Trying to replace an existing ref")
  end
  local tableref_index = P.tmp.tableref_next_free
  local tableref_data = {}
  local tableref = {tableref = tableref_index}
  P.tmp.table_to_tableref[t] = tableref
  P.tmp.tableref_next_free = tableref_index + 1
  P.tableref_to_table[tableref_index] = tableref_data
  return tableref, tableref_index, tableref_data
end
local pickle_stringref_add = function(P, s)
  local stringref = P.tmp.string_to_stringref[s]
  if not stringref then
    stringref = P.tmp.stringref_next_free
    P.tmp.stringref_next_free = stringref + 1
    P.tmp.string_to_stringref[s] = stringref
    P.stringref_to_string[stringref] = s
  end
  return stringref
end
local pickle_classref_add = function(P, CLASS)
  local classref = P.tmp.classname_to_classref[CLASS.CLASSNAME]
  if not classref then
    classref = {
      classref = pickle_stringref_add(P, CLASS.CLASSNAME)
    }
    P.tmp.classname_to_classref[CLASS.CLASSNAME] = classref
  end
  return classref
end
local key_is_allowed_in_pickle = function(k)
  return k ~= "__identity"
end
local pickle_table_pickle
local pickle_value_pickle = function(P, v)
  local tt = type(v)
  if tt == "table" then
    local tableref = pickle_table_pickle(P, v)
    return tableref
  elseif tt == "string" then
    pickle_stringref_add(P, v)
    return v
  elseif tt == "userdata" then
    if engine.CanPickle and engine.CanPickle(v) then
      return v
    end
  elseif tt ~= "function" then
    return v
  end
end
function pickle_table_pickle(P, t)
  do
    local tableref = P.tmp.table_to_tableref[t]
    if tableref then
      return tableref
    end
  end
  if table_is_class(t) then
    return pickle_classref_add(P, t)
  end
  local tableref, tableref_index, tableref_data = pickle_tableref_add(P, t)
  local mt = getmetatable(t)
  if type(mt) == "table" and mt.CLASS == mt then
    pickle_classref_add(P, mt.CLASS)
    P.tableref_to_classname[tableref_index] = mt.CLASSNAME
  end
  for k, v in pairs(t) do
    if key_is_allowed_in_pickle(k) then
      local pk, pv = pickle_value_pickle(P, k), pickle_value_pickle(P, v)
      if pk ~= nil then
        tableref_data[pk] = pv
      elseif engine.CanPickle then
        engine.Warning("Unable to pickle key '", k, "', which is a value of type '", type(v), "'. Script may not correctly restore from checkpoint.")
      end
    end
  end
  return tableref
end
local pickle_init = function()
  return {
    tmp = {
      table_to_tableref = {},
      tableref_next_free = 1,
      classname_to_classref = {},
      string_to_stringref = {},
      stringref_next_free = 1
    },
    tableref_to_table = {},
    tableref_to_classname = {},
    stringref_to_string = {}
  }
end
local pickle_close = function(P)
  P.tmp = nil
end
local unpickle_init = function(P)
  P.tmp = {
    tableref_to_fixup = {},
    refs_to_restore = {}
  }
end
local unpickle_find_fixups
local unpickle_add_fixup = function(P, t, tref)
  local tableref = tref.tableref
  local fixup = P.tmp.tableref_to_fixup[tableref]
  if fixup then
    return
  end
  P.tmp.tableref_to_fixup[tableref] = t
  unpickle_find_fixups(P, t, tableref)
end
function unpickle_find_fixups(P, existing_table, tableref)
  local data = P.tableref_to_table[tableref]
  for k, v in pairs(existing_table) do
    local vref = data[k]
    if type(k) ~= "table" and type(v) == "table" and not table_is_class(v) and type(vref) == "table" and vref.tableref ~= nil then
      unpickle_add_fixup(P, v, vref)
    end
  end
end
local function unpickle_traverse_refs_to_restore(P, tableref)
  P.tmp.refs_to_restore[tableref] = true
  local data = P.tableref_to_table[tableref]
  for k, v in pairs(data) do
    if type(k) == "table" and k.tableref and not P.tmp.refs_to_restore[k.tableref] then
      unpickle_traverse_refs_to_restore(P, k.tableref)
    end
    if type(v) == "table" and v.tableref and not P.tmp.refs_to_restore[v.tableref] then
      unpickle_traverse_refs_to_restore(P, v.tableref)
    end
  end
end
local unpickle_tableref_for_fixup = function(P, fixup)
  for tableref, t in pairs(P.tmp.tableref_to_fixup) do
    if t == fixup then
      return tableref
    end
  end
end
local unpickle_prepare_target_table = function(P, tableref)
  local fixup = P.tmp.tableref_to_fixup[tableref]
  if fixup then
    setmetatable(fixup, nil)
    for k in pairs(fixup) do
      fixup[k] = nil
    end
  else
    fixup = {}
    P.tmp.tableref_to_fixup[tableref] = fixup
  end
  return fixup
end
local unpickle_get_table = function(P, tableref)
  local t = P.tmp.tableref_to_fixup[tableref]
  if not t then
    t = {}
    P.tmp.tableref_to_fixup[tableref] = t
  end
  return t
end
local function unpickle_value(P, v)
  local tt = type(v)
  if tt == "table" then
    if v.classref then
      local CLASSNAME = unpickle_value(v.classref)
      local CLASS = classlib.lookup(CLASSNAME)
      return CLASS
    else
      return unpickle_get_table(P, v.tableref)
    end
  else
    return v
  end
end
local unpickle_restore_data = function(P, data, fixup)
  for k, v in pairs(data) do
    local uk, uv = unpickle_value(P, k), unpickle_value(P, v)
    fixup[uk] = uv
  end
end
local unpickle_restore = function(P)
  for tableref in pairs(P.tmp.refs_to_restore) do
    local data = P.tableref_to_table[tableref]
    local fixup = unpickle_prepare_target_table(P, tableref)
    unpickle_restore_data(P, data, fixup)
  end
end
local unpickle_bless = function(P)
  for tableref in pairs(P.tmp.refs_to_restore) do
    local fixup = P.tmp.tableref_to_fixup[tableref]
    local CLASSNAME = P.tableref_to_classname[tableref]
    if CLASSNAME then
      local CLASS = classlib.lookup(CLASSNAME)
      classlib.bless(fixup, CLASS)
    end
  end
end
local unpickle_finalize = function(P)
  for tableref in pairs(P.tmp.refs_to_restore) do
    local fixup = P.tmp.tableref_to_fixup[tableref]
    local CLASSNAME = P.tableref_to_classname[tableref]
    if CLASSNAME then
      local unpickle_fn = fixup.OnUnpickle
      if unpickle_fn then
        unpickle_fn(fixup)
      end
    end
  end
end
local unpickle_close = function(P)
  P.tmp = nil
end
local indent_strings = {}
for i = 0, 20 do
  indent_strings[i] = string.rep("  ", i)
end
local val = function(v)
  if type(v) ~= "table" then
    return v
  end
  return v.tableref and string.format("tableref:%d", v.tableref) or v.__identity or v.CLASSNAME and getmetatable(v).__identity or v
end
local function dumptable(t, indent, refs)
  local indent_string = indent_strings[indent]
  for k, v in pairs(t) do
    if k ~= "__identity" then
      engine.Print(indent_string, val(k), ": ", val(v))
      if type(v) == "table" and not refs[v] and v.CLASS ~= v and not v.tableref then
        refs[v] = true
        dumptable(v, indent + 1, refs)
      end
    end
  end
end
local pickle, unpickle, unpickle_partial
local OnPickleInternal = function(go)
  _G.__new_pickle__ = true
  local UserPickle = _G.Pickle
  if not UserPickle then
    return
  end
  local root = UserPickle(go)
  if not root then
    return
  end
  local P = pickle(root)
  return P
end
local OnUnpickleInternal = function(go, P)
  local UserPickle = _G.Pickle
  local UserUnpickle = _G.Unpickle
  if not UserPickle or not UserUnpickle then
    return
  end
  local root = UserPickle(go)
  if not root then
    return
  end
  unpickle(P, root)
  UserUnpickle(go, root)
end
local install = function()
  assert(not _G.OnPickleInternal, "Pickle hook already defined?")
  assert(not _G.OnUnpickleInternal, "Unpickle hook already defined?")
  _G.OnPickleInternal = OnPickleInternal
  _G.OnUnpickleInternal = OnUnpickleInternal
end
function pickle(root)
  local P = pickle_init()
  P.tableref_root = pickle_table_pickle(P, root).tableref
  pickle_close(P)
  return P
end
function unpickle(P, root)
  unpickle_partial(P, root, root)
end
function unpickle_partial(P, root, partial_root)
  unpickle_init(P)
  unpickle_add_fixup(P, root, {
    tableref = P.tableref_root
  })
  local partial_tableref = unpickle_tableref_for_fixup(P, partial_root)
  unpickle_traverse_refs_to_restore(P, partial_tableref)
  unpickle_restore(P)
  unpickle_bless(P)
  unpickle_finalize(P)
  unpickle_close(P)
end
return {
  Pickle = pickle,
  Unpickle = unpickle,
  Install = install,
  dumptable = dumptable
}
