import "@std/mem"
import "@std/arraylist"
import "@source/ast"

#! owns output bytes and retains its allocator; release with deinit.
Renderer :: struct {
    bytes arraylist.ArrayList(u8)
}

#! retains `allocator` for rendered output; the caller owns the returned Renderer.
init proc(allocator mem.Allocator) Renderer {
    return Renderer{ bytes = arraylist.init(allocator) }
}

#! releases output bytes with the retained allocator.
deinit proc(renderer @mut Renderer) void {
    arraylist.deinit(&renderer.bytes)
}

#! borrows `tree` during rendering; output stays renderer-owned until render or deinit.
render proc(renderer @mut Renderer, tree @ast.Tree) void ! mem.AllocError {
    arraylist.clear(&renderer.bytes)
    try append_text(renderer, "package\n")
    for 0..ast.package_decl_count(tree) |index| {
        declaration :: ast.package_decl(tree, index)
        match ast.decl_kind(tree, declaration) {
            .statement: try render_statement(renderer, tree, declaration, "  ")
            .function: try render_function(renderer, tree, declaration)
            .invalid: try append_text(renderer, "  invalid declaration\n")
        }
    }
}

@hide:file
render_function proc(
    renderer @mut Renderer,
    tree @ast.Tree,
    declaration ast.NodeId,
) void ! mem.AllocError {
    function :: ast.function_decl(tree, declaration)
    try append_text(renderer, "  function ")
    try append_text(renderer, ast.source_text(tree, function.name))
    try append_text(renderer, "\n")
    try render_block(renderer, tree, function.body, 2)
}

@hide:file
render_block proc(renderer @mut Renderer, tree @ast.Tree, block ast.NodeId, depth usize) void ! mem.AllocError {
    for 0..ast.block_statement_count(tree, block) |index| {
        statement :: ast.block_statement(tree, block, index)
        for 0..depth |_| try append_text(renderer, "  ")
        match ast.stmt_kind(tree, statement) {
            .declaration: try render_statement(renderer, tree, statement, "")
            .expression: {
                expression :: ast.statement_expr(tree, statement)
                try append_text(renderer, "expression ")
                try append_text(renderer, ast.source_text(tree, ast.node_span(tree, expression)))
                try append_text(renderer, "\n")
            }
            .return_stmt: {
                try append_text(renderer, "return")
                if (ast.return_stmt(tree, statement).value) |value| {
                    try append_text(renderer, " ")
                    try append_text(renderer, ast.source_text(tree, ast.node_span(tree, value)))
                }
                try append_text(renderer, "\n")
            }
            .if_stmt: {
                conditional :: ast.if_stmt(tree, statement)
                try append_text(renderer, "if ")
                try append_text(renderer, ast.source_text(tree, ast.node_span(tree, conditional.condition)))
                try append_text(renderer, "\n")
                try render_block(renderer, tree, conditional.then_block, depth + 1)
                if (conditional.else_block) |alternative| {
                    for 0..depth |_| try append_text(renderer, "  ")
                    try append_text(renderer, "else\n")
                    try render_block(renderer, tree, alternative, depth + 1)
                }
            }
            .invalid: try append_text(renderer, "invalid statement\n")
        }
    }
}

@hide:file
render_statement proc(
    renderer @mut Renderer,
    tree @ast.Tree,
    declaration ast.NodeId,
    indent []u8,
) void ! mem.AllocError {
    statement :: ast.statement_decl(tree, declaration)
    try append_text(renderer, indent)
    try append_text(renderer, "declaration ")
    try append_text(renderer, ast.source_text(tree, statement.name))
    try append_text(renderer, "\n")
}

@hide:file
append_text proc(renderer @mut Renderer, text []u8) void ! mem.AllocError {
    for (text) |byte| try arraylist.append(&renderer.bytes, byte)
}