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

#include "lexer.h"
#include "value.h"
#include "env.h"

#ifdef ENABLE_MEMDEBUG
#define MEMDEBUG_OUT_OF_BOUNDS
#define MEMDEBUG_IMPLEMENTATION
#define MEMDEBUG_OUTPUT_DIR "files"
#endif

#include "common.h"
#include "builtin.h"

#define LEN(arr) (sizeof(arr)/sizeof(*(arr)))

// TODO:
// - Think about memory leakage and on non fatal errors
//   like failed list creation and faild s-exp parsing
// - Check for functions not wrapped in ERR_* macro
// - Add more and better error messages

#define TOKEN(tctx) (tctx)->token

#define TOKEN_NEXT(tctx)                                 \
    ERR_Z(next_token(tctx), die("Can't get next token"))

#define TOKEN_ASSERT(token, ttype, fail)          \
    if(token->type != ttype) {                    \
        err("Expected token '%s', not '%s'",      \
            token_type_string[ttype],             \
            token_type_string[token->type]);      \
        fail;                                     \
    }

#define TOKEN_SKIP(tctx, ttype, fail) do {        \
        TOKEN_ASSERT(TOKEN(tctx), ttype, fail)    \
        else {                                    \
            TOKEN_NEXT(tctx);                     \
        }                                         \
    } while(0)

#define TOKEN_MATCH(tctx, ttype, fail) do {       \
        TOKEN_NEXT(tctx);                         \
        TOKEN_ASSERT(TOKEN(tctx), ttype, fail);   \
    } while(0)


#define NOT_IMPLEMENTED() die("Not Implemented. ABORTING")

static void print_token(struct token *token)
{
    char buf[256] = {0};
    ERR_Z(token_value_string(token, LEN(buf), buf), return);
    info("%-12s %s", token_type_string[token->type], buf);
}

static void print_value(value_t value)
{
    char buf[256] = {0};
    value_string(value, LEN(buf), buf);
    info("%-12s %s", value ?
         value_type_string[value->type] : "VALUE", buf);
}

struct tctx {
    enum tctx_type {
        TCTX_LEXER,
        TCTX_TOKLIST
    } type;

    struct token *token;

    union {
        struct lex_ctx {
            lexer_t lexer;
            struct token token;
        } lex_ctx;

        struct tok_ctx {
            struct token *list;
            size_t len;
            size_t idx;
        } tok_ctx;
    } value;
};

static void tctx_init_lexer(struct tctx *tctx, lexer_t lexer)
{
    tctx->type = TCTX_LEXER;
    tctx->value.lex_ctx.lexer = lexer;
    tctx->token = &tctx->value.lex_ctx.token;
}

static void tctx_init_toklist(struct tctx *tctx, struct token *toklist, size_t len)
{
    tctx->type = TCTX_TOKLIST;
    tctx->value.tok_ctx.list = toklist;
    tctx->value.tok_ctx.len = len;
    tctx->value.tok_ctx.idx = 0;
    tctx->token = NULL;
}

static struct token *next_token(struct tctx *tctx)
{
    switch(tctx->type) {
    case TCTX_LEXER: ;
        struct lex_ctx *l = &tctx->value.lex_ctx;
        ERR_NZ(lexer_token_next(l->lexer, &l->token), _r, goto fail);
        // tctx->token should already point to l->token
        break;
    case TCTX_TOKLIST: ;
        struct tok_ctx *t = &tctx->value.tok_ctx;
        if(t->idx == t->len) goto fail;
        tctx->token = &t->list[t->idx++];
        break;
    }

    return tctx->token;
fail:
    tctx->token = NULL;
    return NULL;
}


value_t apply(value_t proc, size_t argc, value_t *argv);

value_t evaluate_expr(env_t env, struct tctx *tctx);
value_t evaluate_sexp(env_t env, struct tctx *tctx);

value_t evaluate_id    (env_t env, struct tctx *tctx);
value_t evaluate_lambda(env_t env, struct tctx *tctx);
value_t evaluate_define(env_t env, struct tctx *tctx);
value_t evaluate_if    (env_t env, struct tctx *tctx);

value_t quote_expr(env_t env, struct tctx *tctx);
value_t quote_sexp(env_t env, struct tctx *tctx);

size_t toklist_expr(struct tctx *tctx, struct token **toklist);

static env_t global_env = ENV_EMPTY;
static value_t global_nil = VALUE_EMPTY;

static void destroy_env(char *key, value_t value)
{
    (void)key;
    value_destroy(value);
}

static void destroy_global_env(char *key, value_t value)
{
    free(key);
    value_destroy(value);
}

int main(int argc, char **argv)
{
    int opt, repl_flag = 0;

#ifdef DEBUG
    char *filename = "files/test-lambda.l";
#else
    char *filename = NULL;
#endif

    while((opt = getopt(argc, argv, "rhf:")) != -1) {
        switch(opt) {
        case 'h':
            printf("help????");
            return 0;
        case 'r': repl_flag = 1; break;
        case 'f': filename = optarg; break;
        default:
            err("No such option %c", opt);
            return 1;
        }
    }

    if(optind < argc) {
        filename = argv[optind];
    }

    if(!filename) repl_flag = 1;


    env_t builtin_env = env_create(ENV_EMPTY, destroy_env);
    global_env = env_create(builtin_env, destroy_global_env);
    global_nil = value_create(VALUE_NIL, NULL);

    for(size_t i = 0; i < BUILTIN_PROCEDURES; i++) {
        value_t proc_value = value_create(
            VALUE_PROC_BUILTIN,
            (void *)&builtin_proc_descriptions[i]);
        hashtable_insert(builtin_env->table,
                         (void *)builtin_proc_name_list[i],
                         (void *)proc_value, NULL, NULL);
    }

    FILE *fp = stdin;
    if(!repl_flag) {
        fp = fopen(filename, "r");
        if(!fp) {
            die("fopen: %s", strerror(errno));
        }
    }

    lexer_t lexer = lexer_create(fp);

    struct tctx tctx = {0};
    tctx_init_lexer(&tctx, lexer);

    if(repl_flag) printf("> ");
    while(next_token(&tctx))
    {
        value_t val = evaluate_expr(ENV_EMPTY, &tctx);

        if(val == VALUE_EMPTY) {
            if(!repl_flag) {
                printf("%s:%zu: FAILED\n", filename, lexer->line);
                break;
            } else {
                printf("=> FAILED\n");
                lexer_clear_line(lexer);
            }
        } else {
            char buf[256] = {0};
            value_string(val, sizeof(buf), buf);

            if(!repl_flag) {
                printf("%s:%zu: %s\n", filename, lexer->line, buf);
            } else {
                printf("=> %s\n", buf);
            }
        }

        if(repl_flag) printf("> ");
        value_destroy(val);
    }

    lexer_destroy(lexer);
    fclose(fp);

    value_destroy(global_nil);
    env_destroy(global_env);
    return 0;
}

#define HAS_ENOUGH_ARGS(proc, type, argc, fail)                     \
    if(argc != proc->value.type.argc) {                             \
        err("Wrong number of arguemnts, expected %zu, but got %zu", \
            proc->value.type.argc, argc);                           \
        fail;                                                       \
    }

static value_t apply_lambda(struct proc *proc, value_t *args)
{
    value_t ret = VALUE_EMPTY;
    env_t env = ENV_EMPTY;
    struct tctx tctx = {0};

    ERR_Z(env = env_create(env_copy(proc->parent_env), destroy_env), goto exit);
    tctx_init_toklist(&tctx, proc->body, proc->body_len);

    for(size_t i = 0; i < proc->argc; i++) {
        ERR_NZ(hashtable_insert(env->table,
                                (void*)proc->arg_keys[i]->value.atom,
                                (void*)value_copy(args[i]), NULL, NULL), _r, goto exit);
    }

    TOKEN_NEXT(&tctx);
    ret = evaluate_expr(env, &tctx);

exit:
    env_destroy(env);
    return ret;
}

value_t apply(value_t proc, size_t argc, value_t *argv)
{
    if(proc == VALUE_EMPTY) return value_copy(global_nil);

    switch(proc->type) {
    case VALUE_PROC:
        HAS_ENOUGH_ARGS(proc, proc, argc, return VALUE_EMPTY);
        return apply_lambda(&proc->value.proc, argv);
    case VALUE_PROC_BUILTIN:
        HAS_ENOUGH_ARGS(proc, proc_builtin, argc, return VALUE_EMPTY);
        return proc->value.proc_builtin.proc(argv);
    default:
        err("Value is not a procedure");
        return VALUE_EMPTY;
    }
}

value_t evaluate_expr(env_t env, struct tctx *tctx)
{
    value_t ret = VALUE_EMPTY;

    switch(TOKEN(tctx)->type) {
    case TOKEN_LP:
        ERR_Z(ret = evaluate_sexp(env, tctx), goto exit);
        break;
    case TOKEN_ID:
        ERR_Z(ret = evaluate_id(env, tctx), goto exit);
        break;
    case TOKEN_STR:
    case TOKEN_INT:
        ERR_Z(ret = value_from_token(TOKEN(tctx)), goto exit);
        break;
    case TOKEN_QUOTE:
        TOKEN_NEXT(tctx);
        ERR_Z(ret = quote_expr(env, tctx), goto exit);
        break;
    default:
        err("Did not exptect token '%s'", token_type_string[TOKEN(tctx)->type]);
        break;
    }

exit:
    return ret;
}

#define IS_NOT_FIRST(arg, fail)                                \
    if(arg != 0) {                                             \
        err("special forms can only be the first argument");   \
        fail;                                                  \
    }
#define SPECIAL_FORM(ret, argc, fn, fail) do {  \
        IS_NOT_FIRST(argc, fail);               \
        ERR_Z(ret = fn, fail);                  \
    } while(0)

value_t evaluate_sexp(env_t env, struct tctx *tctx)
{
    value_t ret = VALUE_EMPTY;
    value_t body[256] = {VALUE_EMPTY};
    size_t argc = 0;

    TOKEN_SKIP(tctx, TOKEN_LP, goto exit);

    for(argc = 0; TOKEN(tctx)->type != TOKEN_RP; argc++)
    {
        if(argc >= LEN(body)) {
            err("Too many arguments");
            goto exit;
        }

        switch(TOKEN(tctx)->type) {
        case TOKEN_LAMBDA:
            SPECIAL_FORM(ret, argc, evaluate_lambda(env, tctx), goto exit);
            goto exit;
        case TOKEN_DEFINE:
            SPECIAL_FORM(ret, argc, evaluate_define(env, tctx), goto exit);
            goto exit;
        case TOKEN_QUOTE_FORM:
            TOKEN_NEXT(tctx);
            SPECIAL_FORM(ret, argc, quote_expr(env, tctx), goto exit);
            TOKEN_MATCH(tctx, TOKEN_RP,
                        value_destroy(ret);
                        ret = VALUE_EMPTY;
                        goto exit);
            goto exit;
        case TOKEN_IF:
            SPECIAL_FORM(ret, argc, evaluate_if(env, tctx), goto exit);
            goto exit;
        default:
            ERR_Z(body[argc] = evaluate_expr(env, tctx), goto exit);
            break;
        }

        TOKEN_NEXT(tctx);
    }

    ret = apply(body[0], argc-1, &body[1]);

exit:
    for(size_t i = 0; i < argc; i++)
        value_destroy(body[i]);
    return ret;
}

value_t evaluate_id(env_t env, struct tctx *tctx)
{
    if(env == ENV_EMPTY) {
        return evaluate_id(global_env, tctx);
    }

    value_t ret = VALUE_EMPTY;
    ERR_NZ(hashtable_query(env->table, (void *)TOKEN(tctx)->value.id, (void **)&ret), _r, goto fail);

    return value_copy(ret);

fail:
    if(env == global_env->parent) {
        err("Symbol %s is unbound", TOKEN(tctx)->value.id);
        return VALUE_EMPTY;
    }
    return evaluate_id(env->parent, tctx);
}

value_t evaluate_lambda(env_t env, struct tctx *tctx)
{
    value_t ret = VALUE_EMPTY;

    value_t args[32] = {VALUE_EMPTY};
    size_t argc = 0;
    value_t *arg_keys = NULL;

    struct token *body = NULL;
    size_t body_len = 0;

    TOKEN_SKIP(tctx, TOKEN_LAMBDA, goto fail);
    TOKEN_SKIP(tctx, TOKEN_LP, goto fail);

    while(TOKEN(tctx)->type != TOKEN_RP)
    {
        if(argc >= LEN(args)) {
            err("Too many arguments");
            goto fail;
        }

        if(TOKEN(tctx)->type != TOKEN_ID) {
            err("Token '%s' not expected in lambda args", token_type_string[TOKEN(tctx)->type]);
            goto fail;
        }

        ERR_Z(args[argc++] = value_from_token(TOKEN(tctx)), goto fail);

        TOKEN_NEXT(tctx);
    }

    ERR_Z(body_len = toklist_expr(tctx, &body), goto fail);

    TOKEN_MATCH(tctx, TOKEN_RP, goto fail);

    arg_keys = calloc(argc, sizeof(*arg_keys));
    memcpy(arg_keys, args, argc * sizeof(*arg_keys));

    struct proc proc = {env_copy(env), arg_keys, argc, body, body_len};
    ERR_Z(ret = value_create(VALUE_PROC, &proc),
          env_destroy(env); // remove the copy
          goto fail);

    return ret;

fail:
    err("Procedure creation failed");
    for(size_t i = 0; i < argc; i++) value_destroy(args[i]);
    if(body) {
        for(size_t i = 0; i < body_len; i++) token_dealloc(&body[i]);
        free(body);
    }
    if(arg_keys) free(arg_keys);
    return VALUE_EMPTY;
}

#define STR_ALLOC_COPY(dest, str) do {    \
        size_t len = strlen(str) + 1;     \
        dest = malloc(len);               \
        memcpy((dest), (str), len);       \
    } while(0)

value_t evaluate_define(env_t env, struct tctx *tctx)
{
    // TODO: don't alloc when the key is the same
    char *key = NULL;

    // only in the outside environement
    if(env != ENV_EMPTY) {
        err("define can only be called in the outermost environement");
        goto fail;
    }

    TOKEN_SKIP(tctx, TOKEN_DEFINE, goto fail);

    switch(TOKEN(tctx)->type)
    {
    case TOKEN_ID:
        STR_ALLOC_COPY(key, TOKEN(tctx)->value.id);
        break;
    default:
        err("Did not exptect token '%s'", token_type_string[TOKEN(tctx)->type]);
        goto fail;
    }

    TOKEN_NEXT(tctx);

    value_t val = VALUE_EMPTY;
    ERR_Z(val = evaluate_expr(env, tctx), goto fail);

    TOKEN_MATCH(tctx, TOKEN_RP,
                value_destroy(val);
                goto fail);

    value_t prevval = VALUE_EMPTY;
    char *prevkey = NULL;

    ERR_NZ(
        hashtable_insert(global_env->table, (void *)key, (void *)val,
                         (void**)&prevkey, (void **)&prevval),
        r, {
            err("Couldn't insert symbol into the hashtable due to %s", strerror(r));
            value_destroy(val); // the copy
            goto fail;
        });

    if(prevkey) free(prevkey);
    value_destroy(prevval);
    return value_copy(global_nil);

fail:
    if(key) free(key);
    return VALUE_EMPTY;
}

value_t evaluate_if(env_t env, struct tctx *tctx)
{
    value_t ret = VALUE_EMPTY;
    value_t cond = VALUE_EMPTY;

    TOKEN_SKIP(tctx, TOKEN_IF, goto exit);

    ERR_Z(cond = evaluate_expr(env, tctx), goto exit);

    // if(cond->type == VALUE_NIL) {
    //     toklist_expr(tctx, NULL); // skip one expression
    //     TOKEN_NEXT(tctx);
    //     TOKEN_MATCH(tctx, TOKEN_RP, goto exit);
    //     goto exit;
    // }

    ERR_Z(cond->type == VALUE_INT,
          err("expected condition to evaluate to VALUE_INT");
          goto exit);

    if(cond->value.num) {
        TOKEN_NEXT(tctx);
        ret = evaluate_expr(env, tctx);

        toklist_expr(tctx, NULL); // skip one expression
    } else {
        toklist_expr(tctx, NULL); // skip one expression

        TOKEN_NEXT(tctx);
        ret = evaluate_expr(env, tctx);
    }

    TOKEN_MATCH(tctx, TOKEN_RP, goto exit);

exit:
    value_destroy(cond);
    return ret;
}

value_t quote_expr(env_t env, struct tctx *tctx)
{
    value_t ret = VALUE_EMPTY;

    switch(TOKEN(tctx)->type) {
    case TOKEN_ID:
    case TOKEN_STR:
    case TOKEN_INT:
        ERR_Z(ret = value_from_token(TOKEN(tctx)), goto exit);
        break;
    case TOKEN_QUOTE:
    case TOKEN_LAMBDA:
    case TOKEN_DEFINE:
    case TOKEN_QUOTE_FORM:
    case TOKEN_IF: ;
        char name[64] = {0};
        ERR_Z(token_value_string(TOKEN(tctx), LEN(name), name), goto exit);
        struct token temp = {0}; temp.type = TOKEN_ID; temp.value.id = name;
        ERR_Z(ret = value_from_token(&temp), goto exit);
        break;
    case TOKEN_LP:
        ERR_Z(ret = quote_sexp(env, tctx), goto exit);
        break;
    case TOKEN_UNQUOTE:
        TOKEN_SKIP(tctx, TOKEN_UNQUOTE, goto exit);
        ERR_Z(ret = evaluate_expr(env, tctx), goto exit);
        break;
    default:
        err("Did not exptect token '%s'", token_type_string[TOKEN(tctx)->type]);
        break;
    }
exit:
    return ret;
}

value_t quote_sexp(env_t env, struct tctx *tctx)
{
    value_t ret   = VALUE_EMPTY;
    value_t left  = VALUE_EMPTY;
    value_t right = VALUE_EMPTY;

    TOKEN_SKIP(tctx, TOKEN_LP, goto exit);

    if(TOKEN(tctx)->type == TOKEN_RP) {
        ret = value_copy(global_nil);
        goto exit;
    }

    ERR_Z(left = quote_expr(env, tctx), goto exit);

    TOKEN_NEXT(tctx);
    if(TOKEN(tctx)->type == TOKEN_DOT)
    {
        // Parse cons
        TOKEN_NEXT(tctx);
        ERR_Z(right = quote_expr(env, tctx), return VALUE_EMPTY);
        TOKEN_MATCH(tctx, TOKEN_RP, return VALUE_EMPTY);

        struct cons cons = {value_copy(left), value_copy(right)};
        ret = value_create(VALUE_CONS, &cons);
        goto exit;
    }

    // Parse list
    right = value_copy(global_nil);
    value_t *rightmost = &right; // the final nil
    while(TOKEN(tctx)->type != TOKEN_RP)
    {
        value_t new = VALUE_EMPTY;

        if(TOKEN(tctx)->type == TOKEN_DOT) {
            TOKEN_NEXT(tctx);

            ERR_Z(new = quote_expr(env, tctx), goto exit);

            TOKEN_MATCH(tctx, TOKEN_RP,
                        value_destroy(new);
                        goto exit);

            value_destroy(*rightmost); // destroy the copied nil
            *rightmost = new;
            break;
        }

        ERR_Z(new = quote_expr(env, tctx),
              goto exit);

        value_t new_cons = VALUE_EMPTY;
        struct cons cons = {new, *rightmost};
        ERR_Z(new_cons = value_create(VALUE_CONS, &cons),
              value_destroy(new);
              goto exit);

        *rightmost = new_cons;
        rightmost = &new_cons->value.cons.right;

        TOKEN_NEXT(tctx);
    }

    struct cons cons = {value_copy(left), value_copy(right)};
    ret = value_create(VALUE_CONS, &cons);
exit:
    value_destroy(left);
    value_destroy(right);
    return ret;
}

#define IS_OP(token) (token->type == TOKEN_QUOTE)

size_t toklist_expr(struct tctx *tctx, struct token **toklist)
{
    struct token tokens[256];
    size_t tokens_len = 0;

    size_t depth = 0;
    do {
        TOKEN_NEXT(tctx);

        if(tokens_len >= LEN(tokens)) {
            err("Too many tokens in expr");
            goto fail;
        }

        if(TOKEN(tctx)->type == TOKEN_LP) depth++;
        else if(TOKEN(tctx)->type == TOKEN_RP) depth--;

        token_clone(&tokens[tokens_len++], TOKEN(tctx));

    } while(depth > 0 || IS_OP(TOKEN(tctx)));

    if(!toklist) goto fail;

    *toklist = calloc(tokens_len, sizeof(*tokens));
    memcpy(*toklist, tokens, tokens_len * sizeof(*tokens));

    return tokens_len;

fail:
    for(size_t i = 0; i < tokens_len; i++) token_dealloc(&tokens[i]);
    return 0;
}