--
-- Minetest asyncio: Core
--

local coroutines = {}
local iterators = {}

-- Make coroutines and iterators "weak tables".
setmetatable(coroutines, {__mode = 'k'})
setmetatable(iterators,  {__mode = 'kv'})

-- Resume a coroutine
local last_error
local function resume(co, ...)
    local good, msg = coroutine.resume(co, ...)
    local status = coroutine.status(co)

    if status == 'dead' then coroutines[co], iterators[co] = nil, nil end

    if not good then
        assert(status == 'dead')
        if last_error then
            last_error.msg = msg
            minetest.after(0, last_error.cleanup)
        end
        minetest.log('error', 'Error in coroutine: ' .. tostring(msg))
        for _, line in ipairs(debug.traceback(co):split('\n')) do
            minetest.log('error', line)
        end
    elseif iterators[co] then
        if msg == nil then return false end
        return msg
    end
end

-- Run coroutines directly
local function run(func, ...)
    local thread = coroutine.running()
    if coroutines[thread] then
        return func(...)
    end
    local co = coroutine.create(func)
    coroutines[co] = true
    resume(co, ...)
end

function asyncio.run(func, ...)
    if type(func) ~= 'function' then
        error('Invalid function passed to asyncio.run.', 2)
    end
    return run(func, ...)
end

-- Wrap a function
function asyncio.coroutine(func)
    if type(func) ~= 'function' then
        error('Invalid function passed to asyncio.coroutine.', 2)
    end

    return function(...)
        return run(func, ...)
    end
end

-- Run a coroutine in a separate "thread"
function asyncio.run_async(func, ...)
    local co = coroutine.create(func)
    coroutines[co] = true
    resume(co, ...)
end

-- Make sure functions are being run inside coroutines.
local function get_coroutine()
    local co = coroutine.running()
    if not coroutines[co] then
        error('asyncio function called outside coroutine.', 3)
    end
    return co
end

-- This table is used internally for iterators
local suspend = {}
function asyncio.await()
    get_coroutine()
    return coroutine.yield(suspend)
end

-- Create asyncio.resume on-the-fly when required.
setmetatable(asyncio, {__index = function(_, name)
    if name == 'resume' then
        local co = coroutine.running()
        if not coroutines[co] then return end
        while iterators[co] do co = iterators[co] end
        return function(...)
            if coroutines[co] then
                resume(co, ...)
            end
        end
    elseif name == 'running' then
        return coroutines[coroutine.running()] and true or false
    end
end})

-- Wait a specified amount of seconds before continuing.
function asyncio.sleep(delay)
    if type(delay) ~= 'number' or delay ~= delay or delay >= math.huge then
        error('Invalid asyncio.sleep() invocation.', 2)
    end

    asyncio.await(minetest.after(delay, asyncio.resume))
end

-- Defer execution until the next globalstep.
function asyncio.defer()
    asyncio.await(minetest.after(0, asyncio.resume))
end

-- Only defer if it has been 5ms since the last optional_defer()
local next_defer = minetest.get_us_time()
function asyncio.optional_defer()
    if minetest.get_us_time() > next_defer then
        asyncio.defer()
        next_defer = minetest.get_us_time() + 5000
        return true
    end
    return false
end
asyncio.odefer = asyncio.optional_defer

-- Pack and unpack using nil
local function pack_len(...)
    return select('#', ...), {...}
end

local function unpack_len(length, t, i)
    i = i or 1
    if i <= length then
        return t[i], unpack_len(length, t, i + 1)
    end
end

-- asyncio.yield returns a table
function asyncio.yield(...)
    if not iterators[get_coroutine()] then
        error('asyncio.yield() called outside of async iterator.', 2)
    end

    local l, t = pack_len(...)
    t.length = l
    coroutine.yield(t)
end

-- Iterate over a string
local function str_iter(s)
    for n = 1, #s do
        asyncio.yield(n, s:sub(n, n))
    end
end

-- The standard iteration function (LuaJIT, Lua 5.2+)
local rawequal = rawequal
local function iter(func, ...)
    local co = get_coroutine()
    if type(func) == 'string' and select('#', ...) == 0 then
        return iter(str_iter, func)
    elseif type(func) ~= 'function' then
        error('Invalid asyncio.iter() invocation.')
    end

    local iterco = coroutine.create(func)
    coroutines[iterco], iterators[iterco] = true, co
    local l, m = pack_len(...)

    return function()
        if iterco == nil then return end
        assert(get_coroutine() == co)
        local msg = resume(iterco, unpack_len(l, m))
        l, m = 0, nil

        -- Suspend the parent coroutine if required
        while msg and rawequal(msg, suspend) do
            msg = resume(iterco, coroutine.yield(suspend))
        end

        -- Unpack asyncio.yield()s
        if type(msg) == 'table' and type(msg.length) == 'number' then
            if msg[1] == nil then
                coroutines[iterco], iterators[iterco] = nil, nil
                iterco = nil
                return
            end
            return unpack_len(msg.length, msg)
        end

        -- Return the result
        if msg == nil then
            coroutines[iterco], iterators[iterco], iterco = nil, nil, nil
        end
        return msg
    end
end
asyncio.iter = iter
asyncio._iter_start = iter

-- Lua 5.1 (excluding LuaJIT) compatibility
asyncio.lua51 = _VERSION < 'Lua 5.2'
if asyncio.lua51 then
    asyncio.lua51 = not pcall(coroutine.wrap(function()
        for i in coroutine.yield do return end
    end))
end

if asyncio.lua51 then
    last_error = {}
    assert(loadfile((minetest.get_modpath('asyncio') ..
        '/lua51.lua')))(last_error)
end

-- Defer pairs() and ipairs() until the next globalstep
function asyncio.pairs(n, defer_every)
    if not iterators[get_coroutine()] then
        error('asyncio.pairs() called outside of an async iterator.', 2)
    end

    if type(defer_every) == 'number' and defer_every == defer_every and
            defer_every > 1 then
        local i = 0
        for k, v in pairs(n) do
            i = i + 1
            if i >= defer_every then asyncio.defer(); i = 0 end
            asyncio.yield(k, v)
        end
    else
        for k, v in pairs(n) do
            asyncio.optional_defer()
            asyncio.yield(k, v)
        end
    end
end

function asyncio.ipairs(n, defer_every)
    if not iterators[get_coroutine()] then
        error('asyncio.ipairs() called outside of an async iterator.', 2)
    end

    if type(defer_every) == 'number' and defer_every == defer_every and
            defer_every > 1 then
        asyncio.defer()
        defer_every = math.ceil(defer_every)
        for k, v in ipairs(n) do
            if k % defer_every == 0 then asyncio.defer() end
            asyncio.yield(k, v)
        end
    else
        -- Defer every 5ms
        for k, v in ipairs(n) do
            asyncio.optional_defer()
            asyncio.yield(k, v)
        end
    end
end