#include "desugar.h"
#include <stdio.h>
#include <string.h>

// === match → let + if-else 链 ===
AstNode* desugar_match(Parser* p, const Token* match_tok,
                       AstNode* matched_expr,
                       AstNode** arm_patterns, bool* arm_is_wildcard,
                       AstNode** arm_bodies, int arm_count) {
    const char* varname = arena_strdup_impl(p->arena, "__match_val", 12);

    // 从最后一个分支往前构建 if-else 链
    AstNode* result = NULL;
    for (int i = arm_count - 1; i >= 0; i--) {
        if (arm_is_wildcard[i]) {
            AstNode* true_cond = ast_make_literal_bool(p->arena, true, tok_loc(match_tok));
            result = ast_make_if(p->arena, true_cond, arm_bodies[i], result, tok_loc(match_tok));
        } else {
            AstNode* cond = ast_make_binary(p->arena, OP_EQ,
                ast_make_ident(p->arena, varname, tok_loc(match_tok)),
                arm_patterns[i], tok_loc(match_tok));
            result = ast_make_if(p->arena, cond, arm_bodies[i], result, tok_loc(match_tok));
        }
    }

    AstNode* let_stmt = ast_make_let(p->arena, varname, TYPE_UNKNOWN,
        false, false, matched_expr, NULL, 0, NULL, 0, tok_loc(match_tok));

    AstNode* stmts_arr[2] = { let_stmt, result };
    AstNode** stmts = arena_alloc_impl(p->arena, 2 * sizeof(AstNode*));
    memcpy(stmts, stmts_arr, 2 * sizeof(AstNode*));
    return ast_make_block(p->arena, stmts, 2, tok_loc(match_tok));
}

// === guard → if !(cond) { body } ===
AstNode* desugar_guard(Parser* p, const Token* guard_tok,
                       AstNode* cond, AstNode* body) {
    AstNode* not_cond = ast_make_unary(p->arena, OP_NOT, cond, tok_loc(guard_tok));
    return ast_make_if(p->arena, not_cond, body, NULL, tok_loc(guard_tok));
}

// === for i in start to end [step N] { body } → { var i=start; while i<end { body; i=i+step; } } ===
AstNode* desugar_for(Parser* p, const Token* for_tok,
                     const char* var_name, AstNode* start_expr,
                     AstNode* end_expr, AstNode* step_expr, AstNode* body) {
    AstNode* let_stmt = ast_make_let(p->arena, var_name, TYPE_UNKNOWN,
        false, true, start_expr, NULL, 0, NULL, 0, tok_loc(for_tok));

    AstNode* cond = ast_make_binary(p->arena, OP_LT,
        ast_make_ident(p->arena, var_name, tok_loc(for_tok)),
        end_expr, tok_loc(for_tok));

    AstNode* step = step_expr ? step_expr
        : ast_make_literal_i64(p->arena, 1, tok_loc(for_tok));
    AstNode* incr = ast_make_assign(p->arena, var_name,
        ast_make_binary(p->arena, OP_ADD,
            ast_make_ident(p->arena, var_name, tok_loc(for_tok)),
            step, tok_loc(for_tok)),
        tok_loc(for_tok));

    // 增量追加到循环体末尾
    AstNode** new_stmts = arena_alloc_impl(p->arena,
        (body->as.block.stmt_count + 1) * sizeof(AstNode*));
    memcpy(new_stmts, body->as.block.stmts, body->as.block.stmt_count * sizeof(AstNode*));
    new_stmts[body->as.block.stmt_count] = incr;
    AstNode* new_body = ast_make_block(p->arena, new_stmts,
        body->as.block.stmt_count + 1, body->loc);

    // while i < end { body; i = i + 1; }
    AstNode* while_loop = ast_make_while(p->arena, cond, new_body, tok_loc(for_tok));

    AstNode* stmts_arr[2] = { let_stmt, while_loop };
    AstNode** stmts = arena_alloc_impl(p->arena, 2 * sizeof(AstNode*));
    memcpy(stmts, stmts_arr, 2 * sizeof(AstNode*));
    return ast_make_block(p->arena, stmts, 2, tok_loc(for_tok));
}

// === if let pattern = expr { then } else { else } ===
AstNode* desugar_if_let(Parser* p, const Token* if_tok,
                        AstNode* pattern, AstNode* match_expr,
                        AstNode* then_block, AstNode* else_block) {
    static int iflet_counter = 0;
    char vname_buf[32];
    snprintf(vname_buf, sizeof(vname_buf), "__iflet_%d", iflet_counter++);
    const char* vname = arena_strdup_impl(p->arena, vname_buf, strlen(vname_buf));

    AstNode* let_stmt = ast_make_let(p->arena, vname, TYPE_UNKNOWN,
        false, false, match_expr, NULL, 0, NULL, 0, tok_loc(if_tok));
    AstNode* cond = ast_make_binary(p->arena, OP_EQ,
        ast_make_ident(p->arena, vname, tok_loc(if_tok)),
        pattern, tok_loc(if_tok));
    AstNode* if_stmt = ast_make_if(p->arena, cond, then_block, else_block, tok_loc(if_tok));

    AstNode* stmts[2] = { let_stmt, if_stmt };
    AstNode** arr = arena_alloc_impl(p->arena, 2 * sizeof(AstNode*));
    memcpy(arr, stmts, 2 * sizeof(AstNode*));
    return ast_make_block(p->arena, arr, 2, tok_loc(if_tok));
}

// === ident += expr → ident = ident + expr ===
AstNode* desugar_compound_assign(Parser* p, const Token* name_tok,
                                 BinaryOp binop, AstNode* rhs) {
    const char* vname = arena_strdup_impl(p->arena, name_tok->start, name_tok->length);
    AstNode* lhs_ident = ast_make_ident(p->arena, vname, tok_loc(name_tok));
    AstNode* bin_expr = ast_make_binary(p->arena, binop, lhs_ident, rhs, tok_loc(name_tok));
    return ast_make_assign(p->arena, vname, bin_expr, tok_loc(name_tok));
}