#include "parser.h"
#include "lexer.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 void* protected_alloc(size_t sz) {
    void* ptr = calloc(1, sz);
    if (ptr == NULL) {
        fprintf(stderr, "ccc: out of memory\n");
        exit(1);
    }
    return ptr;
}

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

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

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

/* "handle" indicates that we've peeked already */
static void handle_expr(struct expr_node* p_node);
static void handle_stmt(struct stmt_node* p_node);

static void parse_type(struct type_node* p_node) {
    expect(TK_IDENT);
    /* TODO: maybe check that this type is real haha */
    p_node->type = tok.data.ident;
}

static void parse_return(struct return_node* p_node) {
    expect(TK_IDENT);
    if (strcmp(tok.data.ident, "return") != 0)
        PARSER_PANIC("unexpected token %s; expected: return", tok.data.ident);

    if (!lexer_peek(&tok))
        PARSER_PANIC("unexpected EOF in return statement");

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

    p_node->ret_val = protected_alloc(sizeof(struct expr_node));
    handle_expr(p_node->ret_val);
}

static void parse_int_lit(struct int_lit_node* p_node) {
    expect(TK_INT_LIT);
    p_node->val = tok.data.int_lit;
}

static void handle_expr(struct expr_node* p_node) {
    switch (tok.type) {
        case TK_SEMI:
            p_node->type = EXPR_EMPTY;
            return;
        case TK_INT_LIT:
            p_node->type = EXPR_INT_LIT;
            parse_int_lit(&p_node->as._int_lit);
            return;
        default:
            PARSER_PANIC("expected expression");
    }
}

static void parse_var_decl(struct var_decl_node* p_node) {
    parse_type(&p_node->type);
    expect(TK_IDENT);
    p_node->ident = tok.data.ident;
}

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

    struct stmt_node** pp_node = &p_node->body_head;
    for (;;) {
        if (!lexer_peek(&tok))
            PARSER_PANIC("unexpected EOF in statement group");

        if (tok.type == TK_RCURLY) break;

        *pp_node = protected_alloc(sizeof(struct stmt_node));
        handle_stmt(*pp_node);
        pp_node = &((*pp_node)->next);
    }

    expect(TK_RCURLY);
}

static void handle_stmt(struct stmt_node* p_node) {
    switch (tok.type) {
        case TK_LCURLY:
            p_node->type = STMT_GROUP;
            parse_group(&p_node->as._group);
            return;
        case TK_IDENT:
            if (strcmp(tok.data.ident, "return") == 0) {
                p_node->type = STMT_RETURN;
                parse_return(&p_node->as._return);
            } else {
                p_node->type = STMT_VAR_DECL;
                parse_var_decl(&p_node->as._var_decl);
            }
            break;
        default:
            p_node->type = STMT_EXPR;
            handle_expr(&p_node->as._expr);
    }
    expect(TK_SEMI);
}

static void parse_arg_list(struct var_decl_node** pp_arg) {
    for (;;) {
        *pp_arg = protected_alloc(sizeof(struct var_decl_node));
        parse_var_decl(*pp_arg);
        pp_arg = &((*pp_arg)->next);

        if (!lexer_peek(&tok))
            PARSER_PANIC("unexpected EOF in argument list");

        if (tok.type == TK_RPAREN) break;
        expect(TK_COMMA);
    }
}

static void parse_fn_decl(struct fn_decl_node* p_node) {
    parse_type(&p_node->return_type);

    expect(TK_IDENT);
    p_node->name = tok.data.ident;

    expect(TK_LPAREN);

    if (!lexer_peek(&tok))
        PARSER_PANIC("unexpected EOF in function declaration");


    if (tok.type != TK_RPAREN) parse_arg_list(&p_node->args_head);

    expect(TK_RPAREN);

    parse_group(&p_node->body);
}

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->as._fn_decl);
    return true;
}

struct root_node* parse(const char* path) {
    lexer_load(path);

    struct root_node* root;
    struct root_node** p_node = &root;

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

    lexer_close();
    return root;
}
