#include "parse.h"

#include <stdlib.h>

static void parser_next(Parser *parser);

static bool parser_expect(Parser *parser, TokenType type);

Parser *parser_new(Lexer *lexer) {
    Parser *parser = malloc(sizeof(Parser));
    if (parser == NULL) {
        return NULL;
    }
    parser->lexer = lexer;
    parser->flag_error = 0;

    parser_next(parser);
    parser_next(parser);

    return parser;
}

void parser_free(Parser *parser) {
    free(parser);
}

static void parser_next(Parser *parser) {
    parser->current = parser->peek;
    parser->peek = lexer_next_token(parser->lexer);
}

static bool parser_expect(Parser *parser, TokenType type) {
    if (parser->current.type == type) {
        parser_next(parser);
        return true;
    } else {
        parser->flag_error = 1;
        return false;
    }
}

/*
    IMPL. PARSER PARSE FUNCTIONS 
*/

ASTNode *parser_parse_program(Parser *parser) {
    ASTNode *root = ast_node_program();

    while (parser->current.type == VAL) {
        ASTNode *defn_node = parser_parse_defn(parser);

        if (defn_node == NULL) {
            ast_node_free(root);
            parser->flag_error = 1;
            return NULL;
        }
        ast_node_add_child(root, defn_node);
    }

    return root;
}

ASTNode *parser_parse_defn(Parser *parser) {
    Token val_tok = parser->current;
    if (!parser_expect(parser, VAL)) {// must start with VAL
        return NULL;
    }

    ASTNode *type_node = parser_parse_type(parser);
    if (type_node == NULL) {
        return NULL;
    }
    if (parser->current.type != ID) {
        parser->flag_error = 1;
        return NULL;
    }
    Token id_tok = parser->current;
    ASTNode *id_node = ast_node_id(id_tok);
    parser_next(parser);

    ASTNode *expr_node = NULL;
    if (parser->current.type != SEMI) {
        expr_node = parser_parse_expr(parser);
        if (expr_node == NULL) {
            parser->flag_error = 1;
            return NULL;
        }
    }
    if (!parser_expect(parser, SEMI)) {
        return NULL;
    }

    ASTNode *defn_node = ast_node_defn(val_tok, type_node, id_node, expr_node);
    return defn_node;
}

/*
    PARSING TYPE
*/

ASTNode *parser_parse_type(Parser *parser) {
    ASTNode *type_node = NULL;
    if (parser->current.type == ID) {
        type_node = parser_parse_type_simple(parser);
    } else if (parser->current.type == LBRACK) {
        type_node = parser_parse_type_complex(parser);
    } else {
        parser->flag_error = 1;
        return NULL;
    }
    return type_node;
}

ASTNode *parser_parse_type_simple(Parser *parser) {
    Token token = parser->current;
    if (!parser_expect(parser, ID)) {
        return NULL;
    }
    ASTNode *type_node = ast_node_type_simple(token);
    while (parser->current.type == STAR) {
        Token star_tok = parser->current;
        parser_next(parser);

        ASTNode *type_star = ast_node_type_star(star_tok);
        ast_node_add_child(type_node, type_star);
    }

    return type_node;
}

ASTNode *parser_parse_type_complex(Parser *parser) {
    Token tok = parser->current;
    parser_next(parser);

    ASTNode *types[256];
    ASTNode *type_ret = NULL;
    size_t cnt = 0;

    ASTNode *ret = NULL;

    while (parser->current.type != RBRACK && parser->current.type != ARROW) {
        ASTNode *type = parser_parse_type(parser);
        if (type == NULL) {
            return NULL;
        }

        types[cnt++] = type;
    }

    if (parser->current.type == ARROW) {

        parser_next(parser);

        if (parser->current.type != RBRACK) {
            type_ret = parser_parse_type(parser);
            if (type_ret == NULL) return NULL;
        } else {
            type_ret = ast_node_type_void(parser);
        }

        ASTNode *type_param = ast_node_type_param();
        for (size_t i = 0; i < cnt; i++) {
            ast_node_add_child(type_param, types[i]);
        }

        ASTNode *type_out = ast_node_type_out();
        ast_node_add_child(type_out, type_ret);
        ret = ast_node_type_complex(tok, type_param, type_out);
    } else if (parser->current.type == RBRACK) {

        if (cnt >= 2) {
            parser->flag_error = 1;// too many args
            return NULL;
        }
        if (cnt == 1) {
            ret = types[0];
        } else {
            ret = ast_node_type_void(parser);
        }
    } else {
        parser->flag_error = 1;
        return NULL;
    }
    parser_next(parser);
    return ret;
}

/*
    PARSING EXPR
*/

ASTNode *parser_parse_expr(Parser *parser) {
    ASTNode *atoms[256];
    size_t cnt = 1;
    ASTNode *atom_head = parser_parse_atom(parser);
    if (atom_head == NULL) {
        parser->flag_error = 1;
        return NULL;
    }
    atoms[0] = atom_head;

    ASTNode *atom;
    while (true) {
        atom = parser_parse_atom(parser);
        if (atom == NULL) {
            break;
        }
        atoms[cnt] = atom;
        cnt++;
    }

    ASTNode *node = ast_node_expr(atom_head->token);
    for (size_t i = 0; i < cnt; i++) {
        ast_node_add_child(node, atoms[i]);
    }

    return node;
}

ASTNode *parser_parse_atom(Parser *parser) {
    if (parser->current.type == NUM) {
        Token num_tok = parser->current;
        parser_next(parser);
        return ast_node_num(num_tok);
    } else if (parser->current.type == ID) {
        Token id_tok = parser->current;
        parser_next(parser);
        return ast_node_id(id_tok);
    } else if (parser->current.type == STRING_LITERAL) {
        Token str_tok = parser->current;
        parser_next(parser);
        return ast_node_str(str_tok);
    } else if (parser->current.type == STAR) {
        Token star_tok = parser->current;
        parser_next(parser);
        return ast_node_star(star_tok);
    } else if (parser->current.type == ANDREF) {
        Token andref_tok = parser->current;
        parser_next(parser);
        return ast_node_andref(andref_tok);
    }
    
    else {
        return NULL;
    }
}