#!/usr/bin/awk -f

# I wanted to implement something non-trivial using awk. 
# If I was clever I wouldn’t implement forth directly in awk,
# instead I’d implement a simple virtual machine using awk, 
# and then implement the forth using the virtual machine’s byte code
# ...but there is only so much brain power I can exert on such a silly project.

BEGIN {
    # Initialize stacks and dictionaries
    stack_ptr = 0
    dict_size = 0
    
    # Built-in words, and some documentation (I could use stack comments,
    # but I find those sort of unintuitive)
    dict["+"] = "+     : Adds the top two numbers on the stack."
    dict["-"] = "-     : Subtracts the top number from the second top number on the stack."
    dict["*"] = "*     : Multiplies the top two numbers on the stack."
    dict["/"] = "/     : Divides the second top number by the top number on the stack."
    dict["."] = ".     : Prints the top of the stack."
    dict[".s"] = ".s    : Shows all values on the stack."
    dict["dup"] = "dup   : Duplicates the top value on the stack."
    dict["drop"] = "drop  : Removes the top value from the stack."
    dict["swap"] = "swap  : Swaps the top two values on the stack."
    dict["over"] = "over  : Copies the second top value to the top of the stack."
    dict["rot"] = "rot   : Rotates the top three values on the stack."
    dict["="] = "=     : Compares the top two values for equality."
    dict["<"] = "<     : Checks if the second top value is less than the top value."
    dict[">"] = ">     : Checks if the second top value is greater than the top value."
    dict["bye"] = "bye   : Exits the interpreter."
    dict["words"] = "words : Lists all available words and their documentation."
    
    # State flags
    compiling = 0
    current_def = ""
    def_name = ""
    
    # If an input file isn't specified, enter REPL mode
    if (ARGC == 1) {
        repl()
    }
}

# Handle file input
{
    if (FILENAME ~ /\.forth$/) {
        interpret($0)
    }
}

function repl() {
    print "f.awk! A forth interpreter.\nUse 'bye' to exit.\nUse 'words' to list all available words.\n"
    while (1) {
        printf "f> "
        if (getline input < "/dev/tty" <= 0) break
        interpret(input)
    }
}

function interpret(line) {
    gsub(/\(.*\)/, "", line)  # Remove everything from ( to )
    gsub(/\\.*$/, "", line)   # Remove backslash comments, too

    n = split(line, words, /[ \t]+/)
    
    for (i = 1; i <= n; i++) {
        word = words[i]
        if (word == "") continue
        
        # print "Processing word: " word
        
        if (word == ":") {
            compiling = 1
            i++
            def_name = words[i]
            current_def = ""
            continue
        }
        
        if (compiling) {
            if (word == ";") {
                # Store user-defined word with its name and definition
                dict[def_name] = "word " current_def
                compiling = 0
                continue
            }
            current_def = current_def " " word
            continue
        }
        
        # Execute the word and skip further processing if it's .s
        if (word == ".s") {
            execute_word(word)
            break  # Exit the loop after executing .s
        }
        
        execute_word(word)
    }
}

function execute_word(word) {
    if (word ~ /^-?[0-9]+$/) {
        push(word + 0)
    } else if (word in dict) {
        if (dict[word] ~ /^word /) {
            # User-defined word
            sequence = substr(dict[word], 6)
            split(sequence, subwords, " ")
            for (sw in subwords) {
                if (subwords[sw] != "") {
                    execute_word(subwords[sw])
                }
            }
        } else {
            # Built-in words
            if (word == "+") math_add()
            else if (word == "-") math_sub()
            else if (word == "*") math_mul()
            else if (word == "/") math_div()
            else if (word == ".") stack_print()
            else if (word == ".s") {
                # print "Executing .s command"
                stack_show()
            }
            else if (word == "dup") stack_dup()
            else if (word == "drop") stack_drop()
            else if (word == "swap") stack_swap()
            else if (word == "over") stack_over()
            else if (word == "rot") stack_rot()
            else if (word == "=") compare_eq()
            else if (word == "<") compare_lt()
            else if (word == ">") compare_gt()
            else if (word == "bye") exit_program()
            else if (word == "words") list_words()
            else if (word == "if") {
                # Handle the if statement
                if_condition = pop()
                if (if_condition == 0) {
                    # Skip to the next part until we find 'then' or 'else'
                    skip_if = 1
                }
            }
            else if (word == "else") {
                # Handle the else statement
                if (skip_if) {
                    skip_if = 0  # Reset the skip flag
                } else {
                    # Skip to the next part until we find 'then'
                    skip_else = 1
                }
            }
            else if (word == "then") {
                # End of the conditional
                skip_if = 0
                skip_else = 0
            }
        }
    } else {
        print "Error: Unknown word '" word "'"
    }
}

function push(val) {
    stack[stack_ptr++] = val
}

function pop() {
    if (stack_ptr <= 0) {
        print "Error: Stack underflow"
        return 0
    }
    return stack[--stack_ptr]
}

function math_add() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a + b)
}

function math_sub() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a - b)
}

function math_mul() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a * b)
}

function math_div() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    if (b == 0) {
        print "Error: Division by zero"
        return
    }
    a = pop()
    push(int(a / b))
}

function stack_print() {
    if (stack_ptr < 1) {
        print "Error: Stack underflow"
        return
    }
    print pop()
}

function stack_show() {
    print "<", stack_ptr, "> "
    for (i = 0; i < stack_ptr; i++) {
        printf "%s ", stack[i]
    }
    print ""
    # print "Stack state after .s: "
    # for (i = 0; i < stack_ptr; i++) {
    #     print stack[i]
    # }
    # print ""
}

function stack_dup() {
    if (stack_ptr < 1) {
        print "Error: Stack underflow"
        return
    }
    val = stack[stack_ptr - 1]
    push(val)
}

function stack_drop() {
    if (stack_ptr < 1) {
        print "Error: Stack underflow"
        return
    }
    pop()
}

function stack_swap() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(b)
    push(a)
}

function stack_over() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a)
    push(b)
    push(a)
}

function stack_rot() {
    if (stack_ptr < 3) {
        print "Error: Stack underflow"
        return
    }
    c = pop()
    b = pop()
    a = pop()
    push(b)
    push(c)
    push(a)
}

function compare_eq() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a == b ? -1 : 0)
}

function compare_lt() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a < b ? -1 : 0)
}

function compare_gt() {
    if (stack_ptr < 2) {
        print "Error: Stack underflow"
        return
    }
    b = pop()
    a = pop()
    push(a > b ? -1 : 0)
}

function exit_program() {
    print "Exiting program."
    exit 0
}

function list_words() {
    print "Available words:"
    
    # Separate arrays to hold built-in and user-defined words
    split("", built_in_words)
    split("", user_defined_words)
    
    for (w in dict) {
        split(dict[w], parts, ": ")
        if (parts[1] ~ /^word /) {
            user_defined_words[w] = parts[2]
        } else {
            built_in_words[w] = parts[2]
        }
    }
    
    # Sort built-in words manually because I'm picky
    n = 0
    for (w in built_in_words) {
        sorted_words[n++] = w
    }
    
    for (i = 0; i < n; i++) {
        for (j = i + 1; j < n; j++) {
            if (sorted_words[i] > sorted_words[j]) {
                temp = sorted_words[i]
                sorted_words[i] = sorted_words[j]
                sorted_words[j] = temp
            }
        }
    }
    
    # First print the built-in words
    for (i = 0; i < n; i++) {
        print sorted_words[i] ": " built_in_words[sorted_words[i]]
    }
    
    # Then print the user-defined words
    for (w in user_defined_words) {
        print w ": " user_defined_words[w] " ( User-defined )"
    }
}