require "lib/metaprogramming.nom"
require "lib/operators.nom"
require "lib/utils.nom"

# Conditionals
compile [if %condition %if_body] to code: ".."
    |if \(%condition as lua) then
    |\(%if_body as lua statements)
    |end --end if
compile [unless %condition %body] to code: ".."
    |if not (\(%condition as lua)) then
    |\(%body as lua statements)
    |end --end if

compile [if %condition %if_body else %else_body, unless %condition %else_body else %if_body] to code: ".."
    |if \(%condition as lua) then
    |\(%if_body as lua statements)
    |else
    |\(%else_body as lua statements)
    |end --end if

# Return
compile [return] to code: "do return; end"
compile [return %return_value] to code: "do return \(%return_value as lua); end"

# GOTOs
compile [-> %label] to code: ".."
    |::label_\(nomsu "var_to_lua_identifier" [%label])::;
compile [go to %label] to code: ".."
    |goto label_\(nomsu "var_to_lua_identifier" [%label]);

rule [tree %tree has function call %call] =:
    lua> ".."
        |local target = (\(%call)).value;
        |for subtree,_ in coroutine.wrap(function() nomsu:walk_tree(\(%tree)); end) do
        |    if type(subtree) == 'table' and subtree.type == "FunctionCall"
        |    and nomsu.utils.equivalent(subtree.value, target, 2) then
        |        return true;
        |    end
        |end
        |do return false; end

# While loops
compile [do next repeat-loop] to code: "goto continue_repeat;"
compile [stop repeat-loop] to code: "goto stop_repeat;"
compile [repeat while %condition %body] to code:
    %continue_labels = (..)
        "\n::continue_repeat::;" if (tree %body has function call \(do next repeat-loop)) else ""
    %code = ".."
        |while \(%condition as lua) do
        |\(%body as lua statements)\(%continue_labels)
        |end --while-loop
    if (tree %body has function call \(stop repeat-loop)):
        return ".."
            |do --while-loop label scope
            |\(%code)
            |::stop_repeat::;
            |end --while-loop label scope
    return %code
parse [repeat %body] as: repeat while (true) %body
parse [repeat until %condition %body] as: repeat while (not %condition) %body

# For loop control flow:
compile [stop for-loop] to code: "goto stop_for;"
compile [stop %var] to code: ".."
    |goto stop_\(nomsu "var_to_lua_identifier" [%var]);
compile [do next for-loop] to code: "goto continue_for;"
compile [do next %var] to code: ".."
    |goto continue_\(nomsu "var_to_lua_identifier" [%var]);

# Numeric range for loops
compile [..]
    for %var from %start to %stop by %step %body
    for %var from %start to %stop via %step %body
..to code:
    %continue_labels = ""
    if (tree %body has function call \(do next for-loop)):
        %continue_labels join= "\n::continue_for::;"
    if (tree %body has function call (nomsu "replaced_vars" [\(do next %), =lua "{['']=\(%var)}"])):
        %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
    %code = ".."
        |for i=\(%start as lua),\(%stop as lua),\(%step as lua) do
        # This trashes the loop variables, just like in Python.
        |\(%var as lua) = i;
        |\(%body as lua statements)\(%continue_labels)
        |end --numeric for-loop
    %stop_labels = ""
    if (tree %body has function call \(stop for-loop)):
        %stop_labels join= "\n::stop_for::;"
    if (tree %body has function call (nomsu "replaced_vars" [\(stop %), =lua "{['']=\(%var)}"])):
        %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%var])::;"
    if (%stop_labels != ""): ".."
        |do --for-loop label scope
        |\(%code)\(%stop_labels)
        |end --for-loop label scope
    ..else: %code
parse [for %var from %start to %stop %body] as: for %var from %start to %stop via 1 %body
parse [..]
    for all %start to %stop by %step %body
    for all %start to %stop via %step %body
..as: for % from %start to %stop via %step %body
parse [for all %start to %stop %body] as: for all %start to %stop via 1 %body

compile [for %var in %iterable %body] to code:
    %continue_labels = ""
    if (tree %body has function call \(do next for-loop)):
        %continue_labels join= "\n::continue_for::;"
    if (tree %body has function call (nomsu "replaced_vars" [\(do next %), =lua "{['']=\(%var)}"])):
        %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
    %code = ".."
        |for i,value in ipairs(\(%iterable as lua)) do
        # This trashes the loop variables, just like in Python.
        |\(%var as lua) = value;
        |\(%body as lua statements)\(%continue_labels)
        |end --foreach-loop
    %stop_labels = ""
    if (tree %body has function call \(stop for-loop)):
        %stop_labels join= "\n::stop_for::;"
    if (tree %body has function call (nomsu "replaced_vars" [\(stop %), =lua "{['']=\(%var)}"])):
        %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%var])::;"
    if (%stop_labels != ""): ".."
        |do --for-loop label scope
        |\(%code)\(%stop_labels)
        |end --for-loop label scope
    ..else: %code
parse [for all %iterable %body] as: for % in %iterable %body


# Switch statement/multi-branch if
compile [when %body] to code:
    %result = ""
    %fallthroughs = []
    %first = (yes)
    for %func_call in (%body's "value"):
        assert ((%func_call's "type") == "FunctionCall") ".."
            |Invalid format for 'when' statement. Only '*' blocks are allowed.
        %tokens = (%func_call's "value")
        %star = (%tokens -> 1)
        assert (=lua "vars.star and vars.star.type == 'Word' and vars.star.value == '*'") ".."
            |Invalid format for 'when' statement. Lines must begin with '*'

        %condition = (%tokens -> 2)
        assert %condition ".."
            |Invalid format for 'when' statement. Lines must begin with '*' and have a condition or the word "else"

        %action = (%tokens -> 3)
        if (%action == (nil)):
            lua do> "table.insert(vars.fallthroughs, vars.condition)"
            do next %func_call

        if (=lua "vars.condition.type == 'Word' and vars.condition.value == 'else'"):
            %result join= ".."
                |
                |else
                |\(%action as lua statements)
            stop for-loop
        ..else:
            %condition = (%condition as lua)
            for all %fallthroughs:
                %condition join= " or \(% as lua)"
            %result join= ".."
                |
                |\("if" if %first else "elseif") \(%condition) then
                |\(%action as lua statements)

        %fallthroughs = []
        %first = (no)

    if (%result != ""):
        %result join= "\nend"
    %result

# Switch statement
compile [when %branch_value == ? %body] to code:
    %result = ""
    %fallthroughs = []
    %first = (yes)
    for %func_call in (%body's "value"):
        assert ((%func_call's "type") == "FunctionCall") ".."
            |Invalid format for 'when' statement. Only '*' blocks are allowed.
        %tokens = (%func_call's "value")
        %star = (%tokens -> 1)
        assert (=lua "vars.star and vars.star.type == 'Word' and vars.star.value == '*'") ".."
            |Invalid format for 'when' statement. Lines must begin with '*'

        %condition = (%tokens -> 2)
        assert %condition ".."
            |Invalid format for 'when' statement. Lines must begin with '*' and have a condition or the word "else"

        %action = (%tokens -> 3)
        if (%action == (nil)):
            lua> "table.insert(vars.fallthroughs, vars.condition)"
            do next %func_call

        if (=lua "vars.condition.type == 'Word' and vars.condition.value == 'else'"):
            %result join= ".."
                |
                |else
                |\(%action as lua statements)
            stop for-loop
        ..else:
            %condition = "branch_value == (\(%condition as lua))"
            for all %fallthroughs:
                %condition join= " or (branch_value == \(% as lua))"
            %result join= ".."
                |
                |\("if" if %first else "elseif") \(%condition) then
                |\(%action as lua statements)

        %fallthroughs = []
        %first = (no)

    if (%result != ""):
        %result = ".."
            |do --when == ?
            |local branch_value = \(%branch_value as lua);\(%result)
            |end
            |end --when == ?
    %result

# Try/except
compile [..]
    try %action and if it succeeds %success or if it fails %fallback
    try %action and if it fails %fallback or if it succeeds %success
..to code: ".."
    |do
    |    local fell_through = false;
    |    local ok, ret1, ret2 = pcall(function(nomsu, vars)
    |        \(%action as lua statements)
    |        fell_through = true;
    |    end, nomsu, vars);
    |    if ok then
    |        \(%success as lua statements)
    |    end
    |    if not ok then
    |        \(%fallback as lua statements)
    |    elseif not fell_through then
    |        return ret1, ret2;
    |    end
    |end
parse [try %action] as:
    try %action and if it succeeds {pass} or if it fails {pass}
parse [try %action and if it fails %fallback] as:
    try %action and if it succeeds {pass} or if it fails %fallback
parse [try %action and if it succeeds %success] as:
    try %action and if it succeeds %success or if it fails {pass}

# Do/finally:
compile [do %action then always %final_action] to code: ".."
    |do
    |    local fell_through = false;
    |    local ok, ret1, ret2 = pcall(function(nomsu, vars)
    |        \(%action as lua statements)
    |        fell_through = true;
    |    end, nomsu, vars);
    |    local ok2, _ = pcall(function(nomsu, vars)
    |        \(%final_action as lua statements)
    |    end, nomsu, vars);
    |    if not ok then nomsu:error(ret1); end
    |    if not ok2 then nomsu:error(ret2); end
    |    if not fell_through then
    |        return ret1, ret2;
    |    end
    |end