#!/usr/bin/awk -f

# Forth interpreter in AWK

BEGIN {
    print "Welcome to the AWK Forth Interpreter!"
    print "Type your commands below. Use 'bye' to quit."
    # Initialize variables
    top = -1  # Initialize stack pointer
    
    # Initialize the dictionary with basic words
    words["+"] = "+"
    words["-"] = "-"
    words["*"] = "*"
    words["/"] = "/"
    words["dup"] = "dup"
    words["over"] = "over"
    words["swap"] = "swap"
    words["."] = "."
    words["bye"] = "bye"
    words["rot"] = "rot"
    words["drop"] = "drop"
    words["nip"] = "nip"
    words["tuck"] = "tuck"
    words["roll"] = "roll"
    words["pick"] = "pick"
    words["negate"] = "negate"
    words["abs"] = "abs"
    words["max"] = "max"
    words["min"] = "min"
    words["mod"] = "mod"
    words["="] = "="
    words["see"] = "see"
    words["if"] = "if"
    words["then"] = "then"
    words["else"] = "else"
    words[">"] = ">"

    # Add descriptions for words
    desc["+"] = "( n1 n2 -- sum ) Add top two numbers"
    desc["-"] = "( n1 n2 -- diff ) Subtract top number from second"
    desc["*"] = "( n1 n2 -- prod ) Multiply top two numbers"
    desc["/"] = "( n1 n2 -- quot ) Divide second by top"
    desc["dup"] = "( n -- n n ) Duplicate top of stack"
    desc["over"] = "( n1 n2 -- n1 n2 n1 ) Copy second item to top"
    desc["swap"] = "( n1 n2 -- n2 n1 ) Swap top two items"
    desc["rot"] = "( n1 n2 n3 -- n2 n3 n1 ) Rotate top three items"
    desc["drop"] = "( n -- ) Discard top item"
    desc["nip"] = "( n1 n2 -- n2 ) Remove second item"
    desc["tuck"] = "( n1 n2 -- n2 n1 n2 ) Copy top item below second"
    desc["roll"] = "( nk ... n1 n0 k -- nk-1 ... n1 n0 nk ) Move kth item to top"
    desc["pick"] = "( nk ... n1 n0 k -- nk ... n1 n0 nk ) Copy kth item to top"
    desc["negate"] = "( n -- -n ) Negate number"
    desc["abs"] = "( n -- |n| ) Absolute value"
    desc["max"] = "( n1 n2 -- max ) Maximum of top two numbers"
    desc["min"] = "( n1 n2 -- min ) Minimum of top two numbers"
    desc["mod"] = "( n1 n2 -- rem ) Remainder of n1/n2"
    desc["="] = "( n1 n2 -- flag ) Test if equal, leaves 1 if true, 0 if false"
    desc["if"] = "( flag -- ) Begin conditional execution"
    desc["then"] = "( -- ) End conditional execution"
    desc["else"] = "( -- ) Execute if previous condition was false"
    desc[">"] = "( n1 n2 -- flag ) Returns true if n1 is greater than n2"

    # Add variable-related words to dictionary
    words["variable"] = "variable"
    words["!"] = "!"
    words["@"] = "@"
    words["?"] = "?"

    # Add descriptions for the new words
    desc["variable"] = "( -- addr ) Create a variable"
    desc["!"] = "( n addr -- ) Store n at addr"
    desc["@"] = "( addr -- n ) Fetch contents of addr"
    desc["?"] = "( addr -- ) Print contents of addr"

    # Initialize test-related words
    words["test"] = "test"
    words["testing"] = "testing"
    desc["test"] = "( n1 n2 -- ) Compare values and report test result"
    desc["testing"] = "( -- ) Print test description"

    # Initialize condition stack
    cond_top = -1

    # Mark these as compile-only words
    compile_only["if"] = 1
    compile_only["then"] = 1
    compile_only["else"] = 1

    # Add dot-quote to dictionary
    words[".\""] = ".\""
    desc[".\""] = "( -- ) Print following string up to closing quote"

    # Handle file input
    if (ARGC > 1) {
        print "Loading file:", ARGV[1]
        while ((getline < ARGV[1]) > 0) {
            if ($0 ~ /^[[:space:]]*$/) continue  # Skip empty lines
            print "> " $0
            process_line()
        }
        close(ARGV[1])
        exit
    }

    # Add to dictionary initialization
    words["<"] = "<"
    words[">"] = ">"
    
    # Add descriptions
    desc["<"] = "( n1 n2 -- flag ) Returns true if n1 is less than n2"
    desc[">"] = "( n1 n2 -- flag ) Returns true if n1 is greater than n2"

    # Add handlers
    handlers["+" ] = "add"
    handlers["-" ] = "subtract"
    handlers["*" ] = "multiply"
    handlers["/" ] = "divide"
    handlers["dup"] = "dup"
    handlers["over"] = "over"
    handlers["swap"] = "swap"
    handlers["."] = "print_top"
    handlers["<"] = "less_than"
    handlers[">"] = "greater_than"

    # Add variable handler
    handlers["variable"] = "create_variable"
    handlers["!"] = "store"
    handlers["@"] = "fetch"
    handlers["?"] = "print_var"
}

# Add these helper functions near the top
function check_stack(min_items, error_msg) {
    if (top < min_items - 1) {
        print error_msg ? error_msg : "Error: Not enough values on stack"
        return 0
    }
    return 1
}

function check_underflow() {
    return check_stack(1, "Error: Stack underflow")
}

# Function to define a new word
function define_word(name, definition) {
    # Store the raw definition
    raw_definitions[name] = definition
    print "DEBUG: Raw definition for " name ": " definition
    
    # Process the definition to handle conditionals
    processed = process_definition(definition)
    print "DEBUG: Processed definition for " name ": " processed
    words[name] = processed
}

function process_definition(definition) {
    # Normalize whitespace in definition
    gsub(/^[[:space:]]+/, "", definition)
    gsub(/[[:space:]]+$/, "", definition)
    gsub(/[[:space:]]+/, " ", definition)
    
    # Split definition into words
    split(definition, def_words, " ")
    processed = ""
    in_string = 0
    string_content = ""
    
    for (i = 1; i <= length(def_words); i++) {
        word = def_words[i]
        
        if (word == ".\"") {
            # Start collecting string
            in_string = 1
            string_content = ""
            continue
        }
        
        if (in_string) {
            # Check if this word contains the closing quote
            if (word ~ /"$/) {
                # End of string found
                string_content = string_content (string_content == "" ? "" : " ") substr(word, 1, length(word)-1)
                processed = processed " <string>" string_content "</string>"
                in_string = 0
            } else {
                # Add to string content
                string_content = string_content (string_content == "" ? "" : " ") word
            }
            continue
        }
        
        # Handle non-string words
        if (word == "if") {
            processed = processed " <if>"
        } else if (word == "else") {
            processed = processed " <else>"
        } else if (word == "then") {
            processed = processed " <then>"
        } else {
            processed = processed " " word
        }
    }
    
    return processed
}

# Stack to hold values
function push(value) {
    stack[++top] = value
}

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

function binary_op(operation) {
    if (!check_stack(2)) return
    second = pop()
    first = pop()
    if (operation == "+") push(first + second)
    else if (operation == "-") push(first - second)
    else if (operation == "*") push(first * second)
    else if (operation == "/") {
        if (second == 0) {
            print "Error: Division by zero"
            push(first)
            push(second)
            return
        }
        push(first / second)
    }
    else if (operation == "mod") push(first % second)
    else if (operation == "=") push(first == second ? 1 : 0)
    else if (operation == "<") push(first < second ? 1 : 0)
    else if (operation == ">") push(first > second ? 1 : 0)
}

# Then simplify the operation functions:
function add() { binary_op("+") }
function subtract() { binary_op("-") }
function multiply() { binary_op("*") }
function divide() { binary_op("/") }
function mod() { binary_op("mod") }
function equals() { binary_op("=") }
function less_than() { binary_op("<") }
function greater_than() { binary_op(">") }

# Function to duplicate the top value
function dup() {
    if (!check_stack(1)) return
    push(stack[top])
}

# Function to copy the second value to the top
function over() {
    if (!check_stack(2)) return
    push(stack[top - 1])
}

# Function to swap the top two values
function swap() {
    if (!check_stack(2)) return
    temp = pop()
    second = pop()
    push(temp)
    push(second)
}

function print_top() {
    if (!check_stack(1)) return
    print stack[top]
}

# Function to list all available words
function list_words() {
    print "Available words:"
    for (word in words) {
        print word
    }
}

function rot() {
    if (!check_stack(3)) return
    third = pop()
    second = pop()
    first = pop()
    push(second)
    push(third)
    push(first)
}

function drop() {
    if (!check_stack(1)) return
    top--
}

function nip() {
    if (!check_stack(2)) return
    temp = stack[top]
    drop()
    drop()
    push(temp)
}

function tuck() {
    if (!check_stack(2)) return
    temp = pop()
    second = pop()
    push(temp)
    push(second)
    push(temp)
}

function roll() {
    if (!check_stack(1)) return
    n = int(pop())
    if (!check_stack(n)) return
    if (n <= 0) return

    temp = stack[top - n + 1]
    for (i = top - n + 1; i < top; i++) {
        stack[i] = stack[i + 1]
    }
    stack[top] = temp
}

function pick() {
    if (!check_stack(1)) return
    n = int(pop())
    if (!check_stack(n)) return
    if (n < 0) return
    push(stack[top - n])
}

function negate() {
    if (!check_stack(1)) return
    push(-pop())
}

function abs() {
    if (!check_stack(1)) return
    n = pop()
    push(n < 0 ? -n : n)
}

function max() {
    if (!check_stack(2)) return
    b = pop()
    a = pop()
    push(a > b ? a : b)
}

function min() {
    if (!check_stack(2)) return
    b = pop()
    a = pop()
    push(a < b ? a : b)
}

function see(word) {
    if (word in desc) {
        print desc[word]
    } else if (word in words) {
        print ": " word " " words[word] " ;"
    } else {
        print "Word '" word "' not found"
    }
}

# Add these new functions for variable handling:

# Function to create a new variable
function create_variable(name) {
    if (name == "") {
        print "Error: Variable name required"
        return
    }
    variables[name] = 0  # Initialize variable to 0
    # No need to push the variable name onto the stack
}

# Function to store a value in a variable
function check_variable(addr, error_msg) {
    if (!(addr in variables)) {
        print error_msg ? error_msg : "Error: Invalid variable '" addr "'"
        return 0
    }
    return 1
}

function store() {
    if (!check_stack(2)) return
    addr = pop()
    if (!check_variable(addr)) return
    variables[addr] = pop()
}

# Function to fetch a value from a variable
function fetch() {
    if (!check_stack(1)) return
    addr = pop()
    if (!check_variable(addr)) return
    push(variables[addr])
}

# Function to print a variable's value
function print_var() {
    if (top < 0) {
        print "Error: Stack underflow"
        return
    }
    addr = pop()
    if (!(addr in variables)) {
        print "Error: Invalid variable '" addr "'"
        return
    }
    print variables[addr]
}

# Add these test-related functions:
function run_test(expected) {
    if (top < 0) {
        print "❌ FAIL: Stack is empty"
        return
    }
    actual = pop()
    if (actual == expected) {
        print "✓ PASS"
    } else {
        print "❌ FAIL: Expected", expected, "but got", actual
    }
}

function print_test_description() {
    if (NF < 2) {
        print "Error: testing requires a description"
        return
    }
    # Collect all words after 'testing' as the description
    test_desc = ""
    for (i = 2; i <= NF; i++) {
        test_desc = test_desc " " $i
    }
    print "\nTesting:" test_desc
}

# Add these new functions for conditional handling:

# Function to handle if statement
function handle_if() {
    if (top < 0) {
        print "Error: Stack underflow"
        return
    }
    # Save condition without consuming it
    cond_stack[++cond_top] = (stack[top] != 0)
}

# Function to handle then statement
function handle_then() {
    if (cond_top < 0) {
        print "Error: Unmatched then"
        return
    }
    cond_top--  # Pop condition
}

# Function to handle else statement
function handle_else() {
    if (cond_top < 0) {
        print "Error: Unmatched else"
        return
    }
    cond_stack[cond_top] = !cond_stack[cond_top]
}

# Add this function
function execute_word(word) {
    if (word in handlers) {
        if (handlers[word] == "add") add()
        else if (handlers[word] == "subtract") subtract()
        else if (handlers[word] == "multiply") multiply()
        else if (handlers[word] == "divide") divide()
        else if (handlers[word] == "dup") dup()
        else if (handlers[word] == "over") over()
        else if (handlers[word] == "swap") swap()
        else if (handlers[word] == "print_top") print_top()
        else if (handlers[word] == "create_variable") {
            # Special handling for variable creation
            if (i + 1 <= NF) {
                create_variable($(++i))
            } else {
                print "Error: variable requires a name"
            }
        }
        else if (handlers[word] == "store") store()
        else if (handlers[word] == "fetch") fetch()
        else if (handlers[word] == "print_var") print_var()
        else {
            print "Error: Handler '" handlers[word] "' not implemented"
            return 0
        }
        return 1
    }
    return 0
}

# Move the main command processing into a function
function process_line() {
    # Process input only if it's not empty
    if (NF > 0) {
        # Remove comments and normalize whitespace
        gsub(/\(.*\)/, "")  # Remove everything from ( to )
        gsub(/^[[:space:]]+/, "")  # Remove leading whitespace
        gsub(/[[:space:]]+$/, "")  # Remove trailing whitespace
        gsub(/[[:space:]]+/, " ")  # Normalize internal whitespace

        # Special handling for testing command
        if ($1 == "testing") {
            # Collect the description string
            if ($2 ~ /^".*"$/) {  # Check if the second token is a quoted string
                test_desc = substr($2, 2, length($2) - 2)  # Remove quotes
                print "\nTesting: " test_desc
            } else {
                print "Error: testing requires a description in quotes"
            }
            return
        }

        # Check for word definitions
        if ($1 == ":") {
            word_name = $2
            definition = ""
            in_definition = 1
            # Remove : and word name from the input
            for (i = 3; i <= NF; i++) {
                definition = definition " " $i
            }
            # If we don't find a semicolon, keep reading lines
            while (definition !~ /;$/) {  # Changed to match semicolon at end
                if ((getline) <= 0) break
                gsub(/\(.*\)/, "")  # Remove comments
                gsub(/^[[:space:]]+/, "")
                gsub(/[[:space:]]+$/, "")
                definition = definition " " $0
            }
            # Remove the semicolon
            sub(/[[:space:]]*;[[:space:]]*$/, "", definition)
            
            if (definition != "") {
                define_word(word_name, definition)
                print "Defined new word: " word_name
            }
            in_definition = 0
            return
        }

        # Process remaining input
        for (i = 1; i <= NF; i++) {
            if ($i ~ /^-?[0-9]+$/) {
                push($i)
            } else if ($i in variables) {  # Check if it's a variable
                push(variables[$i])  # Push the value of the variable onto the stack
            } else if ($i in words) {  # Check if it's a defined word
                if (!in_definition && compile_only[$i]) {
                    print "Error: '" $i "' is compile-only"
                    return
                }
                if (!execute_word($i)) {
                    print "Error: Unknown command '" $i "'"
                }
            } else if ($i != "") {  # Ignore empty tokens
                print "Error: Unknown command '" $i "'"
            }
        }
    }
}

# Main loop to read commands (for interactive mode)
{
    process_line()
}

# Add function to handle dot-quote
function print_string(str) {
    printf "%s", str
}