#include <ctype.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define NOB_IMPLEMENTATION
#include "nob.h"

typedef enum e_TokenKind {
    TOKEN_KIND_PAR_LEFT,
    TOKEN_KIND_PAR_RIGHT,
    TOKEN_KIND_NUMBER,
    TOKEN_KIND_SYMBOL,
} TokenKind;

typedef struct s_Token {
    TokenKind kind;
    Nob_String_View data;
} Token;

typedef struct s_Tokens {
    Token *items;
    size_t count;
    size_t capacity;
} Tokens;

typedef enum e_NodeKind {
    NODE_KIND_SYMBOL,
    NODE_KIND_LIST,
    NODE_KIND_NUMBER,
    NODE_KIND_FUNCTION,
} NodeKind;

typedef struct s_Node Node;
typedef struct s_Nodes Nodes;

typedef Node *(*lisp_function_ptr)(Nodes *args);
typedef struct s_LispFunction {
    Nob_String_View name;
    lisp_function_ptr apply;
} LispFunction;

typedef struct s_Node {
    NodeKind kind;
    union {
        Nob_String_View as_symbol;
        Nodes *as_list;
        long long as_number;
        LispFunction as_function;
    };
} Node;

struct s_Nodes {
    Node *items;
    size_t count;
    size_t capacity;
};

Node *lisp_function_addition(Nodes *args) {
    if (args->count != 2) {
        printf("'+' expects two arguments!\n");
        return NULL;
    }

    if (args->items[0].kind != NODE_KIND_NUMBER ||
        args->items[1].kind != NODE_KIND_NUMBER) {
        printf("'+' can only add numbers, you specified something else!\n");
        return NULL;
    }

    Node *node = malloc(sizeof(Node));
    node->kind = NODE_KIND_NUMBER;

    intptr_t a = args->items[0].as_number;
    intptr_t b = args->items[1].as_number;
    node->as_number = a + b;
    return node;
}

static LispFunction globalScope[] = {
    (LispFunction){.name = (Nob_String_View){.items = "+", .count = 1},
                   .apply = lisp_function_addition}};

Node *lisp_function_lookup(Nob_String_View name) {
    for (size_t i = 0; i < NOB_ARRAY_LEN(globalScope); i++) {
        if (nob_sv_eq(globalScope[i].name, name)) {
            Node *node = malloc(sizeof(Node));
            node->kind = NODE_KIND_FUNCTION;
            node->as_function = globalScope[i];
            return node;
        }
    }

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

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 tokenize(Nob_String_View content, Tokens *tokens) {
    for (size_t i = 0; i < content.count; i++) {
        char c = content.data[i];

        int kind = -1;
        switch (c) {
        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};
        while (isdigit(c)) {
            nob_sb_appendf(&sb, "%c", c);
            c = content.data[++i];
        }

        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
        bool condition =
            isascii(c) && !isdigit(c) && !isblank(c) && c != '(' && c != ')';
        while (condition) {
            nob_sb_appendf(&sb, "%c", c);
            c = content.data[++i];
            condition = isascii(c) && !isblank(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);
        }
    }
}

Node *parse(Tokens *tokens) {
    if (tokens->count == 0) {
        printf("unexpected end of input");
        return NULL;
    }

    Token tok = tokens->items[0];
    tokens->items++;
    tokens->count--;

    Node *node = malloc(sizeof(Node));
    if (tok.kind == TOKEN_KIND_PAR_LEFT) {
        Nodes *nodes = malloc(sizeof(Nodes));
        memset(nodes, 0, sizeof(Nodes));

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

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

            Node *node = parse(tokens);
            if (node == NULL)
                return NULL;

            nob_da_append(nodes, *node);
        }

        node->kind = NODE_KIND_LIST;
        node->as_list = nodes;
        return node;
    }

    if (tok.kind == TOKEN_KIND_PAR_RIGHT) {
        printf("unexpected ')'");
        return NULL;
    }

    switch (tok.kind) {
    case TOKEN_KIND_NUMBER: {
        long long num = strtoll(tok.data.data, NULL, 10);
        // TODO:
        //      - error checking
        //      - double support

        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: {
        UNREACHABLE("");
    } break;
    }
    return node;
}

void node_print(Node node) {
    switch (node.kind) {
    case NODE_KIND_NUMBER: {
        printf("%lld", 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: {
        LispFunction func = node.as_function;
        printf(SV_Fmt, SV_Arg(func.name));
    } break;
    }
}

Node *eval(Node *node) {
    if (node == NULL)
        return NULL;

    switch (node->kind) {
    case NODE_KIND_NUMBER:
        return node;
    case NODE_KIND_SYMBOL: {
        return lisp_function_lookup(node->as_symbol);
    }
    case NODE_KIND_LIST: {
        Nodes *children = node->as_list;
        if (children->count == 0)
            return node;

        Node *head = &children->items[0];
        Node *func_node = eval(head);

        Nodes *args = malloc(sizeof(Nodes));
        for (size_t i = 1; i < children->count; i++) {
            nob_da_append(args, *eval(&children->items[i]));
        }

        if (func_node->kind == NODE_KIND_FUNCTION) {
            LispFunction func = func_node->as_function;
            return func.apply(args);
        }

        printf("Not a function: '");
        node_print(*func_node);
        printf("'\n");
        return NULL;
    }
    default:
        UNREACHABLE("");
    }
    return NULL;
}

int main(int argc, char **argv) {
    if (argc < 2)
        return 1;

    argv++;
    const char *line = *argv++;
    Nob_String_View view = nob_sv_from_cstr(line);

    Tokens tokens = {0};
    tokenize(view, &tokens);

#if 0
    for (size_t i = 0; i < tokens.count; i++) {
        Token tok = tokens.items[i];
        printf("Token(kind=%s, data=" SV_Fmt ")\n", token_kind_name(tok.kind),
               SV_Arg(tok.data));
    }
#endif

    Node *result = parse(&tokens);
    if (result == NULL)
        return 1;

    result = eval(result);
    node_print(*result);
    printf("\n");

    return 0;
}