--
-- A primitive code transpiler
--

-- Find multiple patterns
local function find_multiple(text, ...)
    local n = select('#', ...)
    local s, e, pattern
    for i = 1, n do
        local p = select(i, ...)
        local s2, e2 = text:find(p)
        if s2 and (not s or s2 < s) then
            s, e, pattern = s2, e2 or s2, p
        end
    end
    return s, e, pattern
end

-- Matches
-- These take 2-3 arguments (code, res, char) and should return code and res.
local lua51 = not asyncio or asyncio.lua51
local matches
matches = {
    -- Handle multi-line strings and comments
    ['%[=*%['] = function(code, res, char)
        res = res .. char
        if char:sub(1, 2) == '--' then
            char = char:sub(4, -2)
        else
            char = char:sub(2, -2)
        end
        local s, e = code:find(']' .. char .. ']', nil, true)
        if not s or not e then return code, res end
        return code:sub(e + 1), res .. code:sub(1, e)
    end,

    -- Handle regular comments
    ['--'] = function(code, res, char)
        res = res .. char
        local s, e = code:find('\n', nil, true)
        if not s or not e then return code, res end
        return code:sub(e + 1), res .. code:sub(1, e)
    end,

    -- Handle quoted text
    ['"'] = function(code, res, char)
        res = res .. char

        -- Handle backslashes
        repeat
            local s, e, pattern = find_multiple(code, '\\', char)
            if pattern then
                res = res .. code:sub(1, e + 1)
                code = code:sub(e + 2)
            end
        until not pattern or pattern == char

        return code, res
    end,

    -- Handle "async function" statements
    ['async function'] = function(code, res, char)
        if not code:find('end', nil, true) then return code, res .. char end
        local s = code:find(')', nil, true)
        if not s then return code, res .. char end
        if char:sub(1, 1) ~= 'a' then res = res .. char:sub(1, 1) end

        local params = code:match('^%s*%(([^%)]*)%)')
        if params then
            res = res .. 'asyncio.coroutine(function(' .. params
        else
            -- Check for method functions
            local name, params = code:match('^%s*([%w_%.]+)%s*%(([^%)]*)%)')
            if not name or not params then
                local n, n2, p = code:match(
                    '^%s*([%w_%.]+)%:([%w_]+)%s*%(([^%)]*)%)')
                if not n or not n2 or not p then
                    return code, res .. char
                end
                name = n .. '.' .. n2
                if p:find('%S') then
                    params = 'self, ' .. p
                else
                    params = 'self'
                end
            end

            -- Local functions
            if res:sub(-6, -2) == 'local' then res = res .. name .. '; ' end

            -- Add the first part
            res = res .. name .. ' = asyncio.coroutine(function(' .. params
        end
        code = code:sub(s)

        -- Search for the "end"
        if code:sub(#code) == 'd' then code = code .. '\n' end
        local lvl = 1
        while lvl > 0 do
            local s, e, pattern = find_multiple(code, '[\'"\\]', '%-%-%[=*%[',
                '%-%-', '%[=*%[', '^async function ', '[%s;%[%(]async function',
                '^async for ', '[%s;%[]async for ', '^wait until ',
                '[%s;%[]wait until ', '%sdo%s', '[%s%)]then%s',
                '[%s%(%)]function[%s%(]', '%send[%s%),]', '^end[%s%),]')
            if not s then return code .. ')', res end

            -- Add non-matching characters
            res = res .. code:sub(1, math.max(s - 1, 0))

            local char = code:sub(s, e)
            if pattern == '%sdo%s' or pattern == '[%s%(%)]function[%s%(]' or
                    pattern == '[%s%)]then%s' then
                lvl = lvl + 1
                res = res .. char
                code = code:sub(e + 1)
            elseif pattern == '%send[%s%),]' or pattern == '^end[%s%),]' then
                lvl = lvl - 1
                res = res .. char
                code = code:sub(e + 1)
            else
                -- Call the correct function
                local func = matches[char] or matches[pattern]
                assert(func, 'No function found for pattern!')
                code, res = func(code:sub(e + 1), res, char)

                -- Increase the level for async for loops
                if lua51 and (pattern == '^async for ' or
                        pattern == '[%s;%[]async for ') then
                    lvl = lvl + 1
                end
            end
        end

        return code, res:sub(1, -2) .. ')' .. res:sub(-1)
    end,

    -- Handle "async for" statements
    ['async for '] = function(code, res, char)
        if char:sub(1, 1) ~= 'a' then res = res .. char:sub(1, 1) end
        local var, func, n = code:match(
            '^([%w_,%s]+)%s+in%s*([%w_%.]+)%s*%(%s*(%S)')
        if not var or not func then return code, res .. 'for ' end
        local s = code:find('(', nil, true)
        code = code:sub(s + 1)
        res = res .. 'for ' .. var .. ' in asyncio.iter(' .. func
        if n ~= ')' then res = res .. ', ' end
        return code, res
    end,

    ['wait until '] = function(code, res, char)
        if char:sub(1, 1) ~= 'w' then res = res .. char:sub(1, 1) end
        return code, res .. 'repeat asyncio.defer() until '
    end,
}

-- A workaround for Lua 5.1's iterator limitations.
if lua51 then
    local orig = matches['async for ']
    matches['async for '] = function(code, res, char)
        if char:sub(1, 1) ~= 'a' then res = res .. char:sub(1, 1) end
        local var, func, params = code:match(
            '^([%w_,%s]+)%s+in%s*([%w_%.]+)%s*%(%s*([^%(%)]*%))')

        -- TODO: Handle everything
        local s, e = code:find('%)%s*do')
        if not var or not func or not params or params:sub(-1) ~= ')' or
                not e then
            return orig(code, res, 'async for ')
        end

        code = code:sub(e + 1)

        -- Generate code
        s = var:find(',', nil, true)
        local i, v = '__asyncio_lua51_iter_shim', var:sub(1, (s or 2) - 1)
        if params ~= ')' then params = ', ' .. params end
        res = res .. 'local ' .. i .. ' = asyncio._iter_start(' .. func ..
            params .. '; while true do local ' .. var .. ' = ' .. i ..
            '(); if ' .. v .. ' == nil then break end;'

        return code, res
    end
end

-- Give the functions alternate names
matches['%-%-%[=*%['], matches["'"] = matches['%[=*%['], matches['"']
matches['[%s;%[%(]async function'] = matches['async function']
matches['[%s;%[]async for '] = matches['async for ']
matches['[%s;%[]wait until '] = matches['wait until ']

-- The actual transpiler
local function transpile(code)
    assert(type(code) == 'string')

    local res = ''

    -- Split the code by "tokens"
    while true do
        -- Search for special characters
        local s, e, pattern = find_multiple(code, '[\'"\\]', '%-%-%[=*%[',
            '%-%-', '%[=*%[', '^async function', '[%s;%[%(]async function',
            '^async for ', '[%s;%[]async for ', '^wait until ',
            '[%s;%[]wait until ')
        if not s then break end

        -- Add non-matching characters
        res = res .. code:sub(1, math.max(s - 1, 0))

        -- Call the correct function
        local char = code:sub(s, e)
        local func = matches[char] or matches[pattern]
        assert(func, 'No function found for pattern!')
        code, res = func(code:sub(e + 1), res, char)
    end

    return res .. code
end

-- Handle being called outside of Minetest.
if not asyncio or not minetest then
    -- If this is called via dofile(), just return the function.
    if not arg or debug.getinfo(3) then
        return transpile
    end

    if not arg[1] or arg[1] == '' then
        print('Invalid usage.')
        print('Usage: transpile.lua <file>')
        os.exit(1)
    end

    -- Read the file
    local f, err, errcode = io.open(arg[1], 'r')
    if not f then
        print(err)
        os.exit(errcode)
    end

    local code, err, errcode = f:read('*a')
    f:close()
    if not code then
        print(err)
        os.exit(errcode)
    end

    -- Print the output
    print(transpile(code))
    os.exit(0)
end

-- Lua 5.2+ compatibility
local loadstring = assert(rawget(_G, 'loadstring') or load)

-- Load strings
function asyncio.loadstring(code, ...)
    local f, msg = loadstring(transpile(code), ...)
    return f and asyncio.coroutine(f), msg
end

-- Load strings from files
function asyncio.loadfile(file)
    local f, msg = io.open(file, 'r')
    if not f then return nil, msg end
    return asyncio.loadstring(f:read('*a'), '@' .. file)
end

-- Execute files
function asyncio.dofile(file)
    local func, msg = asyncio.loadfile(file)
    if not func then error(msg, 0) end
    return func()
end

-- Execute a string
function asyncio.exec(...)
    local func, msg = asyncio.loadstring(...)
    if not func then error(msg, 2) end
    return func()
end

-- Load a file in the current mod
local DIR_DELIM = rawget(_G, 'DIR_DELIM') or '/'
function asyncio.domodfile(file)
    local prefix = minetest.get_modpath(minetest.get_current_modname())
    return asyncio.dofile(prefix .. DIR_DELIM .. file)
end