import "@std/mem"
import "@std/arraylist"
import "@std/enummap"
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)
            .import: {
                try append_text(renderer, "  ")
                try append_text(renderer, ast.source_text(tree, ast.import_decl(tree, declaration).span))
                try append_text(renderer, "\n")
            }
            .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")
            }
            .break_stmt: try append_text(renderer, "break\n")
            .continue_stmt: try append_text(renderer, "continue\n")
            .unreachable_stmt: try append_text(renderer, "unreachable\n")
            .block_stmt: {
                try append_text(renderer, "block\n")
                try render_block(renderer, tree, ast.statement_block(tree, statement), depth + 1)
            }
            .defer_stmt, .errdefer_stmt: {
                cleanup :: ast.cleanup_stmt(tree, statement)
                keyword :: if (ast.stmt_kind(tree, statement) == .errdefer_stmt) "errdefer" else "defer"
                try append_text(renderer, keyword)
                if (cleanup.capture) |capture| {
                    try append_text(renderer, " |")
                    try append_text(renderer, ast.source_text(tree, ast.node_span(tree, capture)))
                    try append_text(renderer, "|")
                }
                try append_text(renderer, "\n")
                try render_block(renderer, tree, cleanup.body, depth + 1)
            }
            .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")
            }
            .yield_stmt: {
                try append_text(renderer, "yield ")
                value :: ast.yield_stmt(tree, statement).value?
                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)
                }
            }
            .assignment_stmt: {
                assignment :: ast.assignment_stmt(tree, statement)
                try append_text(renderer, "assignment ")
                try append_text(renderer, ast.source_text(tree, ast.node_span(tree, assignment.target)))
                value := assignment.value
                if (assignment.compound) {
                    binary :: ast.binary_expr(tree, value)
                    try append_text(renderer, " ")
                    try append_text(renderer, enummap.get(&assignment_operator_texts, binary.operator)?)
                    try append_text(renderer, "= ")
                    value = binary.right
                } else {
                    try append_text(renderer, " = ")
                }
                try append_text(renderer, ast.source_text(tree, ast.node_span(tree, value)))
                try append_text(renderer, "\n")
            }
            .while_stmt: {
                loop :: ast.while_stmt(tree, statement)
                try append_text(renderer, "while ")
                try append_text(renderer, ast.source_text(tree, ast.node_span(tree, loop.condition)))
                try append_text(renderer, "\n")
                try render_block(renderer, tree, loop.body, depth + 1)
            }
            .for_stmt: {
                loop :: ast.for_stmt(tree, statement)
                if (loop.is_inline) try append_text(renderer, "inline ")
                try append_text(renderer, "for ")
                try append_text(renderer, ast.source_text(tree, ast.node_span(tree, loop.iterable)))
                try append_text(renderer, " |")
                if (loop.pointer_capture) try append_text(renderer, "@")
                if (loop.capture) |capture| try append_text(renderer, ast.source_text(tree, ast.node_span(tree, capture)))
                else try append_text(renderer, "_")
                if (loop.index_capture) {
                    try append_text(renderer, ", ")
                    if (loop.index) |index_capture| try append_text(renderer, ast.source_text(tree, ast.node_span(tree, index_capture)))
                    else try append_text(renderer, "_")
                }
                try append_text(renderer, "|\n")
                try render_block(renderer, tree, loop.body, 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))
    if (statement.type) |type| {
        try append_text(renderer, " ")
        try append_text(renderer, ast.source_text(tree, ast.node_span(tree, type)))
    }
    binding :: if (statement.mutability == .constant) " :: " else " := "
    try append_text(renderer, binding)
    try append_text(renderer, ast.source_text(tree, ast.node_span(tree, statement.initializer)))
    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)
}

@hide:file
assignment_operator_texts :: $enummap.init_default(
    ast.BinaryOperator, []u8, "", {
        add = "+",
        subtract = "-",
        multiply = "*",
        divide = "/",
        bit_and = "&",
        bit_or = "|",
        bit_xor = "xor",
        shift_left = "<<",
        shift_right = ">>",
        shift_left_saturating = "<<|",
    },
)