import "@std/debug"
import "@std/mem"
import "@std/hashmap"
import "@source/strpool"

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

BuiltinType :: enum {
    int, float,
    bool, string,
    void, i64,
    f64, usize,
}

@hide:file
Builtin :: struct {
    name []u8
    type BuiltinType
}

@hide:file
builtins [_]Builtin :: [
    { name = "int", type = .int },
    { name = "float", type = .float },
    { name = "bool", type = .bool },
    { name = "string", type = .string },
    { name = "void", type = .void },
    { name = "i64", type = .i64 },
    { name = "f64", type = .f64 },
    { name = "usize", type = .usize },
]

@hide:file
intern_builtins proc(
    $N usize,
    pool @mut strpool.StringPool,
    entries @[N]Builtin,
) [N]mut strpool.StringId ! (mem.AllocError | strpool.InternError) {
    ids [N]mut strpool.StringId := undefined
    for (entries) |builtin, index| ids[index] = try strpool.intern(pool, builtin.name)
    return ids
}

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

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

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

#! links a file-local import to a package root; failed loads retain their diagnostic.
Import :: struct {
    node ast.NodeId
    name strpool.StringId
    target ?ast.NodeId
    diagnostic ?diagnostics.DiagnosticId = null
}

@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.AutoHashMap(strpool.StringId, Binding)
    parent ?@Scope
}

@hide:file
State :: struct {
    tree @ast.Tree
    diagnostics @mut diagnostics.Store
    bindings []mut Binding
    builtin_ids []strpool.StringId
    allocator mem.Allocator
    namespaces hashmap.AutoHashMap(ast.NodeId, Scope)
}

#! borrows `tree`, its identifier `pool`, and `diagnostic_store` during resolution.
#! `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,
    pool @mut strpool.StringPool,
    diagnostic_store @mut diagnostics.Store,
) Resolution ! (
    mem.AllocError
    | strpool.InternError
    | diagnostics.CapacityError
) {
    roots [1]ast.NodeId :: [tree.root]
    return try resolve_packages(
        allocator,
        tree,
        roots[..],
        &[],
        pool,
        diagnostic_store,
    )
}

#! resolves each package root in its own namespace; AST IDs share the compilation arena.
resolve_packages proc(
    allocator mem.Allocator,
    tree @ast.Tree,
    roots []ast.NodeId,
    imports []Import,
    pool @mut strpool.StringPool,
    diagnostic_store @mut diagnostics.Store,
) Resolution ! (mem.AllocError | strpool.InternError | diagnostics.CapacityError) {
    bindings :: try mem.alloc(Binding, allocator, ast.node_count(tree))
    errdefer mem.free(allocator, bindings)
    mem.set(bindings, Binding{})
    builtin_ids :: try intern_builtins(pool, &builtins)

    state := State{
        tree = tree,
        diagnostics = diagnostic_store,
        bindings = bindings,
        builtin_ids = builtin_ids[..],
        allocator = allocator,
        namespaces = hashmap.init(allocator),
    }
    defer {
        for (state.namespaces.entries) |entry| {
            if (entry.hash == 0) continue
            scope := entry.value
            hashmap.deinit(&scope.names)
        }
        hashmap.deinit(&state.namespaces)
    }

    package_scopes :: try mem.alloc(Scope, allocator, roots.len)
    defer mem.free(allocator, package_scopes)
    file_scopes :: try mem.alloc(Scope, allocator, ast.source_files(tree).len)
    defer mem.free(allocator, file_scopes)
    for (file_scopes) |@scope| scope^ = Scope{ names = hashmap.init(allocator), parent = null }
    defer { for (file_scopes) |@scope| hashmap.deinit(&scope.names) }
    for (roots) |root, index| {
        package_scopes[index] = Scope{ names = hashmap.init(allocator), parent = null }
        scope := &package_scopes[index]
        errdefer hashmap.deinit(&scope.names)
        for (0..ast.package_decl_count_at(tree, root)) |i| {
            id :: ast.package_decl_at(tree, root, i)
            file_scopes[usize(ast.node_span(tree, id).source)].parent = scope
            match ast.decl_kind(tree, id) {
                .statement: try predeclare(&state, scope, id, ast.statement_decl(tree, id).name)
                .function: try predeclare(&state, scope, id, ast.function_decl(tree, id).name)
                .import, .invalid: {}
            }
        }
        hashmap.put(&state.namespaces, root, scope^) catch |err| {
            match err {
                .key_exists: unreachable
                else: return err
            }
        }
    }
    for (imports) |imported| try resolve_import(&state, &file_scopes[usize(ast.node_span(tree, imported.node).source)], imported)
    for (roots) |root| try resolve_package(&state, root, file_scopes)

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

@hide:file
resolve_package proc(
    state @mut State, root ast.NodeId, file_scopes []Scope,
) void ! (mem.AllocError | diagnostics.CapacityError) {
    tree := state.tree^
    tree.root = root
    for (0..ast.package_decl_count(&tree)) |i| {
        id :: ast.package_decl(&tree, i)
        scope :: &file_scopes[usize(ast.node_span(&tree, id).source)]
        match ast.decl_kind(&tree, id) {
            .statement: try resolve_statement_decl(state, scope, id)
            .function: try resolve_function(state, scope, id)
            .invalid: resolve_invalid(state, id)
            .import: {}
        }
    }
}

@hide:file
resolve_import proc(
    state @mut State, scope @mut Scope, imported Import,
) void ! (mem.AllocError | diagnostics.CapacityError) {
    span :: ast.import_decl(state.tree, imported.node).span
    value := if (imported.diagnostic) |diagnostic|
        Binding{ kind = .invalid, data = BindingData{ diagnostic = diagnostic } }
    else if (imported.target) |target|
        Binding{ kind = .namespace, data = BindingData{ namespace = target } }
    else
        Binding{ kind = .invalid, data = BindingData{
            diagnostic = try diagnostics.add(state.diagnostics, .package_unavailable, span),
        } }
    if (imported.name == strpool.no_str) {
        state.bindings[usize(imported.node)] = value
        return
    }
    if hashmap.get(&scope.names, imported.name) |_| {
        state.bindings[usize(imported.node)] = Binding{ kind = .invalid, data = BindingData{
            diagnostic = try diagnostics.add(state.diagnostics, .duplicate_name, span),
        } }
        return
    }
    if lookup(state, scope.parent, imported.name) |_| {
        value = Binding{ kind = .invalid, data = BindingData{
            diagnostic = try diagnostics.add(state.diagnostics, .duplicate_name, span),
        } }
    }
    try put_name(scope, imported.name, value)
    state.bindings[usize(imported.node)] = value
}

#! 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(state @State, scope ?@Scope, name strpool.StringId) ?Binding {
    if scope |current| {
        if hashmap.get(&current.names, name) |found| return found
        return lookup(state, current.parent, name)
    }
    for (state.builtin_ids) |id, index| {
        if (id == name) return Binding{
            kind = .builtin,
            data = BindingData{ builtin = builtins[index].type },
        }
    }
    return null
}

@hide:file
put_name proc(scope @mut Scope, name strpool.StringId, 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.name_id(state.tree, id)
    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(state, 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)
    }
    try resolve_block(state, &scope, function.body)
}

@hide:file
resolve_arm proc(state @mut State, parent @Scope, id ast.NodeId) void ! (mem.AllocError | diagnostics.CapacityError) {
    scope := Scope{
        names = hashmap.init(state.allocator),
        parent = parent,
    }
    defer hashmap.deinit(&scope.names)
    try resolve_block(state, &scope, id)
}

@hide:file
resolve_block proc(state @mut State, scope @mut Scope, id ast.NodeId) void ! (mem.AllocError | diagnostics.CapacityError) {
    tree :: state.tree
    for (0..ast.block_statement_count(tree, id)) |i| {
        statement :: ast.block_statement(tree, id, 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, id)) |i| {
        statement :: ast.block_statement(tree, id, 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))
            .return_stmt: {
                if (ast.return_stmt(tree, statement).value) |value| try resolve_expr(state, scope, value)
            }
            .break_stmt, .continue_stmt: {}
            .if_stmt: {
                conditional :: ast.if_stmt(tree, statement)
                try resolve_expr(state, scope, conditional.condition)
                try resolve_arm(state, scope, conditional.then_block)
                if (conditional.else_block) |other| try resolve_arm(state, scope, other)
            }
            .assignment_stmt: {
                assignment :: ast.assignment_stmt(tree, statement)
                try resolve_expr(state, scope, assignment.target)
                try resolve_expr(state, scope, assignment.value)
            }
            .while_stmt: {
                loop :: ast.while_stmt(tree, statement)
                try resolve_expr(state, scope, loop.condition)
                try resolve_arm(state, scope, loop.body)
            }
            .for_stmt: {
                loop :: ast.for_stmt(tree, statement)
                try resolve_expr(state, scope, loop.iterable)
                body_scope := Scope{
                    names = hashmap.init(state.allocator),
                    parent = scope,
                }
                defer hashmap.deinit(&body_scope.names)
                if (loop.capture) |capture| {
                    try predeclare(state, &body_scope, capture, ast.node_span(tree, capture))
                }
                if (loop.index) |index| {
                    try predeclare(state, &body_scope, index, ast.node_span(tree, index))
                }
                try resolve_block(state, &body_scope, loop.body)
            }
            .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_bool, .literal_string: {}
        .identifier: {
            span :: ast.node_span(tree, id)
            name :: ast.name_id(tree, id)
            if lookup(state, 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 },
                }
            }
        }
        .selection: {
            expression :: ast.selection_expr(tree, id)
            try resolve_expr(state, scope, expression.base)
            base :: state.bindings[usize(expression.base)]
            value := base
            if (base.kind != .invalid) {
                if (base.kind == .namespace) {
                    namespace :: hashmap.get(&state.namespaces, base.data.namespace)?
                    if hashmap.get(&namespace.names, ast.name_id(tree, expression.member)) |found| {
                        value = found
                    } else {
                        value = Binding{ kind = .invalid, data = BindingData{
                            diagnostic = try diagnostics.add(state.diagnostics, .unknown_name, ast.node_span(tree, expression.member)),
                        } }
                    }
                } else {
                    value = Binding{ kind = .invalid, data = BindingData{
                        diagnostic = try diagnostics.add(state.diagnostics, .not_namespace, ast.node_span(tree, expression.base)),
                    } }
                }
            }
            state.bindings[usize(id)] = value
            state.bindings[usize(expression.member)] = value
        }
        .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)
        }
        .range: {
            expression :: ast.range_expr(tree, id)
            try resolve_expr(state, scope, expression.start)
            try resolve_expr(state, scope, expression.end)
        }
        .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))
            }
        }
        .array_type: {
            type :: ast.array_type(tree, id)
            if (type.count) |count| try resolve_expr(state, scope, count)
            try resolve_expr(state, scope, type.element)
        }
        .array_literal: {
            for (0..ast.array_element_count(tree, id)) |i| {
                try resolve_expr(state, scope, ast.array_element(tree, id, i))
            }
        }
        .index: {
            expression :: ast.index_expr(tree, id)
            try resolve_expr(state, scope, expression.base)
            try resolve_expr(state, scope, expression.index)
        }
        .invalid: resolve_invalid(state, id)
    }
}