/*
A.2.3 Statements

(6.8)
    statement:
        labeled-statement
        compound-statement
        expression-statement
        selection-statement
        iteration-statement
        jump-statement
(6.8.1)
    labeled-statement:
        identifier : statement
        case constant-expression : statement
        default : statement
(6.8.2)
    compound-statement:
        { block-item-list(opt) }
(6.8.2)
    block-item-list:
        block-item
        block-item-list block-item
(6.8.2)
    block-item:
        declaration
        statement
(6.8.3)
    expression-statement:
        expression(opt) ;
(6.8.4)
    selection-statement:
        if ( expression ) statement
        if ( expression ) statement else statement
        switch ( expression ) statement
(6.8.5)
    iteration-statement:
        while ( expression ) statement
        do statement while ( expression );
        for ( expression(opt) ; expression(opt) ; expression(opt) ) statement
        for ( declaration expression(opt) ; expression(opt) ) statement
(6.8.6)
    jump-statement:
        goto identifier ;
        continue ;
        break ;
        return expression(opt) ;
*/

#include <parser_utils.h>
#include <scc_parser.h>

static inline scc_ast_stmt_t *ast_stmt_alloc() {
    scc_ast_stmt_t *stmt = (scc_ast_stmt_t *)scc_malloc(sizeof(scc_ast_stmt_t));
    if (stmt == null) {
        LOG_FATAL("Out of memory");
    }
    return stmt;
}

static inline scc_ast_expr_t *ast_parse_paren_expression(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_L_PAREN)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected '(' before like `( expression )` .");
    }

    scc_ast_expr_t *ret = scc_parse_expression(parser);

    if (!scc_parser_consume_if(parser, SCC_TOK_R_PAREN)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected ')' before like `( expression )` .");
    }
    return ret;
}

static scc_ast_stmt_t *parse_label_statement(scc_parser_t *parser) {
    scc_lexer_tok_t tok = {0};
    if (!scc_parser_next_consume(parser, &tok)) {
        return null;
    }

    if (!scc_parser_consume_if(parser, SCC_TOK_COLON)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected identifier before like `identifier : statement` .");
    }

    scc_ast_stmt_t *statement = scc_parse_statement(parser);
    if (statement == null) {
        Panic("expect stmt");
    }

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    Assert(stmt != null);
    scc_ast_stmt_label_init(stmt, scc_cstring_as_cstr(&tok.lexeme), statement);
    return stmt;
}

static scc_ast_stmt_t *parse_case_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_CASE)) {
        return null;
    }

    scc_ast_expr_t *expr = null;
    expr = scc_parser_constant_expression(parser);
    if (expr == null) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected constant expression after case.");
        return null;
    }

    if (!scc_parser_consume_if(parser, SCC_TOK_COLON)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected `:` after case.");
        return null;
    }

    scc_ast_stmt_t *statement = scc_parse_statement(parser);
    if (statement == null) {
        Panic("expect stmt");
    }

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_case_init(stmt, expr, statement);
    return stmt;
}

static scc_ast_stmt_t *parse_default_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_DEFAULT)) {
        return null;
    }

    if (!scc_parser_consume_if(parser, SCC_TOK_COLON)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected constant expression after case.");
        return null;
    }

    scc_ast_stmt_t *statement = scc_parse_statement(parser);
    if (statement == null) {
        Panic("expect stmt");
    }

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_default_init(stmt, statement);
    return stmt;
}

static scc_ast_stmt_t *parse_compound_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_L_BRACE)) {
        return null;
    }
    scc_ast_block_item_vec_t block_items;
    scc_vec_init(block_items);

    while (!scc_parser_consume_if(parser, SCC_TOK_R_BRACE)) {
        /// TODO
        // scc_parse_is_decl();
        scc_ast_node_t *ret = null;
        ret = (scc_ast_node_t *)scc_parse_declaration(parser);
        if (ret == null) {
            ret = (scc_ast_node_t *)scc_parse_statement(parser);
        }
        if (ret == null) {
            SCC_ERROR(scc_parser_got_current_pos(parser), "Invalid statement");
            // TODO free
            parser->errcode = 1;
            return null;
        }
        scc_vec_push(block_items, ret);
    }

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_compound_init(stmt, &block_items);
    return stmt;
}

static scc_ast_stmt_t *parse_if_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_IF)) {
        return null;
    }

    scc_ast_expr_t *expression = ast_parse_paren_expression(parser);
    scc_ast_stmt_t *statement = scc_parse_statement(parser);
    scc_ast_stmt_t *opt_else = null;

    if (scc_parser_consume_if(parser, SCC_TOK_ELSE)) {
        opt_else = scc_parse_statement(parser);
    } else {
        opt_else = null;
    }

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_if_init(stmt, expression, statement, opt_else);
    return stmt;
}

static scc_ast_stmt_t *parse_switch_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_SWITCH)) {
        return null;
    }

    scc_ast_expr_t *expression = ast_parse_paren_expression(parser);
    scc_ast_stmt_t *statement = scc_parse_statement(parser);

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_switch_init(stmt, expression, statement);
    return stmt;
}

static scc_ast_stmt_t *parse_while_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_WHILE)) {
        return null;
    }

    scc_ast_expr_t *expression = ast_parse_paren_expression(parser);
    scc_ast_stmt_t *statement = scc_parse_statement(parser);

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_while_init(stmt, expression, statement);
    return stmt;
}

static scc_ast_stmt_t *parse_do_while_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_DO)) {
        return null;
    }

    scc_ast_stmt_t *statement = scc_parse_statement(parser);

    if (!scc_parser_consume_if(parser, SCC_TOK_WHILE)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected 'while' after do.");
        // TODO 使用更好的错误处理，未来应当采用更好的内存管理器
        scc_free(statement);
        return null;
    }
    scc_ast_expr_t *expression = ast_parse_paren_expression(parser);

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_do_while_init(stmt, expression, statement);
    return stmt;
}

static scc_ast_stmt_t *parse_for_statement(scc_parser_t *parser) {
    if (!scc_parser_consume_if(parser, SCC_TOK_FOR)) {
        return null;
    }

    /*
    for ( expression(opt) ; expression(opt) ; expression(opt) ) statement
    for ( declaration expression(opt) ; expression(opt) ) statement
    */

    if (!scc_parser_consume_if(parser, SCC_TOK_L_PAREN)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected '(' before like `( expression )` .");
    }

    scc_ast_type_t *init = null;
    scc_ast_expr_t *cond = null;
    scc_ast_expr_t *incr = null;
    scc_ast_stmt_t *body = null;

    // TODO use decl or expr
    init = (scc_ast_type_t *)scc_parse_declaration(parser);
    if (init == null) {
        init = (scc_ast_type_t *)scc_parse_expression(parser);
        if (!scc_parser_consume_if(parser, SCC_TOK_SEMICOLON)) {
            SCC_ERROR(scc_parser_got_current_pos(parser),
                      "Expected semicolon in for statement.");
        }
    }

    cond = scc_parse_expression(parser);

    if (!scc_parser_consume_if(parser, SCC_TOK_SEMICOLON)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected semicolon in for statement.");
    }

    incr = scc_parse_expression(parser);

    if (!scc_parser_consume_if(parser, SCC_TOK_R_PAREN)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected ')' after like `( expression )` .");
    }

    body = scc_parse_statement(parser);

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_for_init(stmt, init, cond, incr, body);
    return stmt;
}

static scc_ast_stmt_t *parse_jump_statement(scc_parser_t *parser) {
    scc_ast_stmt_t *stmt = ast_stmt_alloc();

    if (scc_parser_consume_if(parser, SCC_TOK_GOTO)) {
        scc_lexer_tok_t tok = {0};
        if (scc_parser_next_consume(parser, &tok)) {
            scc_ast_stmt_goto_init(stmt, scc_cstring_as_cstr(&tok.lexeme));
        } else {
            SCC_ERROR(scc_parser_got_current_pos(parser),
                      "Expected label after goto.");
        }
    } else if (scc_parser_consume_if(parser, SCC_TOK_CONTINUE)) {
        scc_ast_stmt_continue_init(stmt);
    } else if (scc_parser_consume_if(parser, SCC_TOK_BREAK)) {
        scc_ast_stmt_break_init(stmt);
    } else if (scc_parser_consume_if(parser, SCC_TOK_RETURN)) {
        scc_ast_stmt_return_init(stmt, scc_parse_expression(parser));
    } else {
        UNREACHABLE();
    }

    if (!scc_parser_consume_if(parser, SCC_TOK_SEMICOLON)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected semicolon after jump statement.");
    }
    return stmt;
}

static scc_ast_stmt_t *parse_expression_statement(scc_parser_t *parser) {

    if (scc_parser_consume_if(parser, SCC_TOK_SEMICOLON)) {
        scc_ast_stmt_t *stmt = ast_stmt_alloc();
        scc_ast_stmt_expr_init(stmt, null);
        return stmt;
    }

    scc_ast_expr_t *expr = scc_parse_expression(parser);
    if (expr == null) {
        return null;
    }

    scc_ast_stmt_t *stmt = ast_stmt_alloc();
    scc_ast_stmt_expr_init(stmt, expr);

    if (!scc_parser_consume_if(parser, SCC_TOK_SEMICOLON)) {
        SCC_ERROR(scc_parser_got_current_pos(parser),
                  "Expected semicolon after expression.");
    }
    return stmt;
}

scc_ast_stmt_t *scc_parse_statement(scc_parser_t *parser) {
    scc_ast_stmt_t *stmt;
    const scc_lexer_tok_t *tok_ref;
    tok_ref = scc_parser_peek(parser);
    if (!tok_ref) {
        return null;
    }
    switch (tok_ref->type) {
    /*
    (6.8.1)
        labeled-statement:
            identifier : statement
            case constant-expression : statement
            default : statement
    */
    case SCC_TOK_IDENT:
        scc_parser_next(parser);
        tok_ref = scc_parser_next(parser);
        if (tok_ref == null || tok_ref->type != SCC_TOK_COLON) {
            scc_parser_reset(parser);
            break;
        }
        scc_parser_reset(parser);
        stmt = parse_label_statement(parser);
        goto RETURN;
    case SCC_TOK_CASE: {
        stmt = parse_case_statement(parser);
        goto RETURN;
    }
    case SCC_TOK_DEFAULT:
        stmt = parse_default_statement(parser);
        goto RETURN;
    /*
    (6.8.2)
        compound-statement:
            { block-item-list(opt) }
    (6.8.2)
        block-item-list:
            block-item
            block-item-list block-item
    (6.8.2)
        block-item:
            declaration
            statement
    */
    case SCC_TOK_L_BRACE:
        stmt = parse_compound_statement(parser);
        goto RETURN;
    /*
    (6.8.4)
        selection-statement:
            if ( expression ) statement
            if ( expression ) statement else statement
            switch ( expression ) statement
    */
    case SCC_TOK_IF:
        stmt = parse_if_statement(parser);
        goto RETURN;
    case SCC_TOK_SWITCH:
        stmt = parse_switch_statement(parser);
        goto RETURN;
    /*
    (6.8.5)
        iteration-statement:
            while ( expression ) statement
            do statement while ( expression );
            for ( expression(opt) ; expression(opt) ; expression(opt) )
                statement
            for ( declaration expression(opt) ; expression(opt) )
                statement
    */
    case SCC_TOK_WHILE:
        stmt = parse_while_statement(parser);
        goto RETURN;
    case SCC_TOK_DO:
        stmt = parse_do_while_statement(parser);
        goto RETURN;
    case SCC_TOK_FOR:
        stmt = parse_for_statement(parser);
        goto RETURN;
    /*
    (6.8.6)
        jump-statement:
            goto identifier ;
            continue ;
            break ;
            return expression(opt) ;
    */
    case SCC_TOK_GOTO:
    case SCC_TOK_CONTINUE:
    case SCC_TOK_BREAK:
    case SCC_TOK_RETURN:
        stmt = parse_jump_statement(parser);
        goto RETURN;
    default:
        break;
    }
    /*
    (6.8.3)
        expression-statement:
            expression(opt) ;
    */
    stmt = parse_expression_statement(parser);
RETURN:
    scc_parser_reset(parser);
    if (stmt) {
        parser->sema_callbacks.on_stmt(parser->sema_callbacks.context,
                                       stmt->base.type, stmt);
    }
    return stmt;
}