#include <avr/io.h>
#include <util/delay.h>
#include <string.h>
#include <stdlib.h>

#define F_CPU 16000000UL       // CPU Frequency
#define BAUD 19200             // UART Baud Rate
#define MYUBRR F_CPU/16/BAUD-1 // UBRR Value

// Stack Definitions
#define STACK_SIZE 256         // Stack size for data operations
int16_t stack[STACK_SIZE];
int8_t sp = -1;                // Stack pointer (-1 = empty stack)

// Return Stack for Loops and Control Structures
#define RETURN_STACK_SIZE 256
int16_t return_stack[RETURN_STACK_SIZE];
int8_t rsp = -1;               // Return stack pointer

// Dictionary Definitions
#define DICTIONARY_SIZE 64
#define WORD_BODY_SIZE 256
typedef void (*WordFunction)();

typedef struct {
    const char *name;          // Word name
    WordFunction function;     // Word implementation
} Word;

Word dictionary[DICTIONARY_SIZE];
WordFunction word_body[WORD_BODY_SIZE];
int8_t dict_size = 0;           // Number of words in the dictionary
int16_t word_body_index = 0;    // Index for compiled word instructions

char compiled_word_name[32];    // Name of the word being compiled
int16_t compiled_word_start = 0; // Start of the compiled word body
int8_t compiling = 0;           // Compilation mode flag

// UART Functions
void uart_init(unsigned int ubrr) {
    UBRR0H = (unsigned char)(ubrr >> 8);
    UBRR0L = (unsigned char)ubrr;
    UCSR0B = (1 << RXEN0) | (1 << TXEN0); // Enable UART
    UCSR0C = (1 << UCSZ01) | (1 << UCSZ00); // 8-bit data
}

void uart_transmit(unsigned char data) {
    while (!(UCSR0A & (1 << UDRE0))); // Wait for transmit buffer
    UDR0 = data;
}

unsigned char uart_receive(void) {
    while (!(UCSR0A & (1 << RXC0))); // Wait for received data
    return UDR0;
}

void uart_print(const char *str) {
    while (*str) uart_transmit(*str++);
}

// Stack Operations
void push(int16_t value) {
    if (sp < STACK_SIZE - 1) {
        stack[++sp] = value;
    } else {
        uart_print("Stack Overflow\n");
    }
}

int16_t pop() {
    if (sp >= 0) {
        return stack[sp--];
    } else {
        uart_print("Stack Underflow\n");
        return 0;
    }
}

// Return Stack Operations
void rpush(int16_t value) {
    if (rsp < RETURN_STACK_SIZE - 1) {
        return_stack[++rsp] = value;
    } else {
        uart_print("Return Stack Overflow\n");
    }
}

int16_t rpop() {
    if (rsp >= 0) {
        return return_stack[rsp--];
    } else {
        uart_print("Return Stack Underflow\n");
        return 0;
    }
}

// Dictionary Management
void add_word(const char *name, WordFunction function) {
    if (dict_size < DICTIONARY_SIZE) {
        dictionary[dict_size].name = name;
        dictionary[dict_size].function = function;
        dict_size++;
    } else {
        uart_print("Dictionary Full\n");
    }
}

void execute_word(const char *name) {
    for (int i = 0; i < dict_size; i++) {
        if (strcmp(dictionary[i].name, name) == 0) {
            dictionary[i].function();
            return;
        }
    }
    uart_print("Unknown word: ");
    uart_print(name);
    uart_print("\n");
}

// Interpreter and Tokenizer
void interpret(const char *input) {
    char buffer[64];
    strncpy(buffer, input, sizeof(buffer));
    buffer[sizeof(buffer) - 1] = '\0';

    char *token = strtok(buffer, " ");
    while (token != NULL) {
        // Handle numeric input
        char *endptr;
        int16_t value = strtol(token, &endptr, 10);
        if (*endptr == '\0') { // Valid number
            push(value);
        } else if (strcmp(token, ":") == 0) { // Start compilation
            compiling = 1;
            strncpy(compiled_word_name, strtok(NULL, " "), sizeof(compiled_word_name));
            compiled_word_start = word_body_index;
        } else if (compiling && strcmp(token, ";") == 0) { // End compilation
            compiling = 0;
            add_word(compiled_word_name, *(word_body + compiled_word_start));
        } else {
            execute_word(token); // Try to execute as a word
        }
        token = strtok(NULL, " ");
    }
}

// Core Words
void word_add() { push(pop() + pop()); }
void word_subtract() { int16_t b = pop(); push(pop() - b); }
void word_multiply() { push(pop() * pop()); }
void word_divide() { int16_t b = pop(); if (b != 0) push(pop() / b); else uart_print("Division by zero\n"); }
void word_mod() { int16_t b = pop(); if (b != 0) push(pop() % b); else uart_print("Division by zero\n"); }
void word_dot() { char buffer[16]; itoa(pop(), buffer, 10); uart_print(buffer); uart_print("\n"); }
void word_cr() { uart_print("\n"); }

// Ultrasonic Sensor Words
void word_pin_high() {
    int16_t pin = pop();
    PORTB |= (1 << pin); // Set pin high
}

void word_pin_low() {
    int16_t pin = pop();
    PORTB &= ~(1 << pin); // Set pin low
}

void word_read_pin() {
    int16_t pin = pop();
    push((PINB & (1 << pin)) ? 1 : 0); // Push pin state
}

void word_delay_us() {
    int16_t us = pop();
    while (us--) _delay_us(1);
}

void word_measure_time() {
    int16_t pin = pop();
    uint16_t count = 0;

    // Wait for pin to go high
    while (!(PINB & (1 << pin))) {
        count++;
        if (count > 60000) { // Timeout
            push(0);
            return;
        }
    }

    count = 0;
    // Measure pulse width
    while (PINB & (1 << pin)) {
        count++;
        _delay_us(1); // 1 µs increment
    }

    push(count); // Push measured time in µs
}

// Main Function
int main() {
    char buffer[64];
    uint8_t buffer_pos = 0;

    uart_init(MYUBRR);
    uart_print("Tiny Forth for AVR\n");

    // Add Words
    add_word("+", word_add);
    add_word("-", word_subtract);
    add_word("*", word_multiply);
    add_word("/", word_divide);
    add_word("MOD", word_mod);
    add_word(".", word_dot);
    add_word("CR", word_cr);

    add_word("PIN-HIGH", word_pin_high);
    add_word("PIN-LOW", word_pin_low);
    add_word("READ-PIN", word_read_pin);
    add_word("DELAY-US", word_delay_us);
    add_word("MEASURE-TIME", word_measure_time);

    // Main Interpreter Loop
    while (1) {
        char c = uart_receive();
        if (c == '\n' || c == '\r') {
            buffer[buffer_pos] = '\0';
            interpret(buffer);
            buffer_pos = 0;
        } else {
            if (buffer_pos < sizeof(buffer) - 1) {
                buffer[buffer_pos++] = c;
            }
        }
    }
}