#include "type.h"
#include "scope.h"
#include "lexer.h"
#include "parser.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define PARSER_PANIC(format, ...) {\
    fprintf(\
        stderr,\
        "ccc: parse error: %s: line %lu, column %lu: " format "\n",\
        tok.PATH,\
        tok.LINE,\
        tok.COL __VA_OPT__(,)\
        __VA_ARGS__);\
    exit(1);\
}

static struct token tok;
static struct scope* scope;

static void unexpected_token(enum token_type expected) {
    /* TODO: print what token was expected */
    PARSER_PANIC("unexpected token; expected %d", expected);
}

static void peek_or_panic() {
    if (!lexer_peek(&tok))
        PARSER_PANIC("unexpected EOF");
}

static void expect(enum token_type expected) {
    if (!lexer_pop(&tok))
        PARSER_PANIC("unexpected EOF");

    if (tok.type != expected) unexpected_token(expected);
}

static void expect_kw(const char* kw) {
    if (!lexer_pop(&tok))
        PARSER_PANIC("unexpected EOF, expected %s", kw);

    if (tok.type != TK_IDENT)
        PARSER_PANIC("unexpected token, expected %s", kw);

    if (strcmp(kw, tok.data.ident) != 0)
        PARSER_PANIC(
            "unexpected identifier %s, expected %s", tok.data.ident, kw);

    /* string won't go in the AST, discard it */
    free(tok.data.ident);
    tok.data.ident = NULL;
}

static struct type pointer_type_from_alias(
    const struct type_alias* type_alias
) {
    switch (type_alias->type.type) {
        case TP_DATA:
            return (struct type) {
                .type = TP_PTR,
                .pointer.data_type = type_alias->type.data.data_type,
                .pointer.ptr_level = 0,
            };
        case TP_PTR:
            return type_alias->type;
    }
    PARSER_PANIC("unhandled type of type case");
}

static void parse_type_ref(struct type* p_type) {
    /* TODO: modifiers, void rules, arrays, etc. */
    /* TODO: struct, union, enum */
    expect(TK_IDENT);
    const struct data_type* primitive_type = NULL;
    const struct type_alias* type_alias = NULL;

    for (integral_t i = 0; primitive_types[i] != NULL; i++) {
        if (strcmp(tok.data.ident, primitive_types[i]->name) == 0) {
            primitive_type = primitive_types[i];
            break;
        }
    }

    scope_get_type(scope, &type_alias, tok.data.ident);
    if (primitive_type == NULL && type_alias == NULL)
        PARSER_PANIC("unknown type name: '%s'", tok.data.ident);

    free(tok.data.ident);
    peek_or_panic();

    if (tok.type == TK_STAR) {
        if (type_alias != NULL) *p_type = pointer_type_from_alias(type_alias);
        else *p_type = (struct type) {
            .type = TP_PTR,
            .pointer.data_type = primitive_type,
            .pointer.ptr_level = 0,
        };
    } else {
        if (type_alias != NULL) *p_type = type_alias->type;
        else *p_type = (struct type) {
            .type = TP_DATA,
            .data.data_type = primitive_type,
        };
    }

    while (tok.type == TK_STAR) {
        expect(TK_STAR);
        p_type->pointer.ptr_level++;
        peek_or_panic();
    }
}

static void parse_expr(struct expr_node* p_node);

static void parse_literal(struct expr_node* p_node) {
    peek_or_panic();

    switch (tok.type) {
        case TK_INT_LIT:
            expect(TK_INT_LIT);
            p_node->type = EXPR_INT_LIT;
            p_node->inner.int_lit.val = tok.data.int_lit;
            break;
        case TK_CHAR_LIT:
            expect(TK_CHAR_LIT);
            p_node->type = EXPR_CHAR_LIT;
            p_node->inner.char_lit.val = tok.data.char_lit;
            break;
        default:
            PARSER_PANIC("invalid literal type");
    }
}

static void parse_var_ref(struct var_ref_node* p_node) {
    expect(TK_IDENT);
    struct var_def* var_def;
    if (!scope_get_var(scope, &var_def, tok.data.ident))
        PARSER_PANIC("use of undeclared identifier: '%s'", tok.data.ident);

    free(tok.data.ident);
    p_node->def_ref = var_def;
}

static void parse_expr_assign(struct expr_node* p_node) {
    struct expr_node* lval = ccc_alloc(sizeof(struct expr_node));
    *lval = *p_node;

    p_node->type = EXPR_ASSIGN;
    struct assign_node* a_node = &p_node->inner.assign;
    a_node->lval = lval;

    expect(TK_ASSIGN);
    a_node->rval = ccc_alloc(sizeof(struct expr_node));
    parse_expr(a_node->rval);
}

static void parse_expr_list(struct expr_list_node* p_node) {
    for (;;) {
        p_node->expr = ccc_alloc(sizeof(struct expr_node));
        parse_expr(p_node->expr);

        peek_or_panic();
        if (tok.type != TK_COMMA) break;
        expect(TK_COMMA);

        p_node->next = ccc_alloc(sizeof(struct expr_list_node));
        p_node = p_node->next;
    }
}

static void parse_expr_call(struct expr_node* p_node) {
    switch (p_node->type) {
        case EXPR_VAR_REF:
            struct var_def* var_def = p_node->inner.var_ref.def_ref;
            /* TODO: I would like to include functions in the type model rather than checking storage */
            if (var_def->loc.type != STO_FN)
                PARSER_PANIC("called object is not a function");

            p_node->inner.call.called_fn_ref = var_def->loc.decl;
            break;
        default:
            PARSER_PANIC("expression is not callable");
    }

    p_node->type = EXPR_CALL;
    p_node->inner.call.args = NULL;

    expect(TK_LPAREN);
    peek_or_panic();
    if (tok.type != TK_RPAREN) {
        p_node->inner.call.args =
            ccc_alloc(sizeof(struct expr_list_node));
        parse_expr_list(p_node->inner.call.args);
    }
    expect(TK_RPAREN);
}

static void parse_unary(struct unary_node* p_node) {
    peek_or_panic();
    switch (tok.type) {
        case TK_NEG:
            expect(TK_NEG);
            p_node->op = UNARY_NEG;
            break;
        default:
            PARSER_PANIC("expected unary operator");
    }

    p_node->expr = ccc_alloc(sizeof(struct expr_node));
    parse_expr(p_node->expr);
}

static void parse_expr_binary(struct expr_node* p_node) {
    struct expr_node* lhs = ccc_alloc(sizeof(struct expr_node));
    struct expr_node* rhs = ccc_alloc(sizeof(struct expr_node));

    *lhs = *p_node;
    *p_node = (struct expr_node) {
        .type = EXPR_BINARY,
        .inner.binary = {
            .lhs = lhs,
            .rhs = rhs,
        },
    };

    peek_or_panic();
    switch (tok.type) {
        case TK_PLUS:
            expect(TK_PLUS);
            p_node->inner.binary.op = BINARY_ADD;
            break;
        case TK_NEG:
            expect(TK_NEG);
            p_node->inner.binary.op = BINARY_SUB;
            break;
        case TK_STAR:
            expect(TK_STAR);
            p_node->inner.binary.op = BINARY_MUL;
            break;
        case TK_DIV:
            expect(TK_DIV);
            p_node->inner.binary.op = BINARY_DIV;
            break;
        default:
            PARSER_PANIC("expected binary operator");
    }

    parse_expr(rhs);
}

static void parse_expr(struct expr_node* p_node) {
    peek_or_panic();
    switch (tok.type) {
        case TK_LPAREN:
            expect(TK_LPAREN);
            parse_expr(p_node);
            expect(TK_RPAREN);
            break;
        case TK_NEG:
            p_node->type = EXPR_UNARY;
            parse_unary(&p_node->inner.unary);
            break;
        case TK_INT_LIT:
        case TK_CHAR_LIT:
        case TK_FLOAT_LIT:
        case TK_STR_LIT:
            parse_literal(p_node);
            break;
        case TK_IDENT:
            p_node->type = EXPR_VAR_REF;
            parse_var_ref(&p_node->inner.var_ref);
            break;
        default:
            PARSER_PANIC("expected expression");
    }

    peek_or_panic();
    switch (tok.type) {
        case TK_ASSIGN:
            parse_expr_assign(p_node);
            break;
        case TK_LPAREN:
            parse_expr_call(p_node);
            break;
        case TK_PLUS:
        case TK_NEG:
        case TK_STAR:
        case TK_DIV:
            parse_expr_binary(p_node);
            break;
        default:
    }
}

static void parse_var_decl(struct var_decl_node* p_node) {
    parse_type_ref(&p_node->type);
    expect(TK_IDENT);

    p_node->def_ref = scope_define_var(scope, (struct var_def) {
        .type = &p_node->type,
        .name = tok.data.ident,
        .loc.type = STO_UNRESOLVED,
    });
    if (p_node->def_ref == NULL)
        PARSER_PANIC("redefinition of '%s'", tok.data.ident);

    peek_or_panic();
    if (tok.type == TK_ASSIGN) {
        expect(TK_ASSIGN);
        p_node->initial_value = ccc_alloc(sizeof(struct expr_node));
        parse_expr(p_node->initial_value);
    }
}

static void parse_stmt(struct stmt_node* p_node);

static void parse_return(struct return_node* p_node) {
    expect_kw("return");

    peek_or_panic();
    if (tok.type == TK_SEMI) {
        p_node->ret_val = NULL;
        return;
    }

    p_node->ret_val = ccc_alloc(sizeof(struct expr_node));
    parse_expr(p_node->ret_val);
}

static void parse_group(struct group_node* p_node) {
    expect(TK_LCURLY);

    scope_push(&scope);
    p_node->scope = scope;

    struct stmt_node** pp_node = &p_node->head;
    for (;;) {
        peek_or_panic();
        if (tok.type == TK_RCURLY) break;

        *pp_node = ccc_alloc(sizeof(struct stmt_node));
        parse_stmt(*pp_node);
        pp_node = &((*pp_node)->next);
    }

    scope_pop(&scope);

    expect(TK_RCURLY);
}

static void parse_if(struct if_node* p_node) {
    expect_kw("if");

    scope_push(&scope);
    p_node->scope = scope;

    expect(TK_LPAREN);
    p_node->cond = ccc_alloc(sizeof(struct expr_node));
    parse_expr(p_node->cond);
    expect(TK_RPAREN);

    p_node->true_branch = ccc_alloc(sizeof(struct stmt_node));
    parse_stmt(p_node->true_branch);

    peek_or_panic();
    if (tok.type == TK_IDENT && strcmp(tok.data.ident, "else") == 0) {
        expect_kw("else");

        p_node->false_branch = ccc_alloc(sizeof(struct stmt_node));
        parse_stmt(p_node->false_branch);
    }

    scope_pop(&scope);
}

static void parse_while(struct loop_node* p_node) {
    expect_kw("while");

    scope_push(&scope);
    p_node->scope = scope;

    expect(TK_LPAREN);
    p_node->cond = ccc_alloc(sizeof(struct expr_node));
    parse_expr(p_node->cond);
    expect(TK_RPAREN);

    p_node->body = ccc_alloc(sizeof(struct stmt_node));
    parse_stmt(p_node->body);

    scope_pop(&scope);
}

static bool is_type_token() {
    if (tok.type != TK_IDENT) return false;

    /* check primitive types */
    for (integral_t i = 0; primitive_types[i] != NULL; i++) {
        if (strcmp(tok.data.ident, primitive_types[i]->name) == 0) return true;
    }

    /* check type aliases */
    return scope_get_type(scope, NULL, tok.data.ident);
}

static void parse_loop_init(struct loop_init_node* p_node) {
    peek_or_panic();
    if (!is_type_token()) {
        p_node->type = INIT_EXPR_LIST;
        p_node->expr_list = ccc_alloc(sizeof(struct expr_list_node));
        parse_expr_list(p_node->expr_list);
    } else {
        p_node->type = INIT_DECL;
        p_node->decl = ccc_alloc(sizeof(struct var_decl_node));
        parse_var_decl(p_node->decl);
    }
}

static void parse_for(struct loop_node* p_node) {
    expect_kw("for");

    scope_push(&scope);
    p_node->scope = scope;

    expect(TK_LPAREN);
    peek_or_panic();
    if (tok.type != TK_SEMI) {
        p_node->init = ccc_alloc(sizeof(struct loop_init_node));
        parse_loop_init(p_node->init);
    }
    expect(TK_SEMI);
    peek_or_panic();
    if (tok.type != TK_SEMI) {
        p_node->cond = ccc_alloc(sizeof(struct expr_node));
        parse_expr(p_node->cond);
    }
    expect(TK_SEMI);
    peek_or_panic();
    if (tok.type != TK_RPAREN) {
        p_node->incr = ccc_alloc(sizeof(struct expr_node));
        parse_expr(p_node->incr);
    }
    expect(TK_RPAREN);

    p_node->body = ccc_alloc(sizeof(struct stmt_node));
    parse_stmt(p_node->body);

    scope_pop(&scope);
}

static void parse_stmt(struct stmt_node* p_node) {
    peek_or_panic();
    switch (tok.type) {
        case TK_SEMI:
            p_node->type = STMT_EMPTY;
            break;
        case TK_LCURLY:
            p_node->type = STMT_GROUP;
            parse_group(&p_node->inner.group);
            return;
        case TK_IDENT:
            if (strcmp(tok.data.ident, "if") == 0) {
                p_node->type = STMT_IF;
                parse_if(&p_node->inner.if_);
                return;
            } else if (strcmp(tok.data.ident, "while") == 0) {
                p_node->type = STMT_LOOP;
                parse_while(&p_node->inner.loop);
                return;
            } else if (strcmp(tok.data.ident, "for") == 0) {
                p_node->type = STMT_LOOP;
                parse_for(&p_node->inner.loop);
                return;
            } else if (strcmp(tok.data.ident, "return") == 0) {
                p_node->type = STMT_RETURN;
                parse_return(&p_node->inner.return_);
                break;
            } else if (is_type_token()) {
                p_node->type = STMT_VAR_DECL;
                parse_var_decl(&p_node->inner.var_decl);
                break;
            }
        default:
            p_node->type = STMT_EXPR;
            parse_expr(&p_node->inner.expr);
    }

    expect(TK_SEMI);
}

static void parse_decl_list(struct decl_list_node* p_node) {
    for (;;) {
        p_node->decl = ccc_alloc(sizeof(struct var_decl_node));
        parse_var_decl(p_node->decl);

        peek_or_panic();
        if (tok.type != TK_COMMA) break;
        expect(TK_COMMA);

        p_node->next = ccc_alloc(sizeof(struct decl_list_node));
        p_node = p_node->next;
    }
}

static void parse_fn_decl(struct fn_decl_node* p_node) {
    if (scope->next_out != NULL)
        PARSER_PANIC("functions can only be define in the root scope");

    parse_type_ref(&p_node->return_type);

    expect(TK_IDENT);

    if (scope_define_var(scope, (struct var_def) {
                .name = tok.data.ident,
                .loc = {
                    .type = STO_FN,
                    .decl = p_node,
                },
            }) == NULL)
        PARSER_PANIC("redefinition of '%s'", tok.data.ident)

    p_node->name = tok.data.ident;

    expect(TK_LPAREN);

    scope_push(&scope);
    p_node->scope = scope;

    peek_or_panic();
    if (tok.type != TK_RPAREN) {
        p_node->args = ccc_alloc(sizeof(struct decl_list_node));
        parse_decl_list(p_node->args);
    }

    expect(TK_RPAREN);

    parse_group(&p_node->body);

    scope_pop(&scope);
}

static bool parse_root(struct root_node* p_node) {
    if (!lexer_peek(&tok)) return false;

    p_node->type = ROOT_FN_DECL;
    parse_fn_decl(&p_node->inner.fn_decl);
    return true;
}

void parse(const char* path, struct ast* ast) {
    lexer_load(path);
    scope_push(&scope);
    ast->root_scope = scope;

    struct root_node** p_node = &ast->root_node;

    for (;;) {
        *p_node = ccc_alloc(sizeof(struct root_node));
        if (!parse_root(*p_node)) {
            free(*p_node);
            *p_node = NULL;
            break;
        }
        p_node = &((*p_node)->next);
    }

    scope_pop(&scope);
    lexer_close();
}