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

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

    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
immediately:
    compile [return] to code: "do return; end"
    compile [return %return_value] to code: "do return \(%return_value as lua); end"

# GOTOs
immediately:
    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]);

# Basic loop control
immediately:
    compile [do next] to code: "continue;"
    compile [stop] to code: "break;"

# Helper function
immediately:
    action [tree %tree has function call %call]:
        lua> ".."
            local target = (\%call).stub;
            for subtree,_ in coroutine.wrap(function() nomsu:walk_tree(\%tree); end) do
                if type(subtree) == 'table' and subtree.type == "FunctionCall"
                and subtree.stub == target then
                    return true;
                end
            end
            return false;

# While loops
immediately:
    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:
        set %continue_labels = (..)
            "\n::continue_repeat::;" if (tree %body has function call \(do next repeat-loop)) else ""
        set %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:
immediately:
    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
immediately:
    compile [..]
        for %var from %start to %stop by %step %body
        for %var from %start to %stop via %step %body
    ..to code:
        local [%continue_labels, %code, %stop_labels]
        set %continue_labels = ""
        if (tree %body has function call \(do next for-loop)):
            %continue_labels join= "\n::continue_for::;"
        if (tree %body has function call (tree \(do next %) with {""=%var})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%var])::;"
        # This trashes the loop variables, just like in Python.
        set %code = ".."
            for i=\(%start as lua),\(%stop as lua),\(%step as lua) do
            \(%var as lua) = i;
            \(%body as lua statements)\
            ..\%continue_labels
            end --numeric for-loop
        set %stop_labels = ""
        if (tree %body has function call \(stop for-loop)):
            %stop_labels join= "\n::stop_for::;"
        if (tree %body has function call (tree \(stop %) with {""=%var})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%var])::;"
        return (..)
            ".."
                do --for-loop label scope
                \%code\
                ..\%stop_labels
                end --for-loop label scope
            ..if (%stop_labels is not "") else %code

immediately:
    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

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

immediately:
    compile [..]
        repeat %n times %body, repeat %n x %body
    ..to code:
        local [%continue_labels, %code, %stop_labels]
        set %continue_labels = ""
        if (tree %body has function call \(do next repeat-loop)):
            %continue_labels join= "\n::continue_repeat::;"
        # This trashes the loop variables, just like in Python.
        set %code = ".."
            for i=1,\(%n as lua) do
            \(%body as lua statements)\
            ..\%continue_labels
            end --numeric for-loop
        set %stop_labels = ""
        if (tree %body has function call \(stop repeat-loop)):
            %stop_labels join= "\n::stop_repeat::;"
        return (..)
            ".."
                do --repeat-loop label scope
                \%code\
                ..\%stop_labels
                end --repeat-loop label scope
            ..if (%stop_labels is not "") else %code

# Dict iteration (lua's "pairs()")
immediately:
    compile [for %key = %value in %iterable %body] to code:
        set %continue_labels = ""
        if (tree %body has function call \(do next for-loop)):
            %continue_labels join= "\n::continue_for::;"
        if (tree %body has function call (tree \(do next %x) with {x=%key})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%key])::;"
        if (tree %body has function call (tree \(do next %) with {""=%value})):
            %continue_labels join= "\n::continue_\(nomsu "var_to_lua_identifier" [%value])::;"
        # This trashes the loop variables, just like in Python.
        set %code = ".."
            for key,value in pairs(\(%iterable as lua)) do
            \(%key as lua), \(%value as lua) = key, value;
            \(%body as lua statements)\
            ..\%continue_labels
            end --foreach-loop
        set %stop_labels = ""
        if (tree %body has function call \(stop for-loop)):
            %stop_labels join= "\n::stop_for::;"
        if (tree %body has function call (tree \(stop %) with {""=%key})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%key])::;"
        if (tree %body has function call (tree \(stop %) with {""=%value})):
            %stop_labels join= "\n::stop_\(nomsu "var_to_lua_identifier" [%value])::;"
        return (..)
            ".."
                do --for-loop label scope
                \%code\
                ..\%stop_labels
                end --for-loop label scope
            ..if (%stop_labels is not "") else %code

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

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

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

            if (=lua "\%condition.type == 'Word' and \%condition.value == 'else'"):
                %result join= ".."
                    
                    else
                    \(%action as lua statements)
                stop for-loop
            ..else:
                set %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)

            set %fallthroughs = []
            set %first = (no)

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

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

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

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

            if (=lua "\%condition.type == 'Word' and \%condition.value == 'else'"):
                %result join= ".."
                    
                    else
                    \(%action as lua statements)
                stop for-loop
            ..else:
                set %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)

            set %fallthroughs = []
            set %first = (no)

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

# Try/except
immediately:
    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)
                \(%action as lua statements)
                fell_through = true;
            end, nomsu);
            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: do nothing
        ..or if it fails: do nothing
    parse [try %action and if it fails %fallback] as:
        try %action and if it succeeds: do nothing
        ..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: do nothing

# Do/finally:
immediately:
    compile [do %action then always %final_action] to code: ".."
        do
            local fell_through = false;
            local ok, ret1, ret2 = pcall(function(nomsu)
                \(%action as lua statements)
                fell_through = true;
            end, nomsu);
            local ok2, _ = pcall(function(nomsu)
                \(%final_action as lua statements)
            end, nomsu);
            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