#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "ast.h"

typedef struct {
    unsigned* i;
    Array memory;
} Memory;

struct Cell {
    char* name;
    unsigned index;
    struct Cell* next;
};

typedef struct {
    struct Cell* next;
} Scope;

Memory new_memory() {
    unsigned size = 1024;
    unsigned* i = (unsigned*)malloc(sizeof(unsigned));
    *i = 0;
    int* mem = (int*)malloc(sizeof(int) * size);
    return (Memory){
        .i = i,
        .memory = {
            .length = size,
            .elements = mem,
        }
    };
}

Scope new_scope() {
    return (Scope) { .next = NULL };
}

int lookup_scope(char* name, Scope scope) {
    struct Cell* cell = scope.next;
    while (cell != NULL) {
        if (strcmp(cell->name, name) == 0) {
            return cell->index;
        } else {
            cell = cell->next;
        }
    }
    return -1;
}

Scope insert(char* name, int n, Memory memory, Scope scope) {
    if (*memory.i + 1 < memory.memory.length) {
        ((int*)(memory.memory.elements))[*memory.i] = n;
        struct Cell* cell = (struct Cell*)malloc(sizeof(struct Cell));
        *cell = (struct Cell) {
            .name = name,
            .index = *memory.i,
            .next = scope.next,
        };
        ++*memory.i;
        Scope new_scope = (Scope) {.next = cell};
        return new_scope;
    }
    fprintf(stderr, "Error: out of memory.");
    exit(1);
}

int* lookup_mem(int index, Memory mem) {
    if (index >= 0 && (unsigned)index < *(mem.i)) {
        return &((int*)(mem.memory.elements))[index];
    }
    return NULL;
}

int* lookup(char* name, Memory memory, Scope scope) {
    int index = lookup_scope(name, scope);
    if (index >= 0) {
        int* result = lookup_mem(index, memory);
        if (result != NULL) {
            return result;
        }
    }
    return NULL;
}

int eval_expr(Expr expr, Memory memory, Scope scope) {
    switch (expr.tag) {
    default:
    case LITERAL: {
        return expr.data.integer;
    }
    case VARIABLE: {
        int* result = lookup(expr.data.variable, memory, scope);
        if (result == NULL) {
            fprintf(stderr, "Error: variable not found '%s'\n", expr.data.variable);
            exit(1);
        } else {
            return *result;
        }
    }
    case FUNCTION: {
        if (strcmp(expr.data.function.name, "print") == 0) {
            if (expr.data.function.args.length == 1) {
                int arg = eval_expr(((Expr*)expr.data.function.args.elements)[0], memory, scope);
                printf("%d\n", arg);
                return 0;
            } else {
                fprintf(stderr, "Error: print expects a single argument.\n");
                exit(1);
            }
        }
        else if (strcmp(expr.data.function.name, "add") == 0) {
            if (expr.data.function.args.length == 2) {
                int arg1 = eval_expr(((Expr*)expr.data.function.args.elements)[0], memory, scope);
                int arg2 = eval_expr(((Expr*)expr.data.function.args.elements)[1], memory, scope);
                return arg1 + arg2;
            } else {
                fprintf(stderr, "Error: negate expects a single argument.\n");
                exit(1);
            }
        }
        else if (strcmp(expr.data.function.name, "negate") == 0) {
            if (expr.data.function.args.length == 1) {
                int arg = eval_expr(((Expr*)expr.data.function.args.elements)[0], memory, scope);
                return 0 - arg;
            } else {
                fprintf(stderr, "Error: negate expects a single argument.\n");
                exit(1);
            }
        }
        return 0;
    }
    }
}

Scope interpret_stmt(Stmt stmt, Memory memory, Scope scope) {
    switch (stmt.tag) {
    case SET: {
        int result = eval_expr(stmt.data.Set.expr, memory, scope);
        int* index = lookup(stmt.data.Set.name, memory, scope);
        if (index != NULL) {
            *index = result;
        } else {
            return insert(stmt.data.Set.name, result, memory, scope);
        }
        break;
    }
    case WHILE: {
        while (eval_expr(stmt.data.While.condition, memory, scope)) {
            Scope new_scope = scope;
            for (unsigned i = 0; i < stmt.data.While.block.length; ++i) {
                Stmt current = ((Stmt*)stmt.data.While.block.elements)[i];
                new_scope = interpret_stmt(current, memory, new_scope);
            }
			// free the part of the scope we are done with
			struct Cell* next = new_scope.next;
			while (next != NULL && next != scope.next) {
				struct Cell* nextnext = next->next;
				free(next);
				next = nextnext;
			}
        }
        break;
    }
    case EXPR:
        eval_expr(stmt.data.Expr.expr, memory, scope);
        break;
    }
    return scope;
}

void execute(StmtArray stmts) {
    Memory memory = new_memory();
    Scope scope = new_scope();
    for (unsigned pc = 0; pc < stmts.length; ++pc) {
        Stmt stmt = stmts.stmts[pc];
        scope = interpret_stmt(stmt, memory, scope);
    }
}