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

import "@source/diagnostics"
import "@source/lexer"
import "@source/ast"

@hide:file
SyntaxError :: enum { unexpected_token }

@hide:file
ParseError :: alias SyntaxError | mem.AllocError | diagnostics.CapacityError

@hide:file
BindingPowerScalar :: distinct u32

@hide:file
BindingPower :: struct {
    left BindingPowerScalar
    right BindingPowerScalar
}

@hide:file
State :: struct {
    builder ast.Builder
    diagnostics @mut diagnostics.Store
    tokens []lexer.Token
    allocator mem.Allocator
    next_token usize
    paren_depth usize = 0
}

#! borrows `tokens` and `source` while parsing; the returned Tree borrows `source` until ast.deinit.
#! `allocator` supplies returned node storage and temporary lists released before return.
#! the caller owns the Tree and must release it with ast.deinit.
parse proc(
    allocator mem.Allocator,
    tokens []lexer.Token,
    source []u8,
    diagnostic_store @mut diagnostics.Store,
) ast.Tree ! (mem.AllocError | diagnostics.CapacityError) {
    state := State{
        builder = ast.init_builder(allocator, source),
        diagnostics = diagnostic_store,
        tokens = tokens,
        allocator = allocator,
        next_token = 0,
    }
    defer ast.deinit_builder(&state.builder)
    root :: try parse_package(&state)
    return ast.finish(&state.builder, root)
}

@hide:file
parse_package proc(state @mut State) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    declarations arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&declarations)

    skip_newlines(state)
    while !next_is(state, .eof) {
        start :: state.next_token
        declaration :: parse_decl(state) catch |err| {
            match err {
                .unexpected_token: {
                    recovered :: try recover_decl(state, start)
                    try arraylist.append(&declarations, recovered)
                    skip_newlines(state)
                    continue
                }
                else: return err
            }
        }
        try arraylist.append(&declarations, declaration)
        skip_newlines(state)
    }

    return try ast.add_package(
        &state.builder,
        ast.Span{ start = 0, end = u32(state.builder.source.len) },
        declarations.items,
    )
}

@hide:file
parse_decl proc(state @mut State) ast.NodeId ! ParseError {
    if (next_is(state, .ident) and next_n_is(state, 1, .proc)) return try parse_func_decl(state)
    return try parse_statement_decl(state)
}

@hide:file
parse_func_decl proc(state @mut State) ast.NodeId ! ParseError {
    name :: try consume(state, .ident)
    _ = try consume(state, .proc)
    _ = try consume(state, .open_paren)

    params arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&params)
    names arraylist.ArrayList(ast.Span) := arraylist.init(state.allocator)
    defer arraylist.deinit(&names)

    while !next_is(state, .close_paren) {
        arraylist.clear(&names)
        while true {
            token :: try consume(state, .ident)
            try arraylist.append(&names, token.span)
            if next_is(state, .ident) break
            _ = try consume(state, .comma)
        }

        type_id :: try parse_type(state)
        type_span :: ast.builder_span(&state.builder, type_id)
        for (names.items) |name_span| {
            param_id :: try ast.add_param(
                &state.builder,
                ast.Span{ start = name_span.start, end = type_span.end },
                name_span,
                type_id,
            )
            try arraylist.append(&params, param_id)
        }
        if next_is(state, .comma) _ = try consume(state, .comma)
    }

    _ = try consume(state, .close_paren)
    return_type :: try parse_type(state)
    body :: try parse_block(state)
    body_span :: ast.builder_span(&state.builder, body)
    return try ast.add_function_decl(
        &state.builder,
        ast.Span{ start = name.span.start, end = body_span.end },
        name.span,
        return_type,
        params.items,
        body,
    )
}

@hide:file
parse_block proc(state @mut State) ast.NodeId ! ParseError {
    open :: try consume(state, .open_curly)
    statements arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&statements)

    skip_newlines(state)
    while !next_is(state, .close_curly) and !next_is(state, .eof) {
        start :: state.next_token
        statement :: parse_statement(state, false) catch |err| {
            match err {
                .unexpected_token: {
                    recovered :: try recover_statement(state, start)
                    try arraylist.append(&statements, recovered)
                    skip_newlines(state)
                    continue
                }
                else: return err
            }
        }
        try arraylist.append(&statements, statement)
        skip_newlines(state)
    }

    close :: try consume(state, .close_curly)
    return try ast.add_block(
        &state.builder,
        ast.Span{ start = open.span.start, end = close.span.end },
        statements.items,
    )
}

@hide:file
parse_statement proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    if (next_is(state, .ident) and next_n_is(state, 1, .equal)) {
        return try parse_assignment(state, allow_else)
    }
    if (
        next_is(state, .ident)
        and (
            next_n_is(state, 1, .ident)
            or next_n_is(state, 1, .double_colon)
            or next_n_is(state, 1, .colon_equal)
        )
    ) {
        return try parse_statement_decl(state)
    }
    match current(state).kind {
        .if: return try parse_if(state)
        .while: return try parse_while(state, allow_else)
        .return: return try parse_return(state, allow_else)
        .break, .continue: return try parse_loop_control(state, allow_else)
        .ident, .int, .float, .string, .true, .false, .minus, .bang, .open_paren: {
            expression :: try parse_expr(state)
            return try ast.add_expression_stmt(&state.builder, expression)
        }
        else: return .unexpected_token
    }
}

@hide:file
parse_return proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: try consume(state, .return)
    if (
        next_is(state, .newline) or next_is(state, .close_curly) or next_is(state, .eof)
        or (allow_else and next_is(state, .else))
    ) {
        return try ast.add_return_stmt(&state.builder, keyword.span, null)
    }
    value :: try parse_expr(state)
    if (
        !next_is(state, .newline) and !next_is(state, .close_curly) and !next_is(state, .eof)
        and !(allow_else and next_is(state, .else))
    ) {
        return .unexpected_token
    }
    # include consumed grouping punctuation, which has no separate expression node.
    end :: state.tokens[state.next_token - 1].span.end
    return try ast.add_return_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = end },
        value,
    )
}

@hide:file
parse_loop_control proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: current(state)
    state.next_token += 1
    if (
        !next_is(state, .newline)
        and !next_is(state, .close_curly)
        and !next_is(state, .eof)
        and !(allow_else and next_is(state, .else))
    ) return .unexpected_token

    if (keyword.kind == .break)
        return try ast.add_break_stmt(&state.builder, keyword.span)

    return try ast.add_continue_stmt(&state.builder, keyword.span)
}

@hide:file
parse_assignment proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    name :: try consume(state, .ident)
    target :: try ast.add_identifier(&state.builder, name.span)
    _ = try consume(state, .equal)
    value :: try parse_expr(state)
    if (
        !next_is(state, .newline)
        and !next_is(state, .close_curly)
        and !next_is(state, .eof)
        and !(allow_else and next_is(state, .else))
    ) {
        return .unexpected_token
    }
    end :: state.tokens[state.next_token - 1].span.end
    return try ast.add_assignment_stmt(
        &state.builder,
        ast.Span{ start = name.span.start, end = end },
        target,
        value,
    )
}

@hide:file
parse_if proc(state @mut State) ast.NodeId ! ParseError {
    keyword :: try consume(state, .if)
    condition :: try parse_guard(state)
    braced :: next_is(state, .open_curly)
    then_block :: try parse_arm(state, true)
    # only a braced arm permits its else on a later line.
    after_then :: state.next_token
    if (braced) skip_newlines(state)
    alternative ?ast.NodeId := null
    if next_is(state, .else) {
        state.next_token += 1
        alternative = try parse_arm(state, true)
    } else {
        state.next_token = after_then
    }
    end :: if (alternative) |block|
        ast.builder_span(&state.builder, block).end
    else
        ast.builder_span(&state.builder, then_block).end
    return try ast.add_if_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = end },
        condition,
        then_block,
        alternative,
    )
}

@hide:file
parse_while proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    keyword :: try consume(state, .while)
    condition :: try parse_guard(state)
    body :: try parse_arm(state, allow_else)
    return try ast.add_while_stmt(
        &state.builder,
        ast.Span{ start = keyword.span.start, end = ast.builder_span(&state.builder, body).end },
        condition,
        body,
    )
}

@hide:file
parse_guard proc(state @mut State) ast.NodeId ! ParseError {
    condition_start :: state.next_token
    condition := parse_expr(state) catch |err| {
        match err {
            .unexpected_token: {}
            else: return err
        }
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        yield try ast.add_invalid_expr(&state.builder, current(state).span, diagnostic)
    }
    if (ast.builder_expr_kind(&state.builder, condition) == .invalid) {
        # discard a malformed guard's tail, but leave its body and enclosing block intact.
        while (
            !next_is(state, .open_curly) and !next_is(state, .close_curly)
            and !next_is(state, .newline) and !next_is(state, .eof)
            and !next_is(state, .else) and !next_is(state, .return)
            and !next_is(state, .if) and !next_is(state, .while)
            and !next_is(state, .break) and !next_is(state, .continue)
        ) {
            if (next_is(state, .ident) and next_n_is(state, 1, .equal)) break
            state.next_token += 1
        }
    }
    if (
        !next_is(state, .open_curly)
        and !parenthesized_guard(state, condition_start)
        and ast.builder_expr_kind(&state.builder, condition) != .call
        and ast.builder_expr_kind(&state.builder, condition) != .invalid
    ) {
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        condition = try ast.add_invalid_expr(
            &state.builder,
            ast.builder_span(&state.builder, condition),
            diagnostic,
        )
    }
    return condition
}

@hide:file
parenthesized_guard proc(state @State, start usize) bool {
    if (state.tokens[start].kind != .open_paren) return false
    depth usize := 0
    for start..state.next_token |index| {
        kind :: state.tokens[index].kind
        if (kind == .open_paren) depth += 1
        if (kind == .close_paren) {
            depth -= 1
            if (depth == 0) return index + 1 == state.next_token
        }
    }
    return false
}

@hide:file
parse_arm proc(state @mut State, allow_else bool) ast.NodeId ! ParseError {
    if next_is(state, .open_curly) return try parse_block(state)
    start :: state.next_token
    statement :: parse_statement(state, allow_else) catch |err| {
        match err {
            .unexpected_token: {}
            else: return err
        }
        diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
        span :: synchronize_arm(state, start)
        yield try ast.add_invalid_stmt(&state.builder, span, diagnostic)
    }
    statements [1]ast.NodeId :: [statement]
    return try ast.add_block(
        &state.builder,
        ast.builder_span(&state.builder, statement),
        statements[..],
    )
}

@hide:file
synchronize_arm proc(state @mut State, start usize) ast.Span {
    first :: state.tokens[start].span
    end := first.end
    depth usize := 0
    while !next_is(state, .eof) {
        if (
            depth == 0
            and (next_is(state, .newline) or next_is(state, .close_curly) or next_is(state, .else))
        ) break
        if next_is(state, .open_curly) depth += 1
        if (next_is(state, .close_curly) and depth > 0) depth -= 1
        end = current(state).span.end
        state.next_token += 1
    }
    return ast.Span{ start = first.start, end = end }
}

@hide:file
parse_statement_decl proc(state @mut State) ast.NodeId ! ParseError {
    name :: try consume(state, .ident)
    type_id ?ast.NodeId :: if next_is(state, .ident)
        try parse_type(state)
    else
        null

    assignment :: try consume_either(state, &[.double_colon, .colon_equal])
    initializer :: try parse_expr(state)
    end :: ast.builder_span(&state.builder, initializer).end
    mutability :: if (assignment.kind == .double_colon) .constant else .variable
    return try ast.add_statement_decl(
        &state.builder,
        ast.Span{ start = name.span.start, end = end },
        name.span,
        type_id,
        mutability,
        initializer,
    )
}

@hide:file
parse_type proc(state @mut State) ast.NodeId ! ParseError {
    if (!next_is(state, .ident)) return .unexpected_token
    token :: current(state)
    state.next_token += 1
    return try ast.add_identifier(&state.builder, token.span)
}

@hide:file
parse_expr proc(state @mut State) ast.NodeId ! ParseError {
    return try parse_expr_bp(state, BindingPowerScalar(0))
}

@hide:file
parse_expr_bp proc(state @mut State, min_bp BindingPowerScalar) ast.NodeId ! ParseError {
    skip_operand_newlines(state)
    token := current(state)
    left := if (prefix_binding_power(token.kind)) |right_bp| {
        state.next_token += 1
        operand :: try parse_expr_bp(state, right_bp)
        operand_span :: ast.builder_span(&state.builder, operand)
        operator ast.UnaryOperator :: if (token.kind == .bang) .logical_not else .negate
        yield try ast.add_unary(
            &state.builder,
            ast.Span{ start = token.span.start, end = operand_span.end },
            operator,
            operand,
        )
    } else {
        yield try parse_primary(state)
    }

    while true {
        skip_continuation_newlines(state)
        token = current(state)
        if (token.kind == .open_paren) {
            left = try parse_call(state, left)
            continue
        }
        if (infix_binding_power(token.kind)) |power| {
            if (power.left < min_bp) break
            state.next_token += 1
            right :: try parse_expr_bp(state, power.right)
            left_span :: ast.builder_span(&state.builder, left)
            right_span :: ast.builder_span(&state.builder, right)
            left = try ast.add_binary(
                &state.builder,
                ast.Span{ start = left_span.start, end = right_span.end },
                binary_operator(token.kind),
                left,
                right,
            )
            continue
        }
        break
    }
    return left
}

@hide:file
parse_call proc(state @mut State, callee ast.NodeId) ast.NodeId ! ParseError {
    _ = try consume(state, .open_paren)
    state.paren_depth += 1
    defer state.paren_depth -= 1
    skip_operand_newlines(state)
    arguments arraylist.ArrayList(ast.NodeId) := arraylist.init(state.allocator)
    defer arraylist.deinit(&arguments)

    while (!is_call_boundary(current(state).kind)) {
        argument :: try parse_expr(state)
        try arraylist.append(&arguments, argument)
        if (!next_is(state, .comma)) break
        state.next_token += 1
        skip_operand_newlines(state)
    }
    close :: try consume(state, .close_paren)
    return try ast.add_call(
        &state.builder,
        ast.Span{
            start = ast.builder_span(&state.builder, callee).start,
            end = close.span.end,
        },
        callee,
        arguments.items,
    )
}

@hide:file
parse_primary proc(state @mut State) ast.NodeId ! ParseError {
    token :: current(state)
    match token.kind {
        .open_paren: {
            state.next_token += 1
            state.paren_depth += 1
            defer state.paren_depth -= 1
            expression :: try parse_expr(state)
            if next_is(state, .close_paren) {
                state.next_token += 1
                return expression
            }
            diagnostic :: try diagnostics.add(
                state.diagnostics,
                .unexpected_token,
                current(state).span,
            )
            return try ast.add_invalid_expr(&state.builder, ast.Span{
                start = token.span.start,
                end = ast.builder_span(&state.builder, expression).end,
            }, diagnostic)
        }
        .int: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .int, token.span)
        }
        .float: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .float, token.span)
        }
        .string: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .string, token.span)
        }
        .true, .false: {
            state.next_token += 1
            return try ast.add_literal(&state.builder, .bool, token.span)
        }
        .ident: {
            state.next_token += 1
            return try ast.add_identifier(&state.builder, token.span)
        }
        else: {
            diagnostic :: if (token.data.diagnostic != diagnostics.no_diagnostic)
                token.data.diagnostic
            else
                try diagnostics.add(state.diagnostics, .unexpected_token, token.span)

            if !is_expression_recovery_boundary(token.kind) {
                state.next_token += 1
            }

            return try ast.add_invalid_expr(&state.builder, token.span, diagnostic)
        }
    }
}

@hide:file
is_expression_boundary proc(kind lexer.TokenKind) bool {
    return match kind {
        .close_paren, .newline, .close_curly, .open_curly, .else, .return, .if, .while, .eof: true
        else: false
    }
}

@hide:file
is_call_boundary proc(kind lexer.TokenKind) bool {
    return (
        kind == .break
        or kind == .continue
        or is_expression_boundary(kind)
    )
}

@hide:file
is_expression_recovery_boundary proc(kind lexer.TokenKind) bool {
    return (
        kind == .comma
        or kind == .and
        or kind == .or
        or is_expression_boundary(kind)
    )
}

@hide:file
skip_operand_newlines proc(state @mut State) void {
    if (state.paren_depth == 0 or !next_is(state, .newline)) return
    start :: state.next_token
    skip_newlines(state)
    kind :: current(state).kind
    if (kind != .close_paren and is_expression_boundary(kind)) state.next_token = start
    # a missing operand must not consume the declaration or assignment on the next line.
    if (
        next_is(state, .ident)
        and (
            next_n_is(state, 1, .double_colon)
            or next_n_is(state, 1, .colon_equal)
            or next_n_is(state, 1, .equal)
            or next_n_is(state, 1, .ident)
            or next_n_is(state, 1, .proc)
        )
    ) state.next_token = start
}

@hide:file
skip_continuation_newlines proc(state @mut State) void {
    if (state.paren_depth == 0 or !next_is(state, .newline)) return
    start :: state.next_token
    skip_newlines(state)
    kind :: current(state).kind
    if (
        kind != .close_paren and kind != .comma and kind != .open_paren
        and infix_binding_power(kind) == null
    ) state.next_token = start
}

@hide:file
recover_decl proc(
    state @mut State,
    start usize,
) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
    span :: synchronize(state, start, false)
    return try ast.add_invalid_decl(&state.builder, span, diagnostic)
}

@hide:file
recover_statement proc(
    state @mut State,
    start usize,
) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    diagnostic :: try diagnostics.add(state.diagnostics, .unexpected_token, current(state).span)
    span :: synchronize(state, start, true)
    return try ast.add_invalid_stmt(&state.builder, span, diagnostic)
}

@hide:file
synchronize proc(state @mut State, start usize, stop_at_curly bool) ast.Span {
    first :: state.tokens[start].span
    end := first.end
    depth usize := 0
    while !next_is(state, .eof) {
        if (depth == 0 and next_is(state, .newline)) break
        if (stop_at_curly and depth == 0 and next_is(state, .close_curly)) break
        if (stop_at_curly and next_is(state, .open_curly)) depth += 1
        if (stop_at_curly and next_is(state, .close_curly) and depth > 0) depth -= 1
        end = current(state).span.end
        state.next_token += 1
    }
    if next_is(state, .newline) {
        end = current(state).span.end
        state.next_token += 1
    }
    return ast.Span{ start = first.start, end = end }
}

@hide:file
skip_newlines proc(state @mut State) void {
    while next_is(state, .newline) state.next_token += 1
}

@hide:file
current proc(state @State) lexer.Token {
    if (state.next_token >= state.tokens.len) return state.tokens[state.tokens.len - 1]
    return state.tokens[state.next_token]
}

@hide:file
next_is proc(state @State, kind lexer.TokenKind) bool {
    return current(state).kind == kind
}

@hide:file
next_n_is proc(state @State, count usize, kind lexer.TokenKind) bool {
    index :: state.next_token + count
    return index < state.tokens.len and state.tokens[index].kind == kind
}

@hide:file
consume proc(state @mut State, kind lexer.TokenKind) lexer.Token ! SyntaxError {
    token :: current(state)
    if (token.kind != kind) return .unexpected_token
    state.next_token += 1
    return token
}

@hide:file
consume_either proc(state @mut State, kinds []lexer.TokenKind) lexer.Token ! SyntaxError {
    token :: current(state)
    for (kinds) |kind| if token.kind == kind {
        state.next_token += 1
        return token
    }
    return .unexpected_token
}

@hide:file
prefix_binding_power proc(kind lexer.TokenKind) ?BindingPowerScalar {
    return match kind {
        .minus, .bang: BindingPowerScalar(50)
        else: null
    }
}

@hide:file
infix_binding_power proc(kind lexer.TokenKind) ?BindingPower {
    return match kind {
        .or: BindingPower{ left = BindingPowerScalar(2), right = BindingPowerScalar(3) }
        .and: BindingPower{ left = BindingPowerScalar(4), right = BindingPowerScalar(5) }
        .equal_equal, .not_equal: BindingPower{ left = BindingPowerScalar(10), right = BindingPowerScalar(11) }
        .less, .less_equal, .greater, .greater_equal: BindingPower{ left = BindingPowerScalar(20), right = BindingPowerScalar(21) }
        .plus, .minus: BindingPower{ left = BindingPowerScalar(30), right = BindingPowerScalar(31) }
        .star, .slash: BindingPower{ left = BindingPowerScalar(40), right = BindingPowerScalar(41) }
        else: null
    }
}

@hide:file
binary_operator proc(kind lexer.TokenKind) ast.BinaryOperator {
    return match kind {
        .and: .logical_and
        .or: .logical_or
        .plus: .add
        .minus: .subtract
        .star: .multiply
        .slash: .divide
        .equal_equal: .equal
        .not_equal: .not_equal
        .less: .less
        .less_equal: .less_equal
        .greater: .greater
        .greater_equal: .greater_equal
        else: unreachable
    }
}