local chupacabra = {}

-- Parse a string into a list of tokens
function chupacabra.tokenize(s)
    local tokens = {}
    local array = nil
    for token in s:gmatch("%S+") do
        if array then
            array = array .. " " .. token
            if token:find("]$") then
                table.insert(tokens, array)
                array = nil
            end
        elseif token:find("^%[") then
            array = token
        else
            table.insert(tokens, token)
        end
    end
    return tokens
end

function chupacabra.evaluate(tokens, context)
    local stack = {}

    -- FIXME: explore if there is a better way to approach this, so that you don't need to special case the array math
    for _, token in ipairs(tokens) do
        if tonumber(token) then
            table.insert(stack, tonumber(token))
        elseif token:match("^%b[]$") then -- If the token is an array
            local array = {}
            for number in token:sub(2, -2):gmatch("%S+") do
                table.insert(array, tonumber(number))
            end
            table.insert(stack, array)
        elseif token == ".." then
            local array = {}
            for i = 1, #stack do
                table.insert(array, stack[i])
            end
            table.insert(stack, array)
        elseif token == "@.." then
            local result = {}
            for i = 1, #stack do
                local value = stack[i]
                if type(value) == "table" then
                    for j = 1, #value do
                        table.insert(result, value[j])
                    end
                else
                    table.insert(result, value)
                end
            end
            table.insert(stack, result)    
        elseif token == "." then
            table.remove(stack)
            --TODO: add the ability for this keywor, or a paired @. keyword to remove an array from the stack
        elseif token == "+" then
            local b = table.remove(stack)
            local a = table.remove(stack)
            table.insert(stack, a + b)
        elseif token == "-" then
            local b = table.remove(stack)
            local a = table.remove(stack)
            table.insert(stack, a - b)
        elseif token == "*" then
            local b = table.remove(stack)
            local a = table.remove(stack)
            table.insert(stack, a * b)
        elseif token == "/" then
            local b = table.remove(stack)
            local a = table.remove(stack)
            table.insert(stack, a / b)
        elseif token == "@+" then
            local a = table.remove(stack)
            local b = table.remove(stack)
            if type(a) == "table" and type(b) == "table" then
            if #a ~= #b then
                error("Arrays must have equal length")
            end
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] + b[i])
            end
            table.insert(stack, result)
            elseif type(a) == "table" and type(b) == "number" then
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] + b)
            end
            table.insert(stack, result)
            elseif type(a) == "number" and type(b) == "table" then
            local result = {}
            for i = 1, #b do
                table.insert(result, a + b[i])
            end
            table.insert(stack, result)
            else
            error("Invalid operands for addition")
            end
        elseif token == "@*" then
            local a = table.remove(stack)
            local b = table.remove(stack)
            if type(a) == "table" and type(b) == "table" then
            if #a ~= #b then
                error("Arrays must have equal length")
            end
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] * b[i])
            end
            table.insert(stack, result)
            elseif type(a) == "table" and type(b) == "number" then
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] * b)
            end
            table.insert(stack, result)
            elseif type(a) == "number" and type(b) == "table" then
            local result = {}
            for i = 1, #b do
                table.insert(result, a * b[i])
            end
            table.insert(stack, result)
            else
            error("Invalid operands for multiplication")
            end
        elseif token == "@-" then
            local a = table.remove(stack)
            local b = table.remove(stack)
            if type(a) == "table" and type(b) == "table" then
            if #a ~= #b then
                error("Arrays must have equal length")
            end
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] - b[i])
            end
            table.insert(stack, result)
            elseif type(a) == "table" and type(b) == "number" then
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] - b)
            end
            table.insert(stack, result)
            elseif type(a) == "number" and type(b) == "table" then
            local result = {}
            for i = 1, #b do
                table.insert(result, a - b[i])
            end
            table.insert(stack, result)
            else
            error("Invalid operands for subtraction")
            end
        elseif token == "@/" then
            local a = table.remove(stack)
            local b = table.remove(stack)
            if type(a) == "table" and type(b) == "table" then
            if #a ~= #b then
                error("Arrays must have equal length")
            end
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] / b[i])
            end
            table.insert(stack, result)
            elseif type(a) == "table" and type(b) == "number" then
            local result = {}
            for i = 1, #a do
                table.insert(result, a[i] / b)
            end
            table.insert(stack, result)
            elseif type(a) == "number" and type(b) == "table" then
            local result = {}
            for i = 1, #b do
                table.insert(result, a / b[i])
            end
            table.insert(stack, result)
        else
            error("Invalid operands for division")
            end
        end
    end

    return table.remove(stack)
end

function chupacabra.run(str, context)
    local tokens = chupacabra.tokenize(str)
    return chupacabra.evaluate(tokens, context)
end

return chupacabra