import "@std/mem"
import "@std/testing"

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

parses_expression_precedence_through_tree_views test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source :: "value :: 1 + 2 * 3\n"
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)

    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    try testing.expect_equal(1, ast.package_decl_count(&tree))
    declaration :: ast.package_decl(&tree, 0)
    statement :: ast.statement_decl(&tree, declaration)
    try testing.expect_equal("value", ast.source_text(&tree, statement.name))
    try testing.expect_equal(ast.ExprKind.binary, ast.expr_kind(&tree, statement.initializer))
    add :: ast.binary_expr(&tree, statement.initializer)
    try testing.expect_equal(ast.BinaryOperator.add, add.operator)
    try testing.expect_equal(ast.ExprKind.literal_int, ast.expr_kind(&tree, add.left))
    try testing.expect_equal(ast.ExprKind.binary, ast.expr_kind(&tree, add.right))
    multiply :: ast.binary_expr(&tree, add.right)
    try testing.expect_equal(ast.BinaryOperator.multiply, multiply.operator)
    try testing.expect_equal(0, diagnostics.count(&diagnostic_store))
}

parses_grouped_expression_precedence test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source :: "value :: a + ((b - c)) * 2 / 3\n"
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)

    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    try testing.expect_equal(0, diagnostics.count(&diagnostic_store))
    try testing.expect_equal(1, ast.package_decl_count(&tree))
    statement :: ast.statement_decl(&tree, ast.package_decl(&tree, 0))
    add :: ast.binary_expr(&tree, statement.initializer)
    try testing.expect_equal(ast.BinaryOperator.add, add.operator)
    try testing.expect_equal("a", ast.source_text(&tree, ast.node_span(&tree, add.left)))
    divide :: ast.binary_expr(&tree, add.right)
    try testing.expect_equal(ast.BinaryOperator.divide, divide.operator)
    try testing.expect_equal("3", ast.source_text(&tree, ast.node_span(&tree, divide.right)))
    multiply :: ast.binary_expr(&tree, divide.left)
    try testing.expect_equal(ast.BinaryOperator.multiply, multiply.operator)
    try testing.expect_equal("2", ast.source_text(&tree, ast.node_span(&tree, multiply.right)))
    subtract :: ast.binary_expr(&tree, multiply.left)
    try testing.expect_equal(ast.BinaryOperator.subtract, subtract.operator)
    try testing.expect_equal("b", ast.source_text(&tree, ast.node_span(&tree, subtract.left)))
    try testing.expect_equal("c", ast.source_text(&tree, ast.node_span(&tree, subtract.right)))
}

parses_function_structure_through_tree_views test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source ::
        `grouped proc(a, b T, c U) void {
        `    value int :: 1
        `}
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)

    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    declaration :: ast.package_decl(&tree, 0)
    try testing.expect_equal(ast.DeclKind.function, ast.decl_kind(&tree, declaration))
    function :: ast.function_decl(&tree, declaration)
    try testing.expect_equal("grouped", ast.source_text(&tree, function.name))
    try testing.expect_equal(3, ast.function_param_count(&tree, declaration))
    first :: ast.param(&tree, ast.function_param(&tree, declaration, 0))
    second :: ast.param(&tree, ast.function_param(&tree, declaration, 1))
    third :: ast.param(&tree, ast.function_param(&tree, declaration, 2))
    try testing.expect_equal("a", ast.source_text(&tree, first.name))
    try testing.expect_equal("b", ast.source_text(&tree, second.name))
    try testing.expect_equal("c", ast.source_text(&tree, third.name))
    try testing.expect_equal(first.type, second.type)
    try testing.expect(first.type != third.type)
    try testing.expect_equal(1, ast.block_statement_count(&tree, function.body))
    try testing.expect_equal(0, diagnostics.count(&diagnostic_store))
}

recovers_invalid_declaration_and_preserves_sibling test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source :: "@ broken\nvalid :: 1\n"
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)

    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    try testing.expect_equal(2, ast.package_decl_count(&tree))
    invalid_decl :: ast.package_decl(&tree, 0)
    valid_decl :: ast.package_decl(&tree, 1)
    try testing.expect_equal(ast.DeclKind.invalid, ast.decl_kind(&tree, invalid_decl))
    try testing.expect_equal(ast.DeclKind.statement, ast.decl_kind(&tree, valid_decl))
    invalid_view :: ast.invalid(&tree, invalid_decl)
    try testing.expect(diagnostics.get(&diagnostic_store, invalid_view.diagnostic) != null)
    valid_statement :: ast.statement_decl(&tree, valid_decl)
    try testing.expect_equal("valid", ast.source_text(&tree, valid_statement.name))
}

recovers_malformed_groups_and_preserves_following_statements test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source ::
        `main proc() void {
        `    empty :: ()
        `    operand :: (1 + )
        `    missing :: (2
        `    valid :: (3)
        `}
        `after :: 4
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)

    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    try testing.expect_equal(3, diagnostics.count(&diagnostic_store))
    try testing.expect_equal(2, ast.package_decl_count(&tree))
    function :: ast.function_decl(&tree, ast.package_decl(&tree, 0))
    try testing.expect_equal(4, ast.block_statement_count(&tree, function.body))
    empty :: ast.statement_decl(&tree, ast.block_statement(&tree, function.body, 0))
    try testing.expect_equal(ast.ExprKind.invalid, ast.expr_kind(&tree, empty.initializer))
    operand :: ast.statement_decl(&tree, ast.block_statement(&tree, function.body, 1))
    add :: ast.binary_expr(&tree, operand.initializer)
    try testing.expect_equal(ast.ExprKind.invalid, ast.expr_kind(&tree, add.right))
    missing :: ast.statement_decl(&tree, ast.block_statement(&tree, function.body, 2))
    try testing.expect_equal(ast.ExprKind.invalid, ast.expr_kind(&tree, missing.initializer))
    valid :: ast.statement_decl(&tree, ast.block_statement(&tree, function.body, 3))
    try testing.expect_equal(ast.ExprKind.literal_int, ast.expr_kind(&tree, valid.initializer))
    try testing.expect_equal("3", ast.source_text(&tree, ast.node_span(&tree, valid.initializer)))
    after :: ast.statement_decl(&tree, ast.package_decl(&tree, 1))
    try testing.expect_equal("after", ast.source_text(&tree, after.name))
    try testing.expect_equal("4", ast.source_text(&tree, ast.node_span(&tree, after.initializer)))
}

parses_call_statements_and_postfix_precedence test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source ::
        `main proc() void {
        `    ping()
        `    result int := -factory()(add(1, 2 * 3), (4 + 5),) + 6
        `    (handler)(result)
        `}
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)
    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    try testing.expect_equal(0, diagnostics.count(&diagnostic_store))
    function :: ast.function_decl(&tree, ast.package_decl(&tree, 0))
    try testing.expect_equal(3, ast.block_statement_count(&tree, function.body))
    ping_stmt :: ast.block_statement(&tree, function.body, 0)
    try testing.expect_equal(ast.StmtKind.expression, ast.stmt_kind(&tree, ping_stmt))
    ping :: ast.statement_expr(&tree, ping_stmt)
    try testing.expect_equal(ast.ExprKind.call, ast.expr_kind(&tree, ping))
    try testing.expect_equal(0, ast.call_arg_count(&tree, ping))
    try testing.expect_equal("ping", ast.source_text(&tree, ast.node_span(&tree, ast.call_expr(&tree, ping).callee)))
    try testing.expect_equal("ping()", ast.source_text(&tree, ast.call_expr(&tree, ping).span))

    result :: ast.statement_decl(&tree, ast.block_statement(&tree, function.body, 1))
    try testing.expect_equal(ast.Mutability.variable, result.mutability)
    add :: ast.binary_expr(&tree, result.initializer)
    try testing.expect_equal(ast.BinaryOperator.add, add.operator)
    try testing.expect_equal("6", ast.source_text(&tree, ast.node_span(&tree, add.right)))
    negate :: ast.unary_expr(&tree, add.left)
    call :: negate.operand
    try testing.expect_equal(2, ast.call_arg_count(&tree, call))
    factory :: ast.call_expr(&tree, call).callee
    try testing.expect_equal(0, ast.call_arg_count(&tree, factory))
    try testing.expect_equal("factory", ast.source_text(&tree, ast.node_span(&tree, ast.call_expr(&tree, factory).callee)))
    nested :: ast.call_arg(&tree, call, 0)
    try testing.expect_equal(2, ast.call_arg_count(&tree, nested))
    try testing.expect_equal("add", ast.source_text(&tree, ast.node_span(&tree, ast.call_expr(&tree, nested).callee)))
    try testing.expect_equal("1", ast.source_text(&tree, ast.node_span(&tree, ast.call_arg(&tree, nested, 0))))
    multiply :: ast.binary_expr(&tree, ast.call_arg(&tree, nested, 1))
    try testing.expect_equal(ast.BinaryOperator.multiply, multiply.operator)
    try testing.expect_equal("2", ast.source_text(&tree, ast.node_span(&tree, multiply.left)))
    try testing.expect_equal("3", ast.source_text(&tree, ast.node_span(&tree, multiply.right)))
    grouped :: ast.binary_expr(&tree, ast.call_arg(&tree, call, 1))
    try testing.expect_equal(ast.BinaryOperator.add, grouped.operator)
    try testing.expect_equal("4", ast.source_text(&tree, ast.node_span(&tree, grouped.left)))
    try testing.expect_equal("5", ast.source_text(&tree, ast.node_span(&tree, grouped.right)))

    handler :: ast.statement_expr(&tree, ast.block_statement(&tree, function.body, 2))
    try testing.expect_equal("handler", ast.source_text(&tree, ast.node_span(&tree, ast.call_expr(&tree, handler).callee)))
    try testing.expect_equal(1, ast.call_arg_count(&tree, handler))
    try testing.expect_equal("result", ast.source_text(&tree, ast.node_span(&tree, ast.call_arg(&tree, handler, 0))))
}

recovers_malformed_calls_and_preserves_siblings test {
    strpool.STRINGS = strpool.init(mem.c_allocator)
    defer strpool.deinit(&strpool.STRINGS)
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    source ::
        `main proc() void {
        `    bad(1 2)
        `    good()
        `    missing(1
        `    after :: 3
        `    holes(, 2)
        `}
        `tail :: f()
    scan_state := lexer.init(mem.c_allocator)
    defer lexer.deinit(&scan_state)
    try lexer.scan(&scan_state, source, &diagnostic_store)
    tree := try parse(scan_state.tokens.items, source, &diagnostic_store, mem.c_allocator)
    defer ast.deinit(&tree)

    try testing.expect_equal(3, diagnostics.count(&diagnostic_store))
    try testing.expect_equal(2, ast.package_decl_count(&tree))
    function :: ast.function_decl(&tree, ast.package_decl(&tree, 0))
    try testing.expect_equal(5, ast.block_statement_count(&tree, function.body))
    for ([0, 2]) |index| {
        statement :: ast.block_statement(&tree, function.body, usize(index))
        try testing.expect_equal(ast.StmtKind.invalid, ast.stmt_kind(&tree, statement))
        try testing.expect(diagnostics.get(&diagnostic_store, ast.invalid(&tree, statement).diagnostic) != null)
    }
    good :: ast.statement_expr(&tree, ast.block_statement(&tree, function.body, 1))
    try testing.expect_equal("good", ast.source_text(&tree, ast.node_span(&tree, ast.call_expr(&tree, good).callee)))
    after :: ast.statement_decl(&tree, ast.block_statement(&tree, function.body, 3))
    try testing.expect_equal("after", ast.source_text(&tree, after.name))
    holes :: ast.statement_expr(&tree, ast.block_statement(&tree, function.body, 4))
    try testing.expect_equal(2, ast.call_arg_count(&tree, holes))
    try testing.expect_equal(ast.ExprKind.invalid, ast.expr_kind(&tree, ast.call_arg(&tree, holes, 0)))
    try testing.expect_equal("2", ast.source_text(&tree, ast.node_span(&tree, ast.call_arg(&tree, holes, 1))))
    tail :: ast.statement_decl(&tree, ast.package_decl(&tree, 1))
    try testing.expect_equal("tail", ast.source_text(&tree, tail.name))
    try testing.expect_equal(0, ast.call_arg_count(&tree, tail.initializer))
}