-- Copyright (C) 2010 Tyson Brown
-- 
-- This file is part QuestHelper2, an AddOn for World of Warcraft.
-- 
-- QuestHelper2 is free software: you can redistribute it and/or modify
-- it under the terms of the GNU General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
-- 
-- QuestHelper2 is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.
-- 
-- You should have received a copy of the GNU General Public License
-- along with QuestHelper2. If not, see <http://www.gnu.org/licenses/>.

if LibDebug then LibDebug() end
local shared = select(2, ...)

local floor, frexp, ldexp = math.floor, math.frexp, math.ldexp
local sort, insert, type, tostring, pairs, ipairs = table.sort, table.insert, type, tostring, pairs, ipairs

-- If any of these types are encountered, they will be encoded as nil, unless they happen to be an external reference.
local invalid_type = {
  ["function"] = true,
  ["thread"] = true,
  ["userdata"] = true
}

-- Used to sort the keys of the serialized tables.
local function keyCmp(a, b)
  local ta, tb = type(a), type(b)
  
  if ta == tb then
    if ta == "number" then
      -- Will probably explode on NaNs, so try not to use those as table keys.
      -- Actually, if a NaN isn't equal to anything else, would it even be possible to
      -- look one up from a table?
      -- * Does a quest experiment to find out... *
      -- Nope, can't use NaNs as table keys. Excellent, one less bit of stupidity to worry about.
      return a < b
    else
      return tostring(a) < tostring(b)
    end
  else
    return ta < tb
  end
end

-- The integer encoders/decoders use this for searching through the
-- space of all possible integers.
local function getSubPartition(partition, dir)
  local sub = partition[dir]
  
  if not sub then
    sub = {[0]=0,1,2}
    partition[dir] = sub
  end
  
  return sub
end

-- Use if no external references are passed to us.
local empty_external_references = {}

-- Validates external references, and returns a table containing the reversed mapping.
local function generateReverseReferences(references)
  local reversed = {}
  
  for key, val in pairs(references) do
    assert(type(key) == "string")
    
    if reversed[val] ~= nil then
      error(tostring(val)..(" has multiple names: %q and %q"):format(reversed[val], key))
    end
    
    reversed[val] = key
  end
  
  return reversed
end

-- external_references is a bunch of strings mapping to functions or metatables
-- that need to be saved, but which we're not capable of serializing correctly.
-- cb is invoked at the end of every call to encodeObject;
-- onSerialize passes coroutine.yield as cb so it can spread the
-- work of compression over several calls.
function shared.createSerializer(external_references, cb)
  local encoder = shared.createArithmeticEncoder()
  
  -- Strings can be big and take a lot of time to encode or decode, so
  -- This is a little wrapper function that calls cb once every 10 times
  -- it is called. encodeString calls it at the end of each byte.
  local string_callback
  
  if cb then
    local count = 0
    string_callback = function ()
      if count == 9 then
        count = 0
        cb()
      else
        count = count + 1
      end
    end
  else
    -- do nothing
    string_callback = function () end
  end
  
  external_references = external_references or empty_external_references
  local reversed_external_references = generateReverseReferences(external_references)
  
  -- Check if an object would encode to nil, so the
  -- table encoder can skip it.
  local function encodesToNil(obj)
    return obj == nil or (invalid_type[type(obj)] and not reversed_external_references[obj])
  end
  
  local encodeSymbol = encoder.encodeSymbol
  local encodeTable, encodeString, encodeNumber, encodeBoolean, encodeExternal
  
  -- 1=nil, 2=table, 3=string, 4=number, 5=boolean, 6=external
  local typeEncode = shared.createMarkovChainEncoder(encodeSymbol, 6)
  
  local function encodeObject(object)
    local object_type = type(object)
    
    if encodesToNil(object) then
      typeEncode(1)
    elseif reversed_external_references[object] then
      typeEncode(6)
      encodeString(reversed_external_references[object])
    elseif object_type == "table" then
      typeEncode(2)
      encodeTable(object)
    elseif object_type == "string" then
      typeEncode(3)
      encodeString(object)
    elseif object_type == "number" then
      typeEncode(4)
      encodeNumber(object)
    else
      --[===[@debug@
      assert(object_type == "boolean")
      --@end-debug@]===]
      typeEncode(5)
      encodeBoolean(object)
    end
    
    if cb then cb() end
  end
  
  -- [1,256] -> bytes [0,255], 257 -> end of string
  local string_prob = {[0]=0, 1, 2}
  local seen_strings = {[""]=1}
  local byteEncode = shared.createMarkovChainEncoder(encodeSymbol, 257)
  encodeString = function(str)
    --[===[@debug@
    assert(type(str) == "string")
    --@end-debug@]===]
    
    local lookup = seen_strings[str]
    
    if lookup then
      -- Already seen the required string, encode a reference to it.
      encodeSymbol(string_prob, lookup)
    else
      -- This is a new string.
      lookup = #string_prob
      encodeSymbol(string_prob, lookup)
      string_prob[lookup] = string_prob[lookup-1]+1
      seen_strings[str] = lookup
      string_prob[lookup+1] = string_prob[lookup]+1
      
      for i=1,str:len() do
        byteEncode(str:byte(i)+1)
        string_callback()
      end
      
      byteEncode(257)
    end
  end
  
  -- 1=false, 2=true
  local booleanEncode = shared.createMarkovChainEncoder(encodeSymbol, 2)
  encodeBoolean = function(bool)
    --[===[@debug@
    assert(type(bool) == "boolean")
    --@end-debug@]===]
    
    booleanEncode(bool and 2 or 1)
  end
  
  local exponentEncode = shared.createMarkovChainEncoder(encodeSymbol, 96)
  local exponent_partitions = {}
  
  -- This function has to encode numbers in the range (0,∞)
  local function encodeFloat(num)
    -- double numbers have an integer exponent in the range [-1022, 1023],
    -- and significand in the range of [0.5,1) with 52 bits of precision (or 0 if num==0, or (-1,-0.5] if num<0, but those cases don't affect us.)
    
    -- We however don't encode the numbers with that amount of precision.
    -- We use an exponent in the range [-32,63], and a significand with 24 bits of precision,
    -- which allows us to encode numbers in the range of 0.000000000116415321826934814453125 to 9223372036854775808,
    -- with an accuracy just slightly above single precision.
    local significand, exponent = frexp(num)
    
    if exponent < -32 then
      significand, exponent = 0.5, -32
    elseif exponent > 63 then
      significand, exponent = 1, 63
    end
    
    exponentEncode(33+exponent)
    
    significand=floor(33554430*significand-16777214.5)
    --[===[@debug@
    assert(significand >= 0 and significand <= 16777215)
    --@end-debug@]===]
    
    local low, high, partition = 0, 16777215, exponent_partitions[exponent]
    
    if not partition then
      partition = {[0]=0,1,2}
      exponent_partitions[exponent] = partition
    end
    
    while low ~= high do
      local mid = floor((low+high)*0.5)
      if significand > mid then
        encodeSymbol(partition, 2)
        low, partition = mid+1, getSubPartition(partition, "r")
      else
        encodeSymbol(partition, 1)
        high, partition = mid, getSubPartition(partition, "l")
      end
    end
    
    --[===[@debug@
    assert(low == significand)
    --@end-debug@]===]
  end
  
  -- This function has to encode numbers in the range [1,2^31-1]
  -- In fact, it can encode beyond that upper limit, but due
  -- to precision, it can get stuck in infinite loops trying to
  -- encode very large numbers exactly.
  local function encodeIntegerPartition(partition, num)
    local low, high = 1, 8
    
    while num > high do
      encodeSymbol(partition, 2)
      low, high, partition = high+1, high*2, getSubPartition(partition, "r")
    end
    
    encodeSymbol(partition, 1)
    partition = getSubPartition(partition, "l")
    
    while low ~= high do
      local mid = floor((low+high)*0.5)
      
      if num > mid then
        encodeSymbol(partition, 2)
        low, partition = mid+1, getSubPartition(partition, "r")
      else
        encodeSymbol(partition, 1)
        high, partition = mid, getSubPartition(partition, "l")
      end
    end
  end
  
  local positive_integer_partition_data = {[0]=0,1,2}
  local negative_integer_partition_data = {[0]=0,1,2}
  
  -- 1=zero, 2=special, 3=+integer, 4=+float, 5=-integer, 6=-float
  local numclassEncoder = shared.createMarkovChainEncoder(encodeSymbol, 6)
  -- 1=positive infinity, 2=negative infinity, 3=nan
  local specialNumberEncode = shared.createMarkovChainEncoder(encodeSymbol, 3)
  
  encodeNumber = function(num)
    if num == 0 then -- Zero
      numclassEncoder(1)
    elseif num == 1/0 then -- Infinity
      numclassEncoder(2)
      specialNumberEncode(1)
    elseif num == -1/0 then -- Negative Infinity
      numclassEncoder(2)
      specialNumberEncode(2)
    elseif num ~= num then -- NaN
      numclassEncoder(2)
      specialNumberEncode(3)
    elseif num > 0 then
      if num < 2147483648 and floor(num) == num then -- Positive Integer
        numclassEncoder(3)
        encodeIntegerPartition(positive_integer_partition_data, num)
      else -- Positive Float
        numclassEncoder(4)
        encodeFloat(num)
      end
    else
      if num > -2147483648 and floor(num) == num then -- Negative Integer
        numclassEncoder(5)
        encodeIntegerPartition(negative_integer_partition_data, -num)
      else -- Negative Float
        numclassEncoder(6)
        encodeFloat(-num)
      end
    end
  end
  
  local table_prob = {[0]=0, 1}
  local seen_tables = {}
  -- 1=no metatable, 2=has metatable
  local meta_prob = {[0]=0, 1, 2}
  
  -- 1=new key follows, 2=done
  local endOfTableEncode = shared.createMarkovChainEncoder(encodeSymbol, 2)
  
  encodeTable = function(tbl)
    local lookup = seen_tables[tbl]
    if lookup then
      -- Encode a reference to a previously encoded table.
      encodeSymbol(table_prob, lookup)
    else
      -- Haven't seen this table before, encode it properly.
      lookup = #table_prob
      encodeSymbol(table_prob, lookup)
      seen_tables[tbl] = lookup
      table_prob[lookup] = table_prob[lookup-1]+1
      table_prob[lookup+1] = table_prob[lookup]+1
      
      -- Encode the metatable.
      local meta = getmetatable(tbl)
      
      if meta and reversed_external_references[meta] then
        encodeSymbol(meta_prob, 2)
        encodeString(reversed_external_references[meta])
      else
        encodeSymbol(meta_prob, 1)
        meta = nil
      end
      
      if meta and meta.serialize then
        -- This object can serialize itself, and probably do a better job than us.
        meta.serialize(tbl, encoder)
      else
        -- Encode the array portion of the table.
        local len = #tbl
        encodeIntegerPartition(positive_integer_partition_data, len+1)
        
        for i=1,len do
          encodeObject(tbl[i])
        end
        
        -- Encode the associative array portion of the table.
        local tbl_key_list = {}
        
        -- Add any keys that weren't part of the array portion of the table.
        for key in pairs(tbl) do
          if type(key) ~= "number" or key < 1 or key > len or floor(key) ~= key then
            insert(tbl_key_list, key)
          end
        end
        
        -- Sort the keys.
        sort(tbl_key_list, keyCmp)
        
        -- Encode the key/value pairs.
        for i, key in ipairs(tbl_key_list) do
          endOfTableEncode(1)
          encodeObject(key)
          encodeObject(tbl[key])
        end
        
        endOfTableEncode(2)
      end
    end
  end
  
  encoder.encodeString = encodeString
  encoder.encodeNumber = encodeNumber
  encoder.encodeTable = encodeTable
  encoder.encodeObject = encodeObject
  
  return encoder
end

function shared.createUnserializer(encoded_data, external_references, cb)
  local decoder = shared.createArithmeticDecoder(encoded_data)
  
  -- Strings can be big and take a lot of time to encode or decode, so
  -- This is a little wrapper function that calls cb once every 10 times
  -- it is called. encodeString calls it at the end of each byte.
  local string_callback
  
  if cb then
    local count = 0
    string_callback = function ()
      if count == 9 then
        count = 0
        cb()
      else
        count = count + 1
      end
    end
  else
    -- do nothing
    string_callback = function () end
  end
  
  external_references = external_references or empty_external_references
  
  local decodeSymbol = decoder.decodeSymbol
  local decodeTable, decodeString, decodeNumber, decodeBoolean, decodeExternal
  
  -- 1=nil, 2=table, 3=string, 4=number, 5=boolean, 6=external
  local typeDecode = shared.createMarkovChainDecoder(decodeSymbol, 6)
  
  local function decodeObject()
    local object_type = typeDecode()
    local object
    
    if object_type == 1 then
      object = nil
    elseif object_type == 2 then
      object = decodeTable()
    elseif object_type == 3 then
      object = decodeString()
    elseif object_type == 4 then
      object = decodeNumber()
    elseif object_type == 5 then
      object = decodeBoolean()
    else
      local reference = decodeString()
      object = external_references[reference]
      
      if object == nil then
        error("Undefined external reference to: "..reference)
      end
    end
    
    if cb then cb() end
    
    return object
  end
  
  -- [1,256] -> bytes [0,255], 257 -> end of string
  local string_prob = {[0]=0, 1, 2}
  local seen_strings = {[1]=""}
  local byteDecode = shared.createMarkovChainDecoder(decodeSymbol, 257)
  decodeString = function()
    local lookup = decodeSymbol(string_prob)
    local str = seen_strings[lookup]
    
    if not str then
      string_prob[lookup] = string_prob[lookup-1]+1
      string_prob[lookup+1] = string_prob[lookup]+1
      
      str = ""
      
      while true do
        local byte = byteDecode()-1
        if byte == 256 then break end
        str = str .. string.char(byte)
        string_callback()
      end
      
      seen_strings[lookup] = str
    end
    
    return str
  end
  
  -- 1=false, 2=true
  local booleanDecode = shared.createMarkovChainDecoder(decodeSymbol, 2)
  decodeBoolean = function(bool)
    return booleanDecode() == 2
  end
  
  local exponentDecode = shared.createMarkovChainDecoder(decodeSymbol, 96)
  local exponent_partitions = {}
  
  local function decodeFloat()
    local exponent = exponentDecode()-33
    local low, high, partition = 0, 16777215, exponent_partitions[exponent]
    
    if not partition then
      partition = {[0]=0,1,2}
      exponent_partitions[exponent] = partition
    end
    
    while low ~= high do
      local mid = floor((low+high)*0.5)
      if decodeSymbol(partition) == 2 then
        low, partition = mid+1, getSubPartition(partition, "r")
      else
        high, partition = mid, getSubPartition(partition, "l")
      end
    end
    
    local significand = low*0.00000002980232416405226+0.5000000149011621
    
    return ldexp(significand, exponent)
  end
  
  local function decodeIntegerPartition(partition)
    local low, high = 1, 8
    
    while decodeSymbol(partition) == 2 do
      low, high, partition = high+1, high*2, getSubPartition(partition, "r")
    end
    
    partition = getSubPartition(partition, "l")
    
    while low ~= high do
      local mid = floor((low+high)*0.5)
      
      if decodeSymbol(partition) == 2 then
        low, partition = mid+1, getSubPartition(partition, "r")
      else
        high, partition = mid, getSubPartition(partition, "l")
      end
    end
    
    return low
  end
  
  local positive_integer_partition_data = {[0]=0,1,2}
  local negative_integer_partition_data = {[0]=0,1,2}
  
  -- 1=zero, 2=special, 3=+integer, 4=+float, 5=-integer, 6=-float
  local numclassDecode = shared.createMarkovChainDecoder(decodeSymbol, 6)
  -- 1=positive infinity, 2=negative infinity, 3=nan
  local specialNumberDecode = shared.createMarkovChainDecoder(decodeSymbol, 3)
  
  decodeNumber = function()
    local class = numclassDecode()
    
    if class == 1 then -- Zero
      return 0
    elseif class == 2 then -- Special
      local special = specialNumberDecode()
      if special == 1 then -- Infinity
        return 1/0
      elseif special == 2 then -- Negative Infinity
        return -1/0
      else -- NaN
        return 0/0
      end
    elseif class == 3 then -- Positive Integer
      return decodeIntegerPartition(positive_integer_partition_data)
    elseif class == 4 then -- Positive Float
      return decodeFloat()
    elseif class == 5 then -- Negative Integer
      return -decodeIntegerPartition(negative_integer_partition_data)
    else -- Negative Float
      return -decodeFloat()
    end
  end
  
  local table_prob = {[0]=0, 1}
  local seen_tables = {}
  -- 1=no metatable, 2=has metatable
  local meta_prob = {[0]=0, 1, 2}
  
  -- 1=new key follows, 2=done
  local endOfTableDecode = shared.createMarkovChainDecoder(decodeSymbol, 2)
  
  decodeTable = function()
    local lookup = decodeSymbol(table_prob)
    local tbl = seen_tables[lookup]
    
    if not tbl then
      -- Haven't seen this table before, decode it properly.
      tbl = {}
      local meta
      seen_tables[lookup] = tbl
      table_prob[lookup] = table_prob[lookup-1]+1
      table_prob[lookup+1] = table_prob[lookup]+1
      
      if decodeSymbol(meta_prob, 2) == 2 then
        -- Has a metatable.
        local reference = decodeString()
        meta = external_references[reference]
        
        if meta == nil then
          error("Undefined external reference to: "..reference)
        elseif type(meta) ~= "table" then
          error("External reference to: "..reference.." can't be used as metatable.")
        end
      end
      
      if meta and meta.unserialize then
        -- This object can serialize itself.
        meta.unserialize(tbl, decoder)
      else
        -- decode the array portion of the table.
        local len = decodeIntegerPartition(positive_integer_partition_data)-1
        for i=1,len do
          tbl[i] = decodeObject()
        end
        
        while endOfTableDecode() == 1 do
          local key, val = decodeObject(), decodeObject()
          tbl[key] = val
        end
      end
      
      setmetatable(tbl, meta)
    end
    
    return tbl
  end
  
  decoder.decodeString = decodeString
  decoder.decodeNumber = decodeNumber
  decoder.decodeTable = decodeTable
  decoder.decodeObject = decodeObject
  
  return decoder
end
