import "@std"
import "@std/mem"
import "@std/arraylist"
import "@std/enums/enummap"
import "@std/debug"

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

# todo: move error stuff into a separate file

Diagnostic :: struct {
    code ErrorCode
    node NodeId
}

ParseError :: alias ErrorCode | mem.AllocError

ErrorCode :: enum {
    unexpected_token,
}

ErrorDetails :: struct {
    name []u8
    message []u8
}

error_msg_map std.EnumMap(ErrorCode, ErrorDetails) :: enummap.init({
    unexpected_token = ErrorDetails {
        name = "P0",
        message = "unexpected token",
    },
})

@hide:file
Token :: alias lexer.Token

@hide:file
TokenId :: alias lexer.TokenId

@hide:file
TokenKind :: alias lexer.TokenKind

@hide:file
Node :: alias ast.Node

@hide:file
NodeId :: alias ast.NodeId

@hide:file
ExtraId :: alias ast.ExtraId

@hide:file
NodeData :: alias ast.NodeData

@hide:file
NO_NODE :: alias ast.NO_NODE

@hide:file
NO_EXTRA :: alias ast.NO_EXTRA

@hide:file
token_id :: alias lexer.token_id

BindingPowerScalar :: distinct u32

BindingPower :: struct {
    left BindingPowerScalar
    right BindingPowerScalar
}

prefix_binding_power proc(token_kind TokenKind) ?BindingPowerScalar {
    return match token_kind {
        .minus: BindingPowerScalar(30)
        else: null
    }
}

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

postfix_binding_power proc(token_kind TokenKind) ?BindingPowerScalar {
    return match token_kind {
        .open_paren, .open_bracket: BindingPowerScalar(40)
        else: null
    }
}

State :: struct {
    nodes std.ArrayList(Node)

    #! array of raw u32's indexing nodes, tokens, etc.
    extra std.ArrayList(u32)

    tokens []Token = &[]
    next_token usize = 0
}

init proc(allocator mem.Allocator) State {
    return State{
        nodes = arraylist.init(allocator),
        extra = arraylist.init(allocator),
    }
}

deinit proc(state @mut State) void {
    arraylist.deinit(&state.nodes)
    arraylist.deinit(&state.extra)
}

parse proc(state @mut State, tokens []Token) NodeId ! ParseError {
    state.tokens = tokens
    return try parse_pkg(state)
}

parse_pkg proc(state @mut State) NodeId ! ParseError {
    if (state.next_token < state.tokens.len
        and state.tokens[state.next_token].kind == .eof
    ) return NO_NODE

    stmt_ids std.ArrayList(NodeId) := arraylist.init(state.nodes.allocator)
    defer arraylist.deinit(&stmt_ids)

    while (state.next_token < state.tokens.len
        and state.tokens[state.next_token].kind != .eof
    ) {
        stmt_id :: try parse_stmt(state)
        try arraylist.append(&stmt_ids, stmt_id)
    }

    pkg_start :: ast.extra_id(state.extra.items.len)
    for (stmt_ids.items) |stmt_id| _ = try add_extra(state, u32(stmt_id))
    pkg_end :: ast.extra_id(state.extra.items.len)

    package :: try add_node(state, Node{
        kind = .package,
        data0 = NodeData{ extra_id = pkg_start },
        data1 = NodeData{ extra_id = pkg_end },
    })

    return package
}

parse_stmt proc(state @mut State) NodeId ! ParseError {
    return try parse_decl(state)
}

parse_decl proc(state @mut State) NodeId ! ParseError {
    # expect identifier
    ident_token :: try consume(state, .ident)

    # check for type
    # todo: make a parse_type proc for this
    type_id :: if (next_is(state, .ident))
        try parse_primary(state)
    else
        NO_NODE

    # expect `::` or `:=` (immutable or mutable assignment)
    mutability_tok_id :: try consume_either(state, &[.double_colon, .colon_equal])

    # expect expression
    expr_id :: try parse_expr(state)

    # expect statement terminator (newline)
    while (next_is(state, .newline)) _ = try consume(state, .newline)

    extra_start :: try add_extra(state, u32(type_id))
    _ = try add_extra(state, u32(expr_id))

    decl :: try add_node(state, Node{
        kind = .stmt_decl,
        main_token = ident_token,
        data0 = NodeData{ extra_id = extra_start },
        data1 = NodeData{ token_id = mutability_tok_id },
    })

    return decl
}

parse_expr proc(state @mut State) NodeId ! mem.AllocError {
    return try parse_expr_bp(state, BindingPowerScalar(0))
}

parse_expr_bp proc(state @mut State, min_bp BindingPowerScalar) NodeId ! mem.AllocError {
    # parse left-hand side
    kind := state.tokens[state.next_token].kind

    lhs := if (prefix_binding_power(kind)) |right_bp| {
        operator :: token_id(state.next_token)
        state.next_token += 1

        operand :: try parse_expr_bp(state, right_bp)

        yield try add_node(state, Node{
            kind = .expr_unary,
            main_token = operator,
            data0 = NodeData{ node_id = operand },
        })
    } else {
        yield try parse_primary(state)
    }

    # parse right-hand side
    while true {
        kind = state.tokens[state.next_token].kind

        # try postfix (higher precedence)
        if postfix_binding_power(kind) |left_bp| {
            if (left_bp < min_bp) break

            lhs = match state.tokens[state.next_token].kind {
                .open_paren: debug.unimplemented() # todo: parse call
                .open_bracket: debug.unimplemented() # todo: parse index
                else: unreachable
            }
            continue
        }

        # try infex (lower precedence)
        if infix_binding_power(kind) |bp| {
            if (bp.left < min_bp) break

            operator :: token_id(state.next_token)
            state.next_token += 1

            rhs :: try parse_expr_bp(state, bp.right)
            lhs = try add_node(state, Node{
                kind = .expr_binary,
                main_token = operator,
                data0 = NodeData{ node_id = lhs },
                data1 = NodeData{ node_id = rhs },
            })

            continue
        }

        break
    }

    return lhs
}

parse_primary proc(state @mut State) NodeId ! mem.AllocError {
    start_idx :: state.next_token
    start_token :: TokenId(start_idx)
    state.next_token += 1
    return match state.tokens[start_idx].kind {
        .int: try add_node(state, Node{
            kind = .literal_int,
            main_token = start_token,
        })
        .float: try add_node(state, Node{
            kind = .literal_float,
            main_token = start_token,
        })
        .string: try add_node(state, Node{
            kind = .literal_string,
            main_token = start_token,
        })
        .ident: try add_node(state, Node{
            kind = .expr_identifier,
            main_token = start_token,
        })
        else: try add_node(state, Node{
            kind = .invalid,
            main_token = start_token,
        })
    }
}

@hide
next_is proc(state @mut State, token_kind TokenKind) bool {
    return state.tokens[state.next_token].kind == token_kind
}

@hide
next_is_either proc(state @mut State, token_kinds []TokenKind) bool {
    is_either := false
    for (token_kinds) |kind| is_either = is_either or next_is(state, kind)
    return is_either
}

@hide
expect proc(state @mut State, token_kind TokenKind) void ! ErrorCode {
    if (state.tokens[state.next_token].kind != token_kind) return .unexpected_token
}

@hide
expect_either proc(state @mut State, token_kinds []TokenKind) void ! ErrorCode {
    for (token_kinds) |kind| if (state.tokens[state.next_token].kind == kind) return
    return .unexpected_token
}

@hide
consume proc(state @mut State, token_kind TokenKind) TokenId ! ErrorCode {
    try expect(state, token_kind)
    id :: TokenId(state.next_token)
    state.next_token += 1
    return id
}

@hide
consume_either proc(state @mut State, token_kinds []TokenKind) TokenId ! ErrorCode {
    try expect_either(state, token_kinds)
    id :: TokenId(state.next_token)
    state.next_token += 1
    return id
}

@hide
add_node proc(state @mut State, node Node) NodeId ! mem.AllocError {
    id :: ast.node_id(state.nodes.items.len)
    try arraylist.append(&state.nodes, node)
    return id
}

@hide
add_extra proc(state @mut State, data u32) ExtraId ! mem.AllocError {
    id :: ast.extra_id(state.extra.items.len)
    try arraylist.append(&state.extra, data)
    return id
}