// Tests for lower.zig

const std = @import("std");
const ast = @import("../ast.zig");
const Node = ast.Node;

const ir_mod = @import("ir.zig");
const TypeId = ir_mod.TypeId;
const Ref = ir_mod.Ref;
const FuncId = ir_mod.FuncId;
const Lowering = ir_mod.Lowering;

test "lower: simple function with arithmetic" {
    const alloc = std.testing.allocator;
    var module = ir_mod.Module.init(alloc);
    defer module.deinit();

    // Build a minimal AST: add :: (a: s64, b: s64) -> s64 { return a + b; }
    const a_type = alloc.create(Node) catch unreachable;
    a_type.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .type_expr = .{ .name = "s64", .is_generic = false } } };
    const b_type = alloc.create(Node) catch unreachable;
    b_type.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .type_expr = .{ .name = "s64", .is_generic = false } } };
    const ret_type = alloc.create(Node) catch unreachable;
    ret_type.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .type_expr = .{ .name = "s64", .is_generic = false } } };

    const a_ident = alloc.create(Node) catch unreachable;
    a_ident.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .identifier = .{ .name = "a" } } };
    const b_ident = alloc.create(Node) catch unreachable;
    b_ident.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .identifier = .{ .name = "b" } } };

    const add_expr = alloc.create(Node) catch unreachable;
    add_expr.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .binary_op = .{
        .op = .add,
        .lhs = a_ident,
        .rhs = b_ident,
    } } };

    const ret_stmt = alloc.create(Node) catch unreachable;
    ret_stmt.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .return_stmt = .{ .value = add_expr } } };

    const body = alloc.create(Node) catch unreachable;
    const stmts: []const *Node = &.{ret_stmt};
    body.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .block = .{ .stmts = stmts } } };

    defer alloc.destroy(a_type);
    defer alloc.destroy(b_type);
    defer alloc.destroy(ret_type);
    defer alloc.destroy(a_ident);
    defer alloc.destroy(b_ident);
    defer alloc.destroy(add_expr);
    defer alloc.destroy(ret_stmt);
    defer alloc.destroy(body);

    const params: []const ast.Param = &.{
        .{ .name = "a", .name_span = .{ .start = 0, .end = 0 }, .type_expr = a_type },
        .{ .name = "b", .name_span = .{ .start = 0, .end = 0 }, .type_expr = b_type },
    };

    const fn_decl = ast.FnDecl{
        .name = "add",
        .params = params,
        .return_type = ret_type,
        .body = body,
    };

    var lowering = Lowering.init(&module);
    lowering.lowerFunction(&fn_decl, "add");

    // Verify
    try std.testing.expectEqual(@as(usize, 1), module.functions.items.len);
    const func = module.getFunction(FuncId.fromIndex(0));
    try std.testing.expectEqual(@as(usize, 2), func.params.len);
    try std.testing.expectEqual(TypeId.s64, func.ret);
    try std.testing.expect(func.blocks.items.len > 0);

    // Print the IR to verify it looks reasonable
    const print_mod = @import("print.zig");
    var aw = std.Io.Writer.Allocating.init(alloc);
    try print_mod.printModule(&module, &aw.writer);
    var result = aw.writer.toArrayList();
    defer result.deinit(alloc);

    const output = result.items;
    try std.testing.expect(std.mem.indexOf(u8, output, "func @add") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "entry:") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "add %") != null or std.mem.indexOf(u8, output, "ret %") != null);
}

test "lower: if/else generates basic blocks" {
    const alloc = std.testing.allocator;
    var module = ir_mod.Module.init(alloc);
    defer module.deinit();

    // Build AST: test :: () -> s64 { if true { return 1; } else { return 2; } }
    const cond_node = alloc.create(Node) catch unreachable;
    defer alloc.destroy(cond_node);
    cond_node.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .bool_literal = .{ .value = true } } };

    const ret1_val = alloc.create(Node) catch unreachable;
    defer alloc.destroy(ret1_val);
    ret1_val.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .int_literal = .{ .value = 1 } } };

    const ret2_val = alloc.create(Node) catch unreachable;
    defer alloc.destroy(ret2_val);
    ret2_val.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .int_literal = .{ .value = 2 } } };

    const then_ret = alloc.create(Node) catch unreachable;
    defer alloc.destroy(then_ret);
    then_ret.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .return_stmt = .{ .value = ret1_val } } };

    const else_ret = alloc.create(Node) catch unreachable;
    defer alloc.destroy(else_ret);
    else_ret.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .return_stmt = .{ .value = ret2_val } } };

    const then_body = alloc.create(Node) catch unreachable;
    defer alloc.destroy(then_body);
    then_body.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .block = .{ .stmts = &.{then_ret} } } };

    const else_body = alloc.create(Node) catch unreachable;
    defer alloc.destroy(else_body);
    else_body.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .block = .{ .stmts = &.{else_ret} } } };

    const if_node = alloc.create(Node) catch unreachable;
    defer alloc.destroy(if_node);
    if_node.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .if_expr = .{
        .condition = cond_node,
        .then_branch = then_body,
        .else_branch = else_body,
        .is_inline = false,
    } } };

    const fn_body = alloc.create(Node) catch unreachable;
    defer alloc.destroy(fn_body);
    fn_body.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .block = .{ .stmts = &.{if_node} } } };

    const ret_type = alloc.create(Node) catch unreachable;
    defer alloc.destroy(ret_type);
    ret_type.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .type_expr = .{ .name = "s64", .is_generic = false } } };

    const fn_decl = ast.FnDecl{
        .name = "test_if",
        .params = &.{},
        .return_type = ret_type,
        .body = fn_body,
    };

    var lowering = Lowering.init(&module);
    lowering.lowerFunction(&fn_decl, "test_if");

    // Verify: should have 4 blocks (entry, if.then, if.else, if.merge)
    const func = module.getFunction(FuncId.fromIndex(0));
    try std.testing.expectEqual(@as(usize, 4), func.blocks.items.len);

    // Print and verify structure
    const print_mod = @import("print.zig");
    var aw = std.Io.Writer.Allocating.init(alloc);
    try print_mod.printModule(&module, &aw.writer);
    var result = aw.writer.toArrayList();
    defer result.deinit(alloc);
    const output = result.items;

    try std.testing.expect(std.mem.indexOf(u8, output, "cond_br") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "if.then") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "if.else") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "if.merge") != null);
}

test "lower: while loop generates header/body/exit blocks" {
    const alloc = std.testing.allocator;
    var module = ir_mod.Module.init(alloc);
    defer module.deinit();

    // Build AST: loop :: () { while true { break; } }
    const cond_node = alloc.create(Node) catch unreachable;
    defer alloc.destroy(cond_node);
    cond_node.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .bool_literal = .{ .value = true } } };

    const break_node = alloc.create(Node) catch unreachable;
    defer alloc.destroy(break_node);
    break_node.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .break_expr };

    const while_body = alloc.create(Node) catch unreachable;
    defer alloc.destroy(while_body);
    while_body.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .block = .{ .stmts = &.{break_node} } } };

    const while_node = alloc.create(Node) catch unreachable;
    defer alloc.destroy(while_node);
    while_node.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .while_expr = .{
        .condition = cond_node,
        .body = while_body,
    } } };

    const fn_body = alloc.create(Node) catch unreachable;
    defer alloc.destroy(fn_body);
    fn_body.* = .{ .span = .{ .start = 0, .end = 0 }, .data = .{ .block = .{ .stmts = &.{while_node} } } };

    const fn_decl = ast.FnDecl{
        .name = "loop_test",
        .params = &.{},
        .return_type = null,
        .body = fn_body,
    };

    var lowering = Lowering.init(&module);
    lowering.lowerFunction(&fn_decl, "loop_test");

    // Verify: should have 4 blocks (entry, while.hdr, while.body, while.exit)
    const func = module.getFunction(FuncId.fromIndex(0));
    try std.testing.expectEqual(@as(usize, 4), func.blocks.items.len);

    // Print and verify structure
    const print_mod = @import("print.zig");
    var aw = std.Io.Writer.Allocating.init(alloc);
    try print_mod.printModule(&module, &aw.writer);
    var result = aw.writer.toArrayList();
    defer result.deinit(alloc);
    const output = result.items;

    try std.testing.expect(std.mem.indexOf(u8, output, "while.hdr") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "while.body") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "while.exit") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "cond_br") != null);
}