import "@std"
import "@std/mem"
import "@std/arraylist"

import "@source/lexer"
import "@source/parser"
import "@source/ast"

@hide:file
TokenId :: alias lexer.TokenId

@hide:file
Token :: alias lexer.Token

@hide:file
Node :: alias ast.Node

@hide:file
NodeId :: alias ast.NodeId

@hide:file
ParseState :: alias parser.State

Renderer :: struct { tokens std.ArrayList(Token) }

init proc(allocator mem.Allocator) Renderer {
    return Renderer{ tokens = arraylist.init(allocator) }
}

deinit proc(renderer @mut Renderer) void {
    arraylist.deinit(&renderer.tokens)
}

render proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token
) void ! mem.AllocError {
    node :: parse_state.nodes.items[usize(node_id)]
    match node.kind {
        .literal_int, .literal_float, .literal_string:
            _ = try render_lit(renderer, node_id, parse_state, tokens)
        .stmt_decl:
            try render_decl(renderer, node_id, parse_state, tokens)
        .package:
            try render_pkg(renderer, node_id, parse_state, tokens)
        else: unreachable
    }
}

@hide:file
render_pkg proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) void ! mem.AllocError {
    node :: parse_state.nodes.items[usize(node_id)]
    extra_start :: usize(node.data0.extra_id)
    extra_end :: usize(node.data1.extra_id)
    for extra_start..extra_end |i| {
        stmt_node_id :: ast.node_id(parse_state.extra.items[i])
        stmt_node :: parse_state.nodes.items[usize(stmt_node_id)]
        match stmt_node.kind {
            .stmt_decl:
                try render_decl(renderer, stmt_node_id, parse_state, tokens)
            else: unreachable
        }
    }
}

@hide:file
render_decl proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) void ! mem.AllocError {
    node :: parse_state.nodes.items[usize(node_id)]
    tok_ident :: tokens[usize(node.main_token)]
    try arraylist.append(&renderer.tokens, tok_ident)

    raw_node_id_type :: parse_state.extra.items[usize(node.data0.extra_id)]
    node_id_type :: ast.maybe_node_id(raw_node_id_type)
    if (node_id_type) |id| try render_type(renderer, id, parse_state, tokens)

    tok_assign :: tokens[usize(node.data1.token_id)]
    try arraylist.append(&renderer.tokens, tok_assign)

    node_id_expr :: ast.node_id(parse_state.extra.items[usize(node.data0.extra_id + 1)])
    last_expr_token :: try render_expr(renderer, node_id_expr, parse_state, tokens)

    terminator_idx :: usize(last_expr_token) + 1
    tok_terminator :: if (terminator_idx < tokens.len) tokens[terminator_idx] else return
    try arraylist.append(&renderer.tokens, tok_terminator)

    eof_idx :: terminator_idx + 1
    if (tok_terminator.kind == .newline and
        eof_idx < tokens.len and
        tokens[eof_idx].kind == .eof
    ) try arraylist.append(&renderer.tokens, tokens[eof_idx])
}

@hide:file
render_expr proc(renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) TokenId ! mem.AllocError {
    # todo: support unary expressions
    node :: parse_state.nodes.items[usize(node_id)]
    return match node.kind {
        .expr_identifier: {
            tok :: tokens[usize(node.main_token)]
            try arraylist.append(&renderer.tokens, tok)
            yield node.main_token
        }
        .expr_unary:
            try render_expr_unary(renderer, node_id, parse_state, tokens)
        .expr_binary:
            try render_expr_binary(renderer, node_id, parse_state, tokens)
        .literal_int, .literal_float, .literal_string:
            try render_lit(renderer, node_id, parse_state, tokens)
        else: unreachable
    }
}

@hide:file
render_expr_unary proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) TokenId ! mem.AllocError {
    node :: parse_state.nodes.items[usize(node_id)]
    operator :: node.main_token
    operand :: node.data0.node_id

    # render operator operand
    try arraylist.append(&renderer.tokens, tokens[usize(operator)])
    last_token :: try render_expr(renderer, operand, parse_state, tokens)

    return last_token
}

@hide:file
render_expr_binary proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) TokenId ! mem.AllocError {
    node :: parse_state.nodes.items[usize(node_id)]
    operator :: node.main_token
    lhs :: node.data0.node_id
    rhs :: node.data1.node_id

    # render lhs operator rhs
    _ = try render_expr(renderer, lhs, parse_state, tokens)
    try arraylist.append(&renderer.tokens, tokens[usize(operator)])
    last_token :: try render_expr(renderer, rhs, parse_state, tokens)

    return last_token
}

@hide:file
render_type proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) void ! mem.AllocError {
    # todo: support more expressive type expressions...
    node :: parse_state.nodes.items[usize(node_id)]
    match node.kind {
        .expr_identifier: {
            tok :: tokens[usize(node.main_token)]
            try arraylist.append(&renderer.tokens, tok)
        }
        else: unreachable
    }
}

@hide:file
render_lit proc(
    renderer @mut Renderer,
    node_id NodeId,
    parse_state @ParseState,
    tokens []Token,
) TokenId ! mem.AllocError {
    node :: parse_state.nodes.items[usize(node_id)]
    tok :: tokens[usize(node.main_token)]
    try arraylist.append(&renderer.tokens, tok)
    return node.main_token
}