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
}

parse proc(
    tokens []lexer.Token,
    source []u8,
    diagnostic_store @mut diagnostics.Store,
    allocator mem.Allocator,
) 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) 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) ast.NodeId ! ParseError {
    if next_is(state, .ident) return try parse_statement_decl(state)
    return .unexpected_token
}

@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 ! (mem.AllocError | diagnostics.CapacityError) {
    return try parse_expr_bp(state, BindingPowerScalar(0))
}

@hide:file
parse_expr_bp proc(state @mut State, min_bp BindingPowerScalar) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    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)
        yield try ast.add_unary(
            &state.builder,
            ast.Span{ start = token.span.start, end = operand_span.end },
            .negate,
            operand,
        )
    } else {
        yield try parse_primary(state)
    }

    while true {
        token = current(state)
        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_primary proc(state @mut State) ast.NodeId ! (mem.AllocError | diagnostics.CapacityError) {
    token :: current(state)
    match token.kind {
        .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)
        }
        .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 (token.kind != .newline and token.kind != .close_curly and token.kind != .eof) {
                state.next_token += 1
            }
            return try ast.add_invalid_expr(&state.builder, token.span, diagnostic)
        }
    }
}

@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
    while !next_is(state, .eof) and !next_is(state, .newline) {
        if (stop_at_curly and next_is(state, .close_curly)) break
        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: BindingPowerScalar(30)
        else: null
    }
}

@hide:file
infix_binding_power proc(kind lexer.TokenKind) ?BindingPower {
    return match kind {
        .plus, .minus: BindingPower{ left = BindingPowerScalar(10), right = BindingPowerScalar(11) }
        .star, .slash: BindingPower{ left = BindingPowerScalar(20), right = BindingPowerScalar(21) }
        else: null
    }
}

@hide:file
binary_operator proc(kind lexer.TokenKind) ast.BinaryOperator {
    return match kind {
        .plus: .add
        .minus: .subtract
        .star: .multiply
        .slash: .divide
        else: unreachable
    }
}