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

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 Token :: alias lexer.Token
hide TokenId :: alias lexer.TokenId
hide TokenKind :: alias lexer.TokenKind
hide Node :: alias ast.Node
hide NodeId :: alias ast.NodeId
hide ExtraId :: alias ast.ExtraId
hide NodeData :: alias ast.NodeData

hide NO_ID_NODE :: alias ast.NO_ID_NODE
hide NO_ID_EXTRA :: alias ast.NO_ID_EXTRA

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

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

    tokens []Token = &[]
    next_token TokenId = TokenId(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) void ! ParseError {
    state.tokens = tokens
    _ = try parse_decl(state)
}

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

    # check for type
    # todo: make a parse_type proc for this
    type_id :: if (state.tokens[state.next_token].kind == .ident)
        try parse_primary(state)
    else
        NO_ID_NODE

    # expect `::` (immutable assignment)
    mutability_tok_id :: try consume(state, .double_colon)

    # expect expression
    expr_id :: try parse_primary(state)

    # expect statement terminator (newline)
    _ = try consume_either(state, &[.newline, .eof])

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

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

    return decl
}

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

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)
    state.next_token += 1
    return state.next_token - 1
}

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

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

hide node_id proc(idx uint) NodeId {
    debug.assert(u64(idx) < u64(NO_ID_NODE))
    return NodeId(idx)
}

hide add_extra proc(extra @mut std.ArrayList(u32), idx u32) ExtraId ! mem.AllocError {
    id :: extra_id(extra.items.len)
    try arraylist.append(extra, idx)
    return id
}

hide extra_id proc(idx uint) ExtraId {
    debug.assert(u64(idx) < u64(NO_ID_EXTRA))
    return ExtraId(idx)
}