-- Concise Binary Object Representation (CBOR)
-- RFC 7049

local function softreq(pkg, field)
    local ok, mod = pcall(require, pkg)
    if not ok then
        return
    end
    if field then
        return mod[field]
    end
    return mod
end
local dostring = function(s)
    local ok, f = load(function()
        local ret = s
        s = nil
        return ret
    end)
    if ok and f then
        return f()
    end
end

local setmetatable = setmetatable
local getmetatable = getmetatable
local dbg_getmetatable = debug.getmetatable
local assert = assert
local error = error
local type = type
local pairs = pairs
local ipairs = ipairs
local tostring = tostring
local s_char = string.char
local t_concat = table.concat
local t_sort = table.sort
local m_floor = math.floor
local m_abs = math.abs
local m_huge = math.huge
local m_max = math.max
local maxint = math.maxinteger or 9007199254740992
local minint = math.mininteger or -9007199254740992
local NaN = 0 / 0
local m_frexp = math.frexp
local m_ldexp = math.ldexp or function(x, exp)
    return x * 2.0 ^ exp
end
local m_type = math.type
    or function(n)
        return n % 1 == 0 and n <= maxint and n >= minint and "integer" or "float"
    end
local s_pack = string.pack or softreq("struct", "pack")
local s_unpack = string.unpack or softreq("struct", "unpack")
local b_rshift = function(a, b)
    return a >> b
end

-- sanity check
if s_pack and s_pack(">I2", 0) ~= "\0\0" then
    s_pack = nil
end
if s_unpack and s_unpack(">I2", "\1\2\3\4") ~= 0x102 then
    s_unpack = nil
end

local _ENV = nil -- luacheck: ignore 211

local encoder = {}

local function encode(obj, opts)
    return encoder[type(obj)](obj, opts)
end

-- Major types 0, 1 and length encoding for others
local function integer(num, m)
    if m == 0 and num < 0 then
        -- negative integer, major type 1
        num, m = -num - 1, 32
    end
    if num < 24 then
        return s_char(m + num)
    elseif num < 2 ^ 8 then
        return s_char(m + 24, num)
    elseif num < 2 ^ 16 then
        return s_char(m + 25, b_rshift(num, 8), num % 0x100)
    elseif num < 2 ^ 32 then
        return s_char(
            m + 26,
            b_rshift(num, 24) % 0x100,
            b_rshift(num, 16) % 0x100,
            b_rshift(num, 8) % 0x100,
            num % 0x100
        )
    elseif num < 2 ^ 64 then
        local high = m_floor(num / 2 ^ 32)
        num = num % 2 ^ 32
        return s_char(
            m + 27,
            b_rshift(high, 24) % 0x100,
            b_rshift(high, 16) % 0x100,
            b_rshift(high, 8) % 0x100,
            high % 0x100,
            b_rshift(num, 24) % 0x100,
            b_rshift(num, 16) % 0x100,
            b_rshift(num, 8) % 0x100,
            num % 0x100
        )
    end
    error("int too large")
end

if s_pack then
    function integer(num, m)
        local fmt
        m = m or 0
        if num < 24 then
            fmt, m = ">B", m + num
        elseif num < 256 then
            fmt, m = ">BB", m + 24
        elseif num < 65536 then
            fmt, m = ">BI2", m + 25
        elseif num < 4294967296 then
            fmt, m = ">BI4", m + 26
        else
            fmt, m = ">BI8", m + 27
        end
        return s_pack(fmt, m, num)
    end
end

local simple_mt = {}
function simple_mt:__tostring()
    return self.name or ("simple(%d)"):format(self.value)
end
function simple_mt:__tocbor()
    return self.cbor or integer(self.value, 224)
end

local function simple(value, name, cbor)
    assert(value >= 0 and value <= 255, "bad argument #1 to 'simple' (integer in range 0..255 expected)")
    return setmetatable({ value = value, name = name, cbor = cbor }, simple_mt)
end

local tagged_mt = {}
function tagged_mt:__tostring()
    return ("%d(%s)"):format(self.tag, tostring(self.value))
end
function tagged_mt:__tocbor(opts)
    return integer(self.tag, 192) .. encode(self.value, opts)
end

local function tagged(tag, value)
    assert(tag >= 0, "bad argument #1 to 'tagged' (positive integer expected)")
    return setmetatable({ tag = tag, value = value }, tagged_mt)
end

local null = simple(22, "null") -- explicit null
local undefined = simple(23, "undefined") -- undefined or nil
local BREAK = simple(31, "break", "\255")

-- Number types dispatch
function encoder.number(num)
    return encoder[m_type(num)](num)
end

-- Major types 0, 1
function encoder.integer(num)
    if num < 0 then
        return integer(-1 - num, 32)
    end
    return integer(num, 0)
end

-- Major type 7
function encoder.float(num)
    if num ~= num then -- NaN shortcut
        return "\251\127\255\255\255\255\255\255\255"
    end
    local sign = (num > 0 or 1 / num > 0) and 0 or 1
    num = m_abs(num)
    if num == m_huge then
        return s_char(251, sign * 128 + 128 - 1) .. "\240\0\0\0\0\0\0"
    end
    local fraction, exponent = m_frexp(num)
    if fraction == 0 then
        return s_char(251, sign * 128) .. "\0\0\0\0\0\0\0"
    end
    fraction = fraction * 2
    exponent = exponent + 1024 - 2
    if exponent <= 0 then
        fraction = fraction * 2 ^ (exponent - 1)
        exponent = 0
    else
        fraction = fraction - 1
    end
    return s_char(
        251,
        sign * 2 ^ 7 + m_floor(exponent / 2 ^ 4) % 2 ^ 7,
        exponent % 2 ^ 4 * 2 ^ 4 + m_floor(fraction * 2 ^ 4 % 0x100),
        m_floor(fraction * 2 ^ 12 % 0x100),
        m_floor(fraction * 2 ^ 20 % 0x100),
        m_floor(fraction * 2 ^ 28 % 0x100),
        m_floor(fraction * 2 ^ 36 % 0x100),
        m_floor(fraction * 2 ^ 44 % 0x100),
        m_floor(fraction * 2 ^ 52 % 0x100)
    )
end

if s_pack then
    function encoder.float(num)
        return s_pack(">Bd", 251, num)
    end
end

-- Major type 2 - byte strings
function encoder.bytestring(s)
    return integer(#s, 64) .. s
end

-- Major type 3 - UTF-8 strings
function encoder.utf8string(s)
    return integer(#s, 96) .. s
end

-- Lua strings are byte strings
encoder.string = encoder.bytestring

function encoder.boolean(bool)
    return bool and "\245" or "\244"
end

encoder["nil"] = function()
    return "\246"
end

function encoder.userdata(ud, opts)
    local mt = dbg_getmetatable(ud)
    if mt then
        local encode_ud = opts and opts[mt] or mt.__tocbor
        if encode_ud then
            return encode_ud(ud, opts)
        end
    end
    error("can't encode userdata")
end

function encoder.table(t, opts)
    local mt = getmetatable(t)
    if mt then
        local encode_t = opts and opts[mt] or mt.__tocbor
        if encode_t then
            return encode_t(t, opts)
        end
    end
    -- the table is encoded as an array iff when we iterate over it,
    -- we see successive integer keys starting from 1.  The lua
    -- language doesn't actually guarantee that this will be the case
    -- when we iterate over a table with successive integer keys, but
    -- due an implementation detail in PUC Rio Lua, this is what we
    -- usually observe.  See the Lua manual regarding the # (length)
    -- operator.  In the case that this does not happen, we will fall
    -- back to a map with integer keys, which becomes a bit larger.
    local array, map, i, p = { integer(#t, 128) }, { "\191" }, 1, 2
    local is_array = true
    for k, v in pairs(t) do
        is_array = is_array and i == k
        i = i + 1

        local encoded_v = encode(v, opts)
        array[i] = encoded_v

        map[p], p = encode(k, opts), p + 1
        map[p], p = encoded_v, p + 1
    end
    -- map[p] = "\255";
    map[1] = integer(i - 1, 160)
    return t_concat(is_array and array or map)
end

-- Array or dict-only encoders, which can be set as __tocbor metamethod
function encoder.array(t, opts)
    local array = {}
    for i, v in ipairs(t) do
        array[i] = encode(v, opts)
    end
    return integer(#array, 128) .. t_concat(array)
end

function encoder.map(t, opts)
    local map, p, len = { "\191" }, 2, 0
    for k, v in pairs(t) do
        map[p], p = encode(k, opts), p + 1
        map[p], p = encode(v, opts), p + 1
        len = len + 1
    end
    -- map[p] = "\255";
    map[1] = integer(len, 160)
    return t_concat(map)
end
encoder.dict = encoder.map -- COMPAT

function encoder.ordered_map(t, opts)
    local map = {}
    if not t[1] then -- no predefined order
        local i = 0
        for k in pairs(t) do
            i = i + 1
            map[i] = k
        end
        t_sort(map)
    end
    for i, k in ipairs(t[1] and t or map) do
        map[i] = encode(k, opts) .. encode(t[k], opts)
    end
    return integer(#map, 160) .. t_concat(map)
end

encoder["function"] = function()
    error("can't encode function")
end

-- Decoder
-- Reads from a file-handle like object
local function read_bytes(fh, len)
    return fh:read(len)
end

local function read_byte(fh)
    return fh:read(1):byte()
end

local function read_length(fh, mintyp)
    if mintyp < 24 then
        return mintyp
    elseif mintyp < 28 then
        local out = 0
        for _ = 1, 2 ^ (mintyp - 24) do
            out = out * 256 + read_byte(fh)
        end
        return out
    else
        error("invalid length")
    end
end

local decoder = {}

local function read_type(fh)
    local byte = read_byte(fh)
    return b_rshift(byte, 5), byte % 32
end

local function read_object(fh, opts)
    local typ, mintyp = read_type(fh)
    return decoder[typ](fh, mintyp, opts)
end

local function read_integer(fh, mintyp)
    return read_length(fh, mintyp)
end

local function read_negative_integer(fh, mintyp)
    return -1 - read_length(fh, mintyp)
end

local function read_string(fh, mintyp)
    if mintyp ~= 31 then
        return read_bytes(fh, read_length(fh, mintyp))
    end
    local out = {}
    local i = 1
    local v = read_object(fh)
    while v ~= BREAK do
        out[i], i = v, i + 1
        v = read_object(fh)
    end
    return t_concat(out)
end

local function read_unicode_string(fh, mintyp)
    return read_string(fh, mintyp)
    -- local str = read_string(fh, mintyp);
    -- if have_utf8 and not utf8.len(str) then
    -- TODO How to handle this?
    -- end
    -- return str;
end

local function read_array(fh, mintyp, opts)
    local out = {}
    if mintyp == 31 then
        local i = 1
        local v = read_object(fh, opts)
        while v ~= BREAK do
            out[i], i = v, i + 1
            v = read_object(fh, opts)
        end
    else
        local len = read_length(fh, mintyp)
        for i = 1, len do
            out[i] = read_object(fh, opts)
        end
    end
    return out
end

local function read_map(fh, mintyp, opts)
    local out = {}
    local k
    if mintyp == 31 then
        local i = 1
        k = read_object(fh, opts)
        while k ~= BREAK do
            out[k], i = read_object(fh, opts), i + 1
            k = read_object(fh, opts)
        end
    else
        local len = read_length(fh, mintyp)
        for _ = 1, len do
            k = read_object(fh, opts)
            out[k] = read_object(fh, opts)
        end
    end
    return out
end

local tagged_decoders = {}

local function read_semantic(fh, mintyp, opts)
    local tag = read_length(fh, mintyp)
    local value = read_object(fh, opts)
    local postproc = opts and opts[tag] or tagged_decoders[tag]
    if postproc then
        return postproc(value)
    end
    return tagged(tag, value)
end

local function read_half_float(fh)
    local exponent = read_byte(fh)
    local fraction = read_byte(fh)
    local sign = exponent < 128 and 1 or -1 -- sign is highest bit

    fraction = fraction + (exponent * 256) % 1024 -- copy two(?) bits from exponent to fraction
    exponent = b_rshift(exponent, 2) % 32 -- remove sign bit and two low bits from fraction;

    if exponent == 0 then
        return sign * m_ldexp(fraction, -24)
    elseif exponent ~= 31 then
        return sign * m_ldexp(fraction + 1024, exponent - 25)
    elseif fraction == 0 then
        return sign * m_huge
    else
        return NaN
    end
end

local function read_float(fh)
    local exponent = read_byte(fh)
    local fraction = read_byte(fh)
    local sign = exponent < 128 and 1 or -1 -- sign is highest bit
    exponent = exponent * 2 % 256 + b_rshift(fraction, 7)
    fraction = fraction % 128
    fraction = fraction * 256 + read_byte(fh)
    fraction = fraction * 256 + read_byte(fh)

    if exponent == 0 then
        return sign * m_ldexp(exponent, -149)
    elseif exponent ~= 0xff then
        return sign * m_ldexp(fraction + 2 ^ 23, exponent - 150)
    elseif fraction == 0 then
        return sign * m_huge
    else
        return NaN
    end
end

local function read_double(fh)
    local exponent = read_byte(fh)
    local fraction = read_byte(fh)
    local sign = exponent < 128 and 1 or -1 -- sign is highest bit

    exponent = exponent % 128 * 16 + b_rshift(fraction, 4)
    fraction = fraction % 16
    fraction = fraction * 256 + read_byte(fh)
    fraction = fraction * 256 + read_byte(fh)
    fraction = fraction * 256 + read_byte(fh)
    fraction = fraction * 256 + read_byte(fh)
    fraction = fraction * 256 + read_byte(fh)
    fraction = fraction * 256 + read_byte(fh)

    if exponent == 0 then
        return sign * m_ldexp(exponent, -149)
    elseif exponent ~= 0xff then
        return sign * m_ldexp(fraction + 2 ^ 52, exponent - 1075)
    elseif fraction == 0 then
        return sign * m_huge
    else
        return NaN
    end
end

if s_unpack then
    function read_float(fh)
        return s_unpack(">f", read_bytes(fh, 4))
    end
    function read_double(fh)
        return s_unpack(">d", read_bytes(fh, 8))
    end
end

local function read_simple(fh, value, opts)
    if value == 24 then
        value = read_byte(fh)
    end
    if value == 20 then
        return false
    elseif value == 21 then
        return true
    elseif value == 22 then
        return null
    elseif value == 23 then
        return undefined
    elseif value == 25 then
        return read_half_float(fh)
    elseif value == 26 then
        return read_float(fh)
    elseif value == 27 then
        return read_double(fh)
    elseif value == 31 then
        return BREAK
    end
    if opts and opts.simple then
        return opts.simple(value)
    end
    return simple(value)
end

decoder[0] = read_integer
decoder[1] = read_negative_integer
decoder[2] = read_string
decoder[3] = read_unicode_string
decoder[4] = read_array
decoder[5] = read_map
decoder[6] = read_semantic
decoder[7] = read_simple

-- opts.more(n) -> want more data
-- opts.simple -> decode simple value
-- opts[int] -> tagged decoder
local function decode(s, opts)
    local fh = {}
    local pos = 1

    local more
    if type(opts) == "function" then
        more = opts
    elseif type(opts) == "table" then
        more = opts.more
    elseif opts ~= nil then
        error(("bad argument #2 to 'decode' (function or table expected, got %s)"):format(type(opts)))
    end
    if type(more) ~= "function" then
        function more()
            error("input too short")
        end
    end

    function fh:read(bytes)
        local ret = s:sub(pos, pos + bytes - 1)
        if #ret < bytes then
            ret = more(bytes - #ret, fh, opts)
            if ret then
                self:write(ret)
            end
            return self:read(bytes)
        end
        pos = pos + bytes
        return ret
    end

    function fh:write(bytes) -- luacheck: no self
        s = s .. bytes
        if pos > 256 then
            s = s:sub(pos + 1)
            pos = 1
        end
        return #bytes
    end

    return read_object(fh, opts)
end

return {
    -- en-/decoder functions
    encode = encode,
    decode = decode,
    decode_file = read_object,

    -- tables of per-type en-/decoders
    type_encoders = encoder,
    type_decoders = decoder,

    -- special treatment for tagged values
    tagged_decoders = tagged_decoders,

    -- constructors for annotated types
    simple = simple,
    tagged = tagged,

    -- pre-defined simple values
    null = null,
    undefined = undefined,
}
