#include "function.h"

#include "hashtable.h"
#include "lisp.h"
#include "list.h"
#include "read.h"
#include "stack.h"

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

const char *llps_strerror(LambdaListParseStatus status) {
    static const char *MSGS[LLPS_N_ERROS] = {
        [LLPS_OK] = "No error",
        [LLPS_DOTTED] = "Dotted list",
        [LLPS_REPEAT_SECTION] = "Repeated section",
        [LLPS_REPEAT_NAME] = "Repeated name",
        [LLPS_ORDER] = "Section out of order",
        [LLPS_BAD_NAME] = "Invalid variable name",
        [LLPS_REPEAT_REST] = "Too many rest variables",
        [LLPS_AFTER_ALLOW_OTHER_KEYS] = "Variable after &allow-other-keys",
        [LLPS_INVALID_OPT_SPEC] = "Invalid optional spec",
    };
    return MSGS[status];
}

DEFINE_SYMBOL(and_optional, "&optional");
DEFINE_SYMBOL(and_rest, "&rest");
DEFINE_SYMBOL(and_key, "&key");
DEFINE_SYMBOL(and_allow_other_keys, "&allow-other-keys");

static bool is_valid_variable_name(LispVal *val) {
    return SYMBOLP(val) && !NILP(val) && val != Qt && val != Qunbound;
}

static LispVal *intern_as_keyword(LispVal *name) {
    assert(SYMBOLP(name));
    LispString *name_str = ((LispSymbol *) name)->name;
    char *kw_name = lisp_malloc(name_str->length + 2);
    kw_name[0] = ':';
    memcpy(kw_name + 1, name_str->data, name_str->length);
    kw_name[name_str->length + 1] = '\0';
    return Fintern(
        make_lisp_string(kw_name, name_str->length + 1, true, false));
}

// on error, put the object that caused the problem in entry
static LambdaListParseStatus
parse_optional_arg_spec(LispVal *used_names, LispVal **out, LispVal *entry) {
    // single symbol
    if (SYMBOLP(entry)) {
        if (!is_valid_variable_name(entry)) {
            *out = entry;
            return LLPS_BAD_NAME;
        } else if (!NILP(Fgethash(used_names, entry, Qnil))) {
            *out = entry;
            return LLPS_REPEAT_NAME;
        }
        *out = LIST(entry, Qnil, Qnil);
        Fputhash(used_names, entry, Qt);
        return LLPS_OK;
    } else if (!CONSP(entry)) {
        *out = entry;
        return LLPS_BAD_NAME;
    }
    // list
    LispVal *name = XCAR(entry);
    if (!is_valid_variable_name(name)) {
        *out = name;
        return LLPS_BAD_NAME;
    } else if (!NILP(Fgethash(used_names, name, Qnil))) {
        *out = name;
        return LLPS_REPEAT_NAME;
    }
    Fputhash(used_names, name, Qt);
    if (list_length_eq(entry, 1)) {
        *out = LIST(XCAR(entry), Qnil, Qnil);
        return LLPS_OK;
    } else if (list_length_eq(entry, 2)) {
        *out = LIST(name, XCAR(XCDR(entry)), Qnil);
        return LLPS_OK;
    } else if (list_length_eq(entry, 3)) {
        LispVal *pvar = XCAR(XCDR(XCDR(entry)));
        if (!is_valid_variable_name(pvar)) {
            *out = pvar;
            return LLPS_BAD_NAME;
        } else if (!NILP(Fgethash(used_names, pvar, Qnil))) {
            *out = pvar;
            return LLPS_REPEAT_NAME;
        }
        Fputhash(used_names, pvar, Qt);
        *out = LIST(XCAR(entry), XCAR(XCDR(entry)), pvar);
        return LLPS_OK;
    } else {
        *out = entry;
        return LLPS_INVALID_OPT_SPEC;
    }
}

#define RETURN_ERROR(err, obj)                \
    {                                         \
        release_hash_table_no_gc(used_names); \
        result->status = err;                 \
        result->err_obj = (obj);              \
        return;                               \
    }
void parse_lambda_list(LambdaListParseResult *result, LispVal *list) {
    enum { REQ = 0, OPT = 1, KEY = 2, REST = 4, MUST_CHANGE } mode = REQ;
    unsigned int seen = 0;
    result->err_obj = Qnil;
    result->status = LLPS_OK;
    struct LambdaList *out = &result->lambda_list;
    LispVal *used_names = make_hash_table_no_gc(Qnil, Qnil);
    // TODO check for repeat names
    out->n_req = 0;
    out->n_opt = 0;
    out->allow_other_keys = false;
    out->req = Qnil;
    out->opt = Qnil;
    out->kw = Qnil;
    out->rest = Qnil;
    size_t cur_idx = 0; // for keyword args
    DOTAILS(tail, list) {
        if (!LISTP(tail)) {
            RETURN_ERROR(LLPS_DOTTED, list);
        } else if (out->allow_other_keys) {
            RETURN_ERROR(LLPS_AFTER_ALLOW_OTHER_KEYS, XCAR(tail));
        }
        LispVal *cur = XCAR(tail);
        if (cur == Qand_allow_other_keys) {
            if (out->allow_other_keys) {
                RETURN_ERROR(LLPS_REPEAT_SECTION, cur);
            } else if (!(seen & KEY)) {
                RETURN_ERROR(LLPS_ORDER, cur);
            }
            out->allow_other_keys = true;
        } else if (cur == Qand_rest) {
            if (seen & REST) {
                RETURN_ERROR(LLPS_REPEAT_SECTION, cur)
            } else if (seen & KEY) {
                RETURN_ERROR(LLPS_ORDER, cur);
            }
            seen |= REST;
            mode = REST;
        } else if (cur == Qand_optional) {
            if (seen & OPT) {
                RETURN_ERROR(LLPS_REPEAT_SECTION, cur)
            } else if (seen & KEY || seen & REST) {
                RETURN_ERROR(LLPS_ORDER, cur);
            }
            seen |= OPT;
            mode = OPT;
        } else if (cur == Qand_key) {
            if (seen & KEY) {
                RETURN_ERROR(LLPS_REPEAT_SECTION, list)
            }
            seen |= KEY;
            mode = KEY;
            out->kw = Fmake_hash_table(Qnil, Qnil);
        } else if (mode == REST) {
            if (!NILP(out->rest)) {
                RETURN_ERROR(LLPS_REPEAT_REST, cur);
            } else if (!is_valid_variable_name(cur)) {
                RETURN_ERROR(LLPS_BAD_NAME, cur)
            } else if (!NILP(Fgethash(used_names, cur, Qnil))) {
                RETURN_ERROR(LLPS_REPEAT_NAME, cur);
            }
            Fputhash(used_names, cur, Qt);
            out->rest = cur;
            ++cur_idx;
        } else if (mode == OPT || mode == KEY) {
            LispVal *entry;
            LambdaListParseStatus status =
                parse_optional_arg_spec(used_names, &entry, cur);
            if (status != LLPS_OK) {
                RETURN_ERROR(status, entry)
            }
            if (mode == OPT) {
                out->opt = CONS(entry, out->opt);
                ++out->n_opt;
            } else {
                Fputhash(out->kw, intern_as_keyword(XCAR(entry)),
                         CONS(MAKE_FIXNUM(cur_idx), entry));
            }
            ++cur_idx;
        } else if (!is_valid_variable_name(cur)) {
            RETURN_ERROR(LLPS_BAD_NAME, cur);
        } else if (!NILP(Fgethash(used_names, cur, Qnil))) {
            RETURN_ERROR(LLPS_REPEAT_NAME, cur);
        } else {
            Fputhash(used_names, cur, Qt);
            out->req = CONS(cur, out->req);
            ++out->n_req;
            ++cur_idx;
        }
    }
    out->req = Fnreverse(out->req);
    out->opt = Fnreverse(out->opt);
    release_hash_table_no_gc(used_names);
}
#undef RETURN_ERROR

LispVal *make_builtin_function(LispVal *name, LispVal *(*cfunc)(void),
                               const char *lisp_args, size_t args_len,
                               LispVal *docstr) {
    LispFunction *obj = lisp_alloc_object(sizeof(LispFunction), TYPE_FUNCTION);
    obj->name = name;
    obj->type = FUNCTION_NATIVE;
    obj->docstr = docstr;
    obj->impl.native.no_eval_args = false;
    obj->impl.native.addr.zero = cfunc;
    ReadStream stream;
    read_stream_init(&stream, lisp_args, args_len);
    LispVal *args_form = read(&stream);
    if (!args_form) {
        fprintf(stderr, "Builtin function lambda list had a syntax error\n");
        fprintf(stderr, "Name: ");
        debug_print(stderr, name);
        fprintf(stderr, "\nLambda list: \"%s\"\n", lisp_args);
        exit(1);
    }
    LambdaListParseResult result;
    parse_lambda_list(&result, args_form);
    if (result.status != LLPS_OK) {
        fprintf(stderr, "Error parsing builtin lambda list: %s\n",
                llps_strerror(result.status));
        fprintf(stderr, "Name: ");
        debug_print(stderr, name);
        fprintf(stderr, "\nLambda list: \"%s\"\n", lisp_args);
        exit(1);
    }
    obj->args = result.lambda_list;
    return obj;
}

// Calling functions
static ALWAYS_INLINE LispVal *evaluate_function_arguments(LispVal *args) {
    LispVal *start = Qnil;
    LispVal *end;
    DOLIST(arg, args) {
        if (NILP(start)) {
            start = CONS(Feval(arg, Vlexical_environment), Qnil);
            end = start;
        } else {
            RPLACD(end, CONS(Feval(arg, Vlexical_environment), Qnil));
            end = XCDR(end);
        }
    }
    return start;
}

enum ProcessArgsResult {
    PROCESS_ARGS_OK,
    PROCESS_ARGS_TOO_FEW,
    PROCESS_ARGS_TOO_MANY,
    PROCESS_ARGS_NO_KEY_VALUE,
    PROCESS_ARGS_BAD_KEY,
    PROCESS_ARGS_N_ERRORS,
};

static const char *process_args_strerror(enum ProcessArgsResult status) {
    static const char *MSGS[PROCESS_ARGS_N_ERRORS] = {
        [PROCESS_ARGS_OK] = "No error",
        [PROCESS_ARGS_TOO_FEW] = "Not enough arguments",
        [PROCESS_ARGS_TOO_MANY] = "Too many arguments",
        [PROCESS_ARGS_NO_KEY_VALUE] = "Key without a value",
        [PROCESS_ARGS_BAD_KEY] = "Unknown key",
    };
    return MSGS[status];
}

static ALWAYS_INLINE size_t NATIVE_FUNCTION_TOTAL_ARG_COUNT(LispVal *val) {
    assert(FUNCTIONP(val));
    LispFunction *fobj = val;
    return fobj->args.n_req + fobj->args.n_opt + !NILP(fobj->args.rest)
           + (NILP(fobj->args.kw) ? 0 : HASH_TABLE_COUNT(fobj->args.kw));
}

static ALWAYS_INLINE enum ProcessArgsResult
process_complex_native_args(LispFunction *fobj, LispVal *args,
                            LispVal *restrict out[MAX_NATIVE_FUNCTION_ARGS],
                            intptr_t *rest_idx) {
    size_t rem_req = fobj->args.n_req;
    size_t rem_opt = fobj->args.n_opt;
    size_t idx = 0;
    while (rem_req--) {
        if (NILP(args)) {
            return PROCESS_ARGS_TOO_FEW;
        }
        out[idx++] = XCAR(args);
        args = XCDR(args);
    }
    while (rem_opt--) {
        if (NILP(args)) {
            return PROCESS_ARGS_OK;
        }
        out[idx++] = XCAR(args);
        args = XCDR(args);
    }
    if (!NILP(args) && (NILP(fobj->args.kw)) && NILP(fobj->args.rest)) {
        return PROCESS_ARGS_TOO_MANY;
    }
    if (!NILP(fobj->args.rest)) {
        *rest_idx = idx;
        out[idx++] = args;
    } else {
        *rest_idx = -1;
    }
    if (NILP(fobj->args.kw)) { // we are not a keyword function
        return PROCESS_ARGS_OK;
    }
    while (!NILP(args)) {
        if (NILP(XCDR(args))) {
            return PROCESS_ARGS_NO_KEY_VALUE;
        }
        LispVal *entry = Fgethash(fobj->args.kw, XCAR(args), Qnil);
        if (!NILP(entry)) {
            fixnum_t idx = XFIXNUM(XCAR(entry));
            if (!out[idx]) {
                out[idx] = XCAR(XCDR(args));
            }
        } else if (!fobj->args.allow_other_keys) {
            return PROCESS_ARGS_BAD_KEY;
        }
        args = XCDR(XCDR(args));
    }
    return PROCESS_ARGS_OK;
}

static ALWAYS_INLINE LispVal *call_native(LispVal *orig_func,
                                          LispFunction *fobj, LispVal *args) {
    StackFrame *stack_ref = LISP_STACK_REF();
    push_call_frame(orig_func, fobj, args);
    if (!fobj->impl.native.no_eval_args) {
        args = evaluate_function_arguments(args);
    }
    set_stack_evaluated_args(LISP_STACK_REF(), args);
    LispVal *arg_arr[MAX_NATIVE_FUNCTION_ARGS] = {NULL};
    size_t count = NATIVE_FUNCTION_TOTAL_ARG_COUNT(fobj);
    intptr_t rest_idx;
    enum ProcessArgsResult res =
        process_complex_native_args(fobj, args, arg_arr, &rest_idx);
    if (res != PROCESS_ARGS_OK) {
        // TODO better errors
        printf("Bad arguments to builtin \"");
        debug_print(stdout, orig_func);
        printf("\": %s\n", process_args_strerror(res));
        abort();
    }
    for (intptr_t i = 0; i < count; ++i) {
        if (!arg_arr[i]) {
            arg_arr[i] = Qnil;
        }
        add_local_reference(LISP_STACK_REF(), arg_arr[i]);
    }
    LispVal *retval;
    switch (count) {
    case 0:
        retval = fobj->impl.native.addr.zero();
        break;
    case 1:
        retval = fobj->impl.native.addr.one(arg_arr[0]);
        break;
    case 2:
        retval = fobj->impl.native.addr.two(arg_arr[0], arg_arr[1]);
        break;
    case 3:
        retval =
            fobj->impl.native.addr.three(arg_arr[0], arg_arr[1], arg_arr[2]);
        break;
    case 4:
        retval = fobj->impl.native.addr.four(arg_arr[0], arg_arr[1], arg_arr[2],
                                             arg_arr[3]);
        break;
    case 5:
        retval = fobj->impl.native.addr.five(arg_arr[0], arg_arr[1], arg_arr[2],
                                             arg_arr[3], arg_arr[4]);
        break;
    default:
        abort();
    }
    the_stack.nogc_retval = retval;
    unwind_to(stack_ref);
    add_local_reference(LISP_STACK_REF(), the_stack.nogc_retval);
    return retval;
}

static ALWAYS_INLINE void push_optional_argument_to_lexenv(LispVal *spec,
                                                           LispVal *value) {
    new_lexical_variable(XCAR(spec), value);
    if (!NILP(THIRD(spec))) {
        new_lexical_variable(THIRD(spec), Qt);
    }
}

static ALWAYS_INLINE void
push_missing_optional_argument_to_lexenv(LispVal *spec) {
    new_lexical_variable(XCAR(spec), Feval(SECOND(spec), Vlexical_environment));
    if (!NILP(THIRD(spec))) {
        new_lexical_variable(THIRD(spec), Qnil);
    }
}

static ALWAYS_INLINE enum ProcessArgsResult
push_interpreted_args_to_lexenv(LispFunction *fobj, LispVal *args) {
    LispVal *rem_req = fobj->args.req;
    LispVal *rem_opt = fobj->args.opt;
    while (!NILP(rem_req)) {
        if (NILP(args)) {
            return PROCESS_ARGS_TOO_FEW;
        }
        new_lexical_variable(XCAR(rem_req), XCAR(args));
        args = XCDR(args);
        rem_req = XCDR(rem_req);
    }
    while (!NILP(rem_opt) && !NILP(args)) {
        push_optional_argument_to_lexenv(XCAR(rem_opt), XCAR(args));
        args = XCDR(args);
        rem_opt = XCDR(rem_opt);
    }
    while (!NILP(rem_opt)) {
        push_missing_optional_argument_to_lexenv(XCAR(rem_opt));
        rem_opt = XCDR(rem_opt);
    }
    if (!NILP(fobj->args.rest)) {
        new_lexical_variable(fobj->args.rest, args);
    }
    if (NILP(fobj->args.kw)) {
        return !NILP(args) && NILP(fobj->args.rest) ? PROCESS_ARGS_TOO_MANY
                                                    : PROCESS_ARGS_OK;
    }
    LispVal *seen_kw = make_hash_table_no_gc(Qnil, Qnil);
    while (!NILP(args)) {
        if (NILP(XCDR(args))) {
            release_hash_table_no_gc(seen_kw);
            return PROCESS_ARGS_NO_KEY_VALUE;
        }
        // has index in front
        LispVal *i_spec = Fgethash(fobj->args.kw, XCAR(args), Qnil);
        if (!NILP(i_spec)) {
            Fputhash(seen_kw, XCAR(args), Qt);
            push_optional_argument_to_lexenv(XCDR(i_spec), SECOND(args));
        } else if (!fobj->args.allow_other_keys) {
            release_hash_table_no_gc(seen_kw);
            return PROCESS_ARGS_BAD_KEY;
        }
        args = XCDR(XCDR(args));
    }
    HT_FOREACH_INDEX(fobj->args.kw, i) {
        if (NILP(Fgethash(seen_kw, HASH_KEY(fobj->args.kw, i), Qnil))) {
            push_missing_optional_argument_to_lexenv(
                XCDR(HASH_VALUE(fobj->args.kw, i)));
        }
    }
    release_hash_table_no_gc(seen_kw);
    return PROCESS_ARGS_OK;
}

static ALWAYS_INLINE LispVal *
call_interpreted(LispVal *orig_func, LispFunction *fobj, LispVal *args) {
    StackFrame *stack_ref = LISP_STACK_REF();
    push_call_frame(orig_func, fobj, args);
    LispVal *evaled_args = evaluate_function_arguments(args);
    set_stack_evaluated_args(LISP_STACK_REF(), evaled_args);
    push_dynamic_binding(Qlexical_environment, fobj->impl.interp.lexenv);
    enum ProcessArgsResult par =
        push_interpreted_args_to_lexenv(fobj, evaled_args);
    if (par != PROCESS_ARGS_OK) {
        // TODO better error handling
        fprintf(stderr, "Bad args to interp func: %s\n",
                process_args_strerror(par));
        abort();
    }
    LispVal *rval = Fprogn(fobj->impl.interp.body);
    the_stack.nogc_retval = rval;
    unwind_to(stack_ref);
    add_local_reference(LISP_STACK_REF(), rval);
    return rval;
}

DEFUN(funcall, "funcall", (LispVal * func, LispVal *args), "(func &rest args)",
      "") {
    LispFunction *fobj = func;
    if (SYMBOLP(func)) {
        fobj = Fsymbol_function(func, Qt);
    } else if (CONSP(func) && EQ(XCAR(func), Qlambda)) {
        fobj = Feval(func, Vlexical_environment);
    }
    // include symbol here for the error message
    CHECK_TYPE(fobj, TYPE_FUNCTION, TYPE_SYMBOL);
    assert(FUNCTIONP(fobj));
    switch (fobj->type) {
    case FUNCTION_NATIVE:
        return call_native(func, fobj, args);
    case FUNCTION_INTERP:
        return call_interpreted(func, fobj, args);
    default:
        abort();
    }
}

static LispVal *parse_lambda_declare_form(LispFunction *fobj, LispVal *body) {
    while (CONSP(body) && CONSP(XCAR(body)) && EQ(XCAR(XCAR(body)), Qdeclare)) {
        LispVal *decls = XCDR(XCAR(body));
        DOLIST(decl, decls) {
            if (EQ(XCAR(decl), Qname)) {
                CHECK_TYPE(SECOND(decl), TYPE_SYMBOL);
                if (!list_length_eq(decl, 2)) {
                    // TODO better error
                    fprintf(stderr, "Invalid (declare (name ...)) form!\n");
                    abort();
                }
                fobj->name = SECOND(decl);
            }
        }
        body = XCDR(body);
    }
    return body;
}

DEFSPECIAL(lambda, "lambda", (LispVal * args, LispVal *body),
           "(args &rest body)", "") {
    LambdaListParseResult llpr;
    parse_lambda_list(&llpr, args);
    if (llpr.status != LLPS_OK) {
        // TODO better handling
        fprintf(stderr,
                "Lambda list parse error: %s: ", llps_strerror(llpr.status));
        debug_print(stderr, args);
        fputc('\n', stderr);
        abort();
    }
    CHECK_LISTP(body);
    LispFunction *fobj = lisp_alloc_object(sizeof(LispFunction), TYPE_FUNCTION);
    fobj->name = Qnil;
    fobj->args = llpr.lambda_list;
    fobj->type = FUNCTION_INTERP;
    if (STRINGP(XCAR(body))) {
        fobj->docstr = XCAR(body);
        if (CONSP(XCDR(body))) {
            body = XCDR(body);
        }
    } else {
        fobj->docstr = Qnil;
    }
    body = parse_lambda_declare_form(fobj, body);
    fobj->impl.interp.body = body;
    fobj->impl.interp.lexenv = Vlexical_environment;
    return fobj;
}

DEFINE_SYMBOL(declare, "declare");
DEFINE_SYMBOL(name, "name");