setmetatable(_ENV, { __index=lpeg })

Scopes = { {} }

function eval_expr(expr)
    local accum = eval(expr[2]) -- because 1 is "expr"
    for i = 3, #expr, 2 do
        local operator = expr[i]
        local num2 = eval(expr[i+1])

        if operator == '+' then
            accum = accum + num2
        elseif operator == '-' then
            accum = accum - num2
        elseif operator == '*' then
            accum = accum * num2
        elseif operator == '/' then
            accum = accum / num2
        end
    end
    return accum
end

function eval_bool(expr)
    local num1 = eval(expr[2])
    local operator = expr[3]
    local num2 = eval(expr[4])

    if operator == '<' then
        return num1 < num2
    elseif operator == '<=' then
        return num1 <= num2
    elseif operator == '>' then
        return num1 > num2
    elseif operator == '>=' then
        return num1 >= num2
    elseif operator == '==' then
        return num1 == num2
    elseif operator == '!=' then
        return num1 ~= num2
    end
end

function eval(ast)
    if type(ast) == 'number' then
        return ast
    elseif ast[1] == 'expr' or ast[1] == 'term' then
        return eval_expr(ast)
    elseif ast[1] == 'array' then
        local new = {}
        for _, el in ipairs(ast[2]) do
            table.insert(new, eval(el))
        end
        return new
    elseif ast[1] == 'ref' then
        return lookup(ast)
    elseif ast[1] == 'assign' then
        return assign(ast[2], eval(ast[3]))
    elseif ast[1] == 'list' then
        local last = nil
        for i = 2, #ast do
            last = eval(ast[i])
        end
        return last
    elseif ast[1] == 'if' then
        if eval_bool(ast[2]) then
            return eval(ast[3])
        end
    elseif ast[1] == 'while' then
        while eval_bool(ast[2]) do
            eval(ast[3])
        end
    elseif ast[1] == 'function' then
        return { args=ast[2],
                 code=ast[3],
                 scope=Scopes[#Scopes] }
    elseif ast[1] == 'call' then
        local fn = eval(ast[2])
        local scope = setmetatable({}, {__index=fn.scope})

        for i, name in ipairs(fn.args) do
            scope[name] = eval(ast[3][i])
        end

        table.insert(Scopes, scope)
        local result = eval(fn.code)
        table.remove(Scopes)
        return result
    end
end

function assign(ref, value)
    local current = Scopes[#Scopes]
    for i = 2, #ref do
        local next_index = ref[i]
        if type(next_index) == 'table' then
            next_index = eval(next_index)
        end

        if i == #ref then -- last one, set the value
            -- Special case, assign to something in an outer scope
            while current[next_index] and not rawget(current, next_index) do
                -- Walk the scope chain back until we see it
                current = getmetatable(current).__index
            end

            current[next_index] = value
            return value
        else -- not the last, keep following the chain
            current = current[next_index]
        end
    end
end

function lookup(ref)
    local current = Scopes[#Scopes]
    for i = 2, #ref do
        local next_index = ref[i]
        if type(next_index) == 'table' then
            next_index = eval(next_index)
        end

        current = current[next_index]
    end
    return current
end

spc = S(" \t\n")^0

digit = R('09')
number = C( (P("-") + digit) *
                digit^0 *
                ( P('.') * digit^0 )^-1 ) / tonumber * spc

lparen = "(" * spc
rparen = ")" * spc

lbrack = "[" * spc
rbrack = "]" * spc

lcurly = "{" * spc
rcurly = "}" * spc
comma = "," * spc

expr_op = C( S('+-') ) * spc
term_op = C( S('*/') ) * spc
letter = R('AZ','az')
name = C( letter * (digit+letter+"_")^0 ) * spc
keywords = (P("if") + P("while") + P("function")) * spc
name = name - keywords

boolean = C( S("<>") + "<=" + ">=" + "!=" + "==" ) * spc

stmt = spc * P{
    "LIST";
    LIST =
        V("STMT") +
        Ct( Cc("list") *
                lcurly *
                (V("STMT") * P(";")^0 * spc)^0 *
                rcurly ),
    STMT = 
        Ct( Cc("assign") * V("REF") * "=" * spc * V("VAL") ) +
        V("IF") +
        V("WHILE") +
        V("CALL") +
        V("FUNCTION") +
        V("EXPR"),
    EXPR = Ct( Cc("expr") * V("TERM") * ( expr_op * V("TERM") )^0 ),
    TERM = Ct( Cc("term") * V("FACT") * ( term_op * V("FACT") )^0 ),
    REF = Ct( Cc("ref") * name * (lbrack * V("EXPR") * rbrack)^0 ),
    FACT =
        number +
        lparen * V("EXPR") * rparen +
        V("REF"),
    ARRAY = Ct( Cc("array") * lbrack * V("VAL_LIST") * rbrack ),
    VAL_LIST = Ct( (V("VAL") * comma^-1)^0 ),
    VAL = V("CALL") + V("ARRAY") + V("FUNCTION") + V("EXPR"),
    BOOL = Ct( Cc("bool") * V("EXPR") * boolean * V("EXPR") ),
    IF = Ct( C("if") * spc * lparen * V("BOOL") * rparen * V("LIST") ),
    WHILE = Ct( C("while") * spc * lparen * V("BOOL") * rparen * V("LIST") ),
    CALL = Ct( Cc("call") * V("REF") * spc * lparen * V("VAL_LIST") * rparen ),
    FUNCTION = Ct( C("function") * spc * lparen * V("NAME_LIST") * rparen * V("LIST") ),
    NAME_LIST = Ct( (name * comma^0)^0 )
}

function test(stmt)
    local Global = Scopes[1]

    stmt = stmt / eval
    assert(stmt:match(" 1 + 2 ") == 3)
    assert(stmt:match("1+2+3+4+5") == 15)
    assert(stmt:match("2*3*4 + 5*6*7") == 234)
    assert(stmt:match(" 1 * 2 + 3") == 5)
    assert(stmt:match("( 2 +2) *6") == 24)
    stmt:match("a=3"); assert(Global.a == 3)
    assert(stmt:match("a") == 3)
    assert(stmt:match("a * 5") == 15); Global.a=nil
    stmt:match("a = [ 4, 5, 6 ]");
    assert(Global.a[1] == 4)
    assert(Global.a[2] == 5)
    assert(Global.a[3] == 6)
    Global.a=nil
    stmt:match("b = [ ]");
    assert(Global.b[1] == nil)
    Global.b=nil
    stmt:match("c = [[1,2], [3,4]]")
    assert(Global.c[1][1] == 1)
    assert(Global.c[1][2] == 2)
    assert(Global.c[2][1] == 3)
    assert(Global.c[2][2] == 4)
    assert(stmt:match("c[4/2][1]") == 3)
    stmt:match("c[3] = 5")
    assert(Global.c[3] == 5)
    Global.c=nil
    stmt:match("if(1 < 0) b = 5"); assert(Global.b ~= 5)
    Global.n=0; Global.x=1
    stmt:match("while(n < 8) { x = x * 2; n = n + 1 }")
    assert(Global.x == 256)
    Global.n=nil; Global.x=nil

    stmt:match("f = function(a) a*2+3")
    assert(Global.f ~= nil)
    assert(Global.f.scope == Scopes[1])

    stmt:match("res = f(5)")
    assert(Global.res == 13)
    assert(Global.a == nil)
    assert(#Scopes == 1)

    Scopes[1] = {}; Global = Scopes[1]
    stmt:match("make_acc = function(){ a = 0; function(n) a = a+n; }")
    assert(Global.make_acc)
    stmt:match("acc = make_acc()")
    assert(stmt:match("acc(1)") == 1)
    assert(stmt:match("acc(1)") == 2)
end

function repl(file)
    file = file or io.input()
    parser = stmt
    for line in file:lines() do
        print(parser:match(line))
    end
end