#include "lisp.h"

static Scope global_scope = {0};

static int scope_lookup(Scopes *scopes, Nob_String_View name, Node *result) {
    nob_da_foreach(Scope *, scope, scopes) {
        nob_da_foreach(ScopeObject, scope_object, *scope) {
            if (nob_sv_eq(name, scope_object->name)) {
                *result = scope_object->node;
                return 0;
            }
        }
    }

    printf("Undefined symbol: '" SV_Fmt "'!\n", SV_Arg(name));
    return 1;
}

static void scope_add_function(Scope *scope, const char *name,
                               LispFunction function) {
    Nob_String_View obj_name = nob_sv_from_cstr(name);
    ScopeObject obj = (ScopeObject){
        .name = obj_name,
        .node = (Node){.kind = NODE_KIND_FUNCTION,
                       .as_function = {.name = obj_name, .apply = function}}};
    nob_da_append(scope, obj);
}

static bool node_is_true(Node node) {
    switch (node.kind) {
    case NODE_KIND_NUMBER:
        return node.as_number != 0;
    case NODE_KIND_LIST:
        return node.as_list->count > 0;
    default:
        return true;
    }
}

void build_scopes(Scopes *scopes) {
    scope_add_function(&global_scope, "+", lisp_function_add);
    scope_add_function(&global_scope, "-", lisp_function_sub);
    scope_add_function(&global_scope, "/", lisp_function_div);
    scope_add_function(&global_scope, "*", lisp_function_mul);
    scope_add_function(&global_scope, "=", lisp_function_eq);
    scope_add_function(&global_scope, "<", lisp_function_lt);
    scope_add_function(&global_scope, ">", lisp_function_gt);
    scope_add_function(&global_scope, "<=", lisp_function_lt_eq);
    scope_add_function(&global_scope, ">=", lisp_function_gt_eq);

    scope_add_function(&global_scope, "println", lisp_function_println);
    scope_add_function(&global_scope, "list", lisp_function_list);
    scope_add_function(&global_scope, "append", lisp_function_append);
    scope_add_function(&global_scope, "length", lisp_function_length);
    scope_add_function(&global_scope, "null?", lisp_function_isnull);
    scope_add_function(&global_scope, "builtins", lisp_function_builtins);

    scope_add_function(&global_scope, "floor", lisp_function_floor);
    scope_add_function(&global_scope, "ceil", lisp_function_ceil);

    nob_da_append(scopes, &global_scope);
}

const char *token_kind_name(TokenKind kind) {
    switch (kind) {
    case TOKEN_KIND_PAR_LEFT:
        return "PAR_LEFT";
    case TOKEN_KIND_PAR_RIGHT:
        return "PAR_RIGHT";
    case TOKEN_KIND_NUMBER:
        return "NUMBER";
    case TOKEN_KIND_SYMBOL:
        return "SYMBOL";
    default:
        UNREACHABLE("Unknown token type");
    }
}

void node_print(Node node) {
    switch (node.kind) {
    case NODE_KIND_NUMBER: {
        printf("%f", node.as_number);
    } break;
    case NODE_KIND_SYMBOL: {
        printf(SV_Fmt, SV_Arg(node.as_symbol));
    } break;
    case NODE_KIND_LIST: {
        Nodes *children = node.as_list;
        printf("(");
        for (size_t i = 0; i < children->count; i++) {
            Node node = children->items[i];
            node_print(node);
            if (i != children->count - 1) {
                printf(" ");
            }
        }
        printf(")");
    } break;
    case NODE_KIND_FUNCTION: {
        Nob_String_View func_name = node.as_function.name;
        printf(SV_Fmt, SV_Arg(func_name));
    } break;
    case NODE_KIND_LAMBDA: {
        printf("<lambda>");
    } break;
    default:
        UNREACHABLE("Unknown node type");
    }
}

int tokenize(Nob_String_View content, Tokens *tokens) {
    bool ignored = false;
    for (size_t i = 0; i < content.count; i++) {
        char c = content.data[i];

        if (ignored) {
            if (c == '\n')
                ignored = false;
            continue;
        }

        int kind = -1;
        switch (c) {
        case '/': {
            if (i + 1 < content.count && content.items[i + 1] == '/') {
                i++;
                ignored = true;
                continue;
            }
        } break;
        case '(': {
            kind = TOKEN_KIND_PAR_LEFT;
        } break;
        case ')': {
            kind = TOKEN_KIND_PAR_RIGHT;
        } break;
        }

        if (kind != -1) {
            Token tok = (Token){.kind = kind, .data = {0}};
            nob_da_append(tokens, tok);
            continue;
        }

        Nob_String_Builder sb = {0};
        bool condition = isdigit(c);
        while (condition) {
            nob_sb_appendf(&sb, "%c", c);
            c = content.data[++i];
            condition = isdigit(c) || c == '.';
        }

        if (sb.count > 0) {
            i--;

            Token tok =
                (Token){.kind = TOKEN_KIND_NUMBER, .data = nob_sb_to_sv(sb)};
            nob_da_append(tokens, tok);
            continue;
        }

        // symbols shouldn't start with a digit
        condition = isascii(c) && !isdigit(c) && !isblank(c) && !iscntrl(c) &&
                    c != '(' && c != ')';
        while (condition) {
            nob_sb_appendf(&sb, "%c", c);
            c = content.data[++i];
            condition = isascii(c) && !isblank(c) && !iscntrl(c) && c != '(' &&
                        c != ')';
        }

        if (sb.count > 0) {
            i--;

            Token tok =
                (Token){.kind = TOKEN_KIND_SYMBOL, .data = nob_sb_to_sv(sb)};
            nob_da_append(tokens, tok);
        }
    }
    return 0;
}

int parse_one(Tokens *tokens, Node *node) {
    if (tokens == NULL || node == NULL || tokens->count == 0) {
        printf("unexpected end of input");
        return 1;
    }

    // equivalent to queue.poll
    // TODO: implement a 'da_poll' macro
    Token tok = tokens->items[0];
    tokens->items++;
    tokens->count--;

    switch (tok.kind) {
    case TOKEN_KIND_PAR_LEFT: {
        Nodes *nodes = malloc(sizeof(Nodes));
        if (nodes == NULL) {
            printf("Failed to allocate memory, something is very wrong!\n");
            return 1;
        }
        memset(nodes, 0, sizeof(Nodes));

        while (true) {
            if (tokens->count == 0) {
                printf("expected closing ')'");
                return 1;
            }

            if (tokens->items[0].kind == TOKEN_KIND_PAR_RIGHT) {
                tokens->items++;
                tokens->count--;
                break;
            }

            int ret = parse_one(tokens, node);
            if (ret != 0)
                return ret;

            nob_da_append(nodes, *node);
        }

        node->kind = NODE_KIND_LIST;
        node->as_list = nodes;
    } break;
    case TOKEN_KIND_PAR_RIGHT: {
        printf("unexpected ')'");
        return 1;
    }
    case TOKEN_KIND_NUMBER: {
        double num = strtod(tok.data.data, NULL);
        if (errno != 0) {
            printf("Failed to read number!\n");
            return 1;
        }

        node->kind = NODE_KIND_NUMBER;
        node->as_number = num;
    } break;
    case TOKEN_KIND_SYMBOL: {
        node->kind = NODE_KIND_SYMBOL;
        node->as_symbol = tok.data;
    } break;
    default: {
        // unknown token kind, therefore an error
        // should I use UNREACHABLE instead?
        return 1;
    }
    }
    return 0;
}

int parse(Tokens *tokens, Nodes *nodes) {
    if (tokens == NULL || nodes == NULL)
        return 1;

    struct {
        size_t *items;
        size_t count;
        size_t capacity;
    } stack = {0};

    size_t last_idx = 0;
    for (size_t i = 0; i < tokens->count; i++) {
        TokenKind kind = tokens->items[i].kind;
        if (kind == TOKEN_KIND_PAR_LEFT)
            nob_da_append(&stack, i);
        else if (kind == TOKEN_KIND_PAR_RIGHT)
            last_idx = nob_da_pop(&stack);

        if (stack.count == 0) {
            size_t sub_size = tokens->count - last_idx;
            Tokens toks = {.count = sub_size,
                           .capacity = sub_size,
                           .items = tokens->items + last_idx};

            Node result = {0};
            int ret = parse_one(&toks, &result);
            if (ret != 0) {
                return ret;
            }
            nob_da_append(nodes, result);
        }
    }
    return 0;
}

int eval_one(Node input, Scopes *scopes, Node *result) {
    if (scopes == NULL || result == NULL)
        return 1;

    switch (input.kind) {
    case NODE_KIND_NUMBER:
        *result = input;
        return 0;
    case NODE_KIND_SYMBOL: {
        return scope_lookup(scopes, input.as_symbol, result);
    }
    case NODE_KIND_LIST: {
        Nodes *children = input.as_list;
        if (children->count == 0) {
            *result = input;
            return 0;
        }

        Node head_node = children->items[0];
        if (head_node.kind == NODE_KIND_SYMBOL) {
            if (nob_sv_eq(head_node.as_symbol,
                          (Nob_String_View){.count = 2, .data = "if"})) {
                if (children->count != 3 + 1) {
                    printf("'if' expects three arguments!\n");
                    return 1;
                }

                Node condition = {0};
                int ret = eval_one(children->items[0 + 1], scopes, &condition);
                if (ret != 0) {
                    return ret;
                }

                if (node_is_true(condition)) {
                    ret = eval_one(children->items[1 + 1], scopes, result);
                } else {
                    ret = eval_one(children->items[2 + 1], scopes, result);
                }
                return ret;
            }

            if (nob_sv_eq(head_node.as_symbol,
                          (Nob_String_View){.count = 6, .data = "define"})) {
                if (children->count != 2 + 1) {
                    printf("'define' expects two arguments!\n");
                    return 1;
                }
                Node name_node = children->items[0 + 1];
                if (name_node.kind != NODE_KIND_SYMBOL) {
                    printf("First argument of 'define' must be a symbol!\n");
                    return 1;
                }

                Node value = {0};
                int ret = eval_one(children->items[1 + 1], scopes, &value);
                if (ret != 0) {
                    return ret;
                }

                nob_da_foreach(ScopeObject, scope_object, &global_scope) {
                    if (nob_sv_eq(scope_object->name, name_node.as_symbol)) {
                        scope_object->node = value;
                        *result = children->items[0 + 1];
                        return 0;
                    }
                }

                ScopeObject obj =
                    (ScopeObject){.name = name_node.as_symbol, .node = value};
                nob_da_append(&global_scope, obj);
                *result = children->items[0 + 1];
                return 0;
            }

            if (nob_sv_eq(head_node.as_symbol,
                          (Nob_String_View){.count = 3, .data = "let"})) {
                if (children->count != 2 + 1) {
                    printf("'let' expects two arguments!\n");
                    return 1;
                }

                Node definitions = children->items[0 + 1];
                if (definitions.kind != NODE_KIND_LIST) {
                    printf("First argument of 'let' must be a list!\n");
                    return 1;
                }

                Node expression = children->items[1 + 1];
                if (expression.kind != NODE_KIND_LIST) {
                    printf("Second argument of 'let' must be a list!\n");
                    return 1;
                }

                // it's fine to allocate this on the stack because we will throw
                // it off the scope stack after parsing 'let' anyways therefore
                // it'll outlive the recursive call
                Scope scope = {0};

                nob_da_foreach(Node, def, definitions.as_list) {
                    if (def->kind != NODE_KIND_LIST ||
                        def->as_list->count != 2 ||
                        def->as_list->items[0].kind != NODE_KIND_SYMBOL) {
                        printf("'let' expects a list of key-value pairs, e.g. "
                               "'(let ((x 42) (y 2)) (+ x y))'\n");
                        return 1;
                    }

                    Node key = def->as_list->items[0];
                    Node value_node = def->as_list->items[1];

                    Node value = {0};
                    int ret = eval_one(value_node, scopes, &value);
                    if (ret != 0)
                        return ret;

                    ScopeObject obj =
                        (ScopeObject){.name = key.as_symbol, .node = value};
                    nob_da_append(&scope, obj);
                }

                Node expression_value = {0};
                {
                    nob_da_append(scopes, &scope);

                    int ret = eval_one(expression, scopes, &expression_value);
                    if (ret != 0) {
                        return ret;
                    }

                    UNUSED(nob_da_pop(scopes));
                }

                *result = expression_value;
                return 0;
            }

            if (nob_sv_eq(head_node.as_symbol,
                          (Nob_String_View){.count = 6, .data = "lambda"})) {
                if (children->count != 2 + 1) {
                    printf("'lambda' expects two argument!\n");
                    return 1;
                }

                Node params = children->items[0 + 1];
                if (params.kind != NODE_KIND_LIST) {
                    printf("'lambda' expects a list of parameters as the first "
                           "argument!\n");
                    return 1;
                }

                nob_da_foreach(Node, node, params.as_list) {
                    if (node->kind != NODE_KIND_SYMBOL) {
                        printf("'lambda' expects a list of symbols "
                               "(parameters) as "
                               "the "
                               "first argument!\n");
                        return 1;
                    }
                }

                result->kind = NODE_KIND_LAMBDA;

                Nodes *params_list = malloc(sizeof(Nodes));
                if (params_list == NULL) {
                    printf("Failed to allocate memory, something is really "
                           "wrong!\n");
                    return 1;
                }
                memcpy(params_list, params.as_list, sizeof(Nodes));

                result->as_lambda.params = params_list;

                Node *body = malloc(sizeof(Node));
                if (body == NULL) {
                    printf("Failed to allocate memory, something is really "
                           "wrong!\n");
                    return 1;
                }
                memcpy(body, &children->items[1 + 1], sizeof(Node));

                result->as_lambda.body = body;
                return 0;
            }
        }

        Node head = {0};
        int ret = eval_one(head_node, scopes, &head);
        if (ret != 0) {
            return ret;
        }

        Nodes args = {0};
        for (size_t i = 1; i < children->count; i++) {
            Node n = {0};
            int ret = eval_one(children->items[i], scopes, &n);
            if (ret != 0) {
                return ret;
            }
            nob_da_append(&args, n);
        }

        switch (head.kind) {
        case NODE_KIND_FUNCTION: {
            head.as_function.apply(args, scopes, result);
            return 0;
        }
        case NODE_KIND_LAMBDA: {
            Nodes *lambda_params = head.as_lambda.params;
            if (lambda_params == NULL) {
                printf("Couldn't evaluate lambda, no params!\n");
                return 1;
            }

            if (lambda_params->count != args.count) {
                printf("Invalid amount of arguments to lambda: expected %zu, "
                       "found %zu\n!",
                       lambda_params->count, args.count);
                return 1;
            }

            Node *lambda_body = head.as_lambda.body;
            if (lambda_body == NULL) {
                printf("Couldn't evaluate lambda, no body!\n");
                return 1;
            }

            Scope scope = {0};
            // params->count == args.count
            for (size_t i = 0; i < lambda_params->count; i++) {
                Node param = lambda_params->items[i];
                // this should already be checked when we emit the lambda node
                NOB_ASSERT(param.kind == NODE_KIND_SYMBOL);
                Node value = args.items[i];

                ScopeObject obj = {.name = param.as_symbol, .node = value};
                nob_da_append(&scope, obj);
            }

            Node expression_value = {0};
            {
                nob_da_append(scopes, &scope);

                int ret = eval_one(*lambda_body, scopes, &expression_value);
                if (ret != 0) {
                    return ret;
                }

                UNUSED(nob_da_pop(scopes));
            }

            *result = expression_value;
            return 0;
        } break;
        default: {
            printf("Not a function-like object: '");
            node_print(head);
            printf("'\n");
            return 1;
        }
        }
    }
    case NODE_KIND_LAMBDA: {
        return 0;
    } break;
    default:
        return 1;
    }
}

int eval(Nodes inputs, Scopes *scopes, Nodes *results) {
    nob_da_foreach(Node, input, &inputs) {
        Node result = {0};
        int ret = eval_one(*input, scopes, &result);
        if (ret != 0) {
            return ret;
        }
        nob_da_append(results, result);
    }
    return 0;
}