import "@std/debug"
import "@std/mem"
import "@std/hashmap"
import "@std/strmap"

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

BuiltinType :: enum { int, float, string, void }

@hide:file
builtins strmap.StaticStringMap(BuiltinType) :: strmap.init([
    { "int", .int },
    { "float", .float },
    { "string", .string },
    { "void", .void },
])

BindingKind :: enum { none, declaration, builtin, invalid }

BindingData :: union {
    declaration ast.NodeId
    builtin BuiltinType
    diagnostic diagnostics.DiagnosticId
}

Binding :: struct {
    kind BindingKind = .none
    data BindingData = BindingData{ diagnostic = diagnostics.NO_DIAGNOSTIC }
}

@hide:file
Storage :: struct {
    bindings []mut Binding
    allocator mem.Allocator
}

#! owns binding storage and retains its allocator; release with deinit.
#! node and diagnostic IDs refer to the caller's Tree and diagnostics.Store.
Resolution :: struct { storage Storage }

@hide:file
Scope :: struct {
    names hashmap.StringHashMap(Binding)
    parent ?@Scope
}

@hide:file
State :: struct {
    tree @ast.Tree
    diagnostics @mut diagnostics.Store
    bindings []mut Binding
    allocator mem.Allocator
}

#! borrows `tree` and `diagnostic_store` during resolution, without retaining their storage.
#! `allocator` supplies caller-owned bindings and temporary scopes released before return.
#! release the returned Resolution with deinit, which uses the retained allocator.
resolve proc(
    allocator mem.Allocator,
    tree @ast.Tree,
    diagnostic_store @mut diagnostics.Store,
) Resolution ! (mem.AllocError | diagnostics.CapacityError) {
    bindings :: try mem.alloc(Binding, allocator, ast.node_count(tree))
    errdefer mem.free(allocator, bindings)
    mem.set(bindings, Binding{})

    state := State{
        tree = tree,
        diagnostics = diagnostic_store,
        bindings = bindings,
        allocator = allocator,
    }

    package_scope := Scope{ names = hashmap.init(allocator), parent = null }
    defer hashmap.deinit(&package_scope.names)
    for (0..ast.package_decl_count(tree)) |i| {
        id :: ast.package_decl(tree, i)
        match ast.decl_kind(tree, id) {
            .statement: try predeclare(&state, &package_scope, id, ast.statement_decl(tree, id).name)
            .function: try predeclare(&state, &package_scope, id, ast.function_decl(tree, id).name)
            .invalid: {}
        }
    }

    for (0..ast.package_decl_count(tree)) |i| {
        id :: ast.package_decl(tree, i)
        match ast.decl_kind(tree, id) {
            .statement: try resolve_statement_decl(&state, &package_scope, id)
            .function: try resolve_function(&state, &package_scope, id)
            .invalid: resolve_invalid(&state, id)
        }
    }

    return Resolution{ storage = Storage{ bindings = bindings, allocator = allocator } }
}

#! releases binding storage with the allocator retained by resolve.
deinit proc(resolution @mut Resolution) void {
    mem.free(resolution.storage.allocator, resolution.storage.bindings)
}

binding proc(resolution @Resolution, id ast.NodeId) Binding {
    debug.assert(usize(id) < resolution.storage.bindings.len)
    return resolution.storage.bindings[usize(id)]
}

@hide:file
lookup proc(scope ?@Scope, name []u8) ?Binding {
    if scope |current| {
        if hashmap.get(&current.names, name) |found| return found
        return lookup(current.parent, name)
    }
    if strmap.get(&builtins, name) |type| {
        return Binding{ kind = .builtin, data = BindingData{ builtin = type } }
    }
    return null
}

@hide:file
put_name proc(scope @mut Scope, name []u8, value Binding) void ! mem.AllocError {
    hashmap.put(&scope.names, name, value) catch |err| {
        match err {
            .key_exists: unreachable
            else: return err
        }
    }
}

@hide:file
predeclare proc(
    state @mut State,
    scope @mut Scope,
    id ast.NodeId,
    name_span ast.Span,
) void ! (mem.AllocError | diagnostics.CapacityError) {
    name :: ast.source_text(state.tree, name_span)
    if hashmap.get(&scope.names, name) |_| {
        diagnostic :: try diagnostics.add(state.diagnostics, .duplicate_name, name_span)
        state.bindings[usize(id)] = Binding{
            kind = .invalid,
            data = BindingData{ diagnostic = diagnostic },
        }
        return
    }

    value := Binding{ kind = .declaration, data = BindingData{ declaration = id } }
    if lookup(scope.parent, name) |_| {
        diagnostic :: try diagnostics.add(state.diagnostics, .duplicate_name, name_span)
        # poison the inner name rather than silently binding uses to its ancestor.
        value = Binding{ kind = .invalid, data = BindingData{ diagnostic = diagnostic } }
    }
    try put_name(scope, name, value)
    state.bindings[usize(id)] = value
}

@hide:file
resolve_function proc(
    state @mut State,
    package_scope @Scope,
    id ast.NodeId,
) void ! (mem.AllocError | diagnostics.CapacityError) {
    tree :: state.tree
    function :: ast.function_decl(tree, id)
    # signatures cannot see parameters or body-local declarations.
    try resolve_expr(state, package_scope, function.return_type)
    for (0..ast.function_param_count(tree, id)) |i| {
        parameter :: ast.function_param(tree, id, i)
        try resolve_expr(state, package_scope, ast.param(tree, parameter).type)
    }

    scope := Scope{ names = hashmap.init(state.allocator), parent = package_scope }
    defer hashmap.deinit(&scope.names)
    for (0..ast.function_param_count(tree, id)) |i| {
        parameter :: ast.function_param(tree, id, i)
        try predeclare(state, &scope, parameter, ast.param(tree, parameter).name)
    }
    for (0..ast.block_statement_count(tree, function.body)) |i| {
        statement :: ast.block_statement(tree, function.body, i)
        if ast.stmt_kind(tree, statement) == .declaration {
            try predeclare(state, &scope, statement, ast.statement_decl(tree, statement).name)
        }
    }

    for (0..ast.block_statement_count(tree, function.body)) |i| {
        statement :: ast.block_statement(tree, function.body, i)
        match ast.stmt_kind(tree, statement) {
            .declaration: try resolve_statement_decl(state, &scope, statement)
            .expression: try resolve_expr(state, &scope, ast.statement_expr(tree, statement))
            .invalid: resolve_invalid(state, statement)
        }
    }
}

@hide:file
resolve_statement_decl proc(
    state @mut State,
    scope @Scope,
    id ast.NodeId,
) void ! (mem.AllocError | diagnostics.CapacityError) {
    declaration :: ast.statement_decl(state.tree, id)
    if (declaration.type) |type| try resolve_expr(state, scope, type)
    try resolve_expr(state, scope, declaration.initializer)
}

@hide:file
resolve_invalid proc(state @mut State, id ast.NodeId) void {
    if (state.bindings[usize(id)].kind != .none) return
    state.bindings[usize(id)] = Binding{
        kind = .invalid,
        data = BindingData{ diagnostic = ast.invalid(state.tree, id).diagnostic },
    }
}

@hide:file
resolve_expr proc(
    state @mut State,
    scope @Scope,
    id ast.NodeId,
) void ! (mem.AllocError | diagnostics.CapacityError) {
    # grouped parameters share their type node, including any diagnostic it owns.
    if (state.bindings[usize(id)].kind != .none) return
    tree :: state.tree
    match ast.expr_kind(tree, id) {
        .literal_int, .literal_float, .literal_string: {}
        .identifier: {
            span :: ast.node_span(tree, id)
            name :: ast.source_text(tree, span)
            if lookup(scope, name) |found| {
                state.bindings[usize(id)] = found
            } else {
                diagnostic :: try diagnostics.add(state.diagnostics, .unknown_name, span)
                state.bindings[usize(id)] = Binding{
                    kind = .invalid,
                    data = BindingData{ diagnostic = diagnostic },
                }
            }
        }
        .unary: try resolve_expr(state, scope, ast.unary_expr(tree, id).operand)
        .binary: {
            expression :: ast.binary_expr(tree, id)
            try resolve_expr(state, scope, expression.left)
            try resolve_expr(state, scope, expression.right)
        }
        .call: {
            try resolve_expr(state, scope, ast.call_expr(tree, id).callee)
            for (0..ast.call_arg_count(tree, id)) |i| {
                try resolve_expr(state, scope, ast.call_arg(tree, id, i))
            }
        }
        .invalid: resolve_invalid(state, id)
    }
}