package main

import "core:fmt"
import "core:strconv"

Scope :: struct {
    function_definitions: map[int]^FunctionType, // A map to nodes which are the function definitions
    variable_definitions: map[int]^Type, // A map to types
    variable_mutability_definitions: map[int]bool, // A map to a variable's mutability
    function_return_type: ^Type,
}

@(private = "file")
infer_type :: proc(parent: ^Node, child: ^Node) {
    if child.return_type == nil {
        #partial switch child.kind {
        case .Integer:   child.return_type = type_create_integer(32, true)
        case .Float:     child.return_type = type_create_float(32)
        case .String:    child.return_type = type_create_array(type_create_integer(32, false), 0)
        case .Character: child.return_type = type_create_integer(32, false)
        }
    } else {
        if parent != nil {
            parent.return_type = child.return_type
        }
    }
}

@(private = "file")
is_number :: proc(node: ^Node) -> bool {
    return node.kind == .Integer || node.kind == .Float
}

@(private = "file")
ast_to_type :: proc(node: ^Node) -> ^Type {
    if node == nil {
        return type_create_integer(0, false)
    }
    if node.kind == .Identifier {
        value := node.value.([dynamic]u8)
        if value[0] == 'u' {
            bit_size, ok := strconv.parse_u64_of_base(string(value[1:]), 10)
            if !ok {
                fmt.panicf("Failed to parse integer: %s", value)
            }
            return type_create_integer(u8(bit_size), false)
        } else if value[0] == 'i' {
            bit_size, ok := strconv.parse_u64_of_base(string(value[1:]), 10)
            if !ok {
                fmt.panicf("Failed to parse integer: %s", value)
            }
            return type_create_integer(u8(bit_size), true)
        } else if value[0] == 'f' {
            bit_size, ok := strconv.parse_u64_of_base(string(value[1:]), 10)
            if !ok {
                fmt.panicf("Failed to parse integer: %s", value)
            }
            return type_create_float(u8(bit_size))
        } else {
            fmt.panicf("Unhandled identifier in ast_to_type: %s", value)
        }
    } else if node.kind == .Pointer {
        return type_create_pointer(ast_to_type(node.children[0]))
    } else if node.kind == .Array {
        return type_create_array(ast_to_type(node.children[0]), node.value.(u64))
    } else {
        fmt.panicf("Unhandled node kind in ast_to_type: {}", node.kind)
    }
}

scope_stack := [dynamic]Scope {}

scope_enter :: proc() {
    append(&scope_stack, Scope{})
    scope_stack[len(scope_stack) - 1].function_definitions = make(map[int]^FunctionType)
    scope_stack[len(scope_stack) - 1].variable_definitions = make(map[int]^Type)
    scope_stack[len(scope_stack) - 1].variable_mutability_definitions = make(map[int]bool)
    scope_stack[len(scope_stack) - 1].function_return_type = nil
}

scope_leave :: proc() {
    if len(scope_stack) == 0 {
        fmt.panicf("Tried to leave scope when there are no scopes")
    }
    delete(scope_stack[len(scope_stack) - 1].function_definitions)
    delete(scope_stack[len(scope_stack) - 1].variable_definitions)
    pop(&scope_stack)
}

scope_variable_lookup :: proc(name: [dynamic]u8) -> ^Type {
    name_ := name
    #reverse for &scope in scope_stack {
        type, ok := scope.variable_definitions[get_character_sum_of_dyn_arr(&name_)]
        if ok {
            return type
        }
    }
    return nil
}

scope_variable_lookup_mutable :: proc(name: [dynamic]u8) -> bool {
    name_ := name
    #reverse for &scope in scope_stack {
        type, ok := scope.variable_mutability_definitions[get_character_sum_of_dyn_arr(&name_)]
        if ok {
            return type
        }
    }
    return false
}

scope_function_lookup :: proc(name: [dynamic]u8) -> ^FunctionType {
    name_ := name
    #reverse for &scope in scope_stack {
        type, ok := scope.function_definitions[get_character_sum_of_dyn_arr(&name_)]
        if ok {
            return type
        }
    }
    return nil
}

scope_function_return_type_lookup :: proc() -> ^Type {
    #reverse for &scope in scope_stack {
        if scope.function_return_type != nil {
            return scope.function_return_type
        }
    }
    return nil
}

type_check_function_call :: proc(ast: ^Node, parent_ast: ^Node, must_be_function := true) -> ^FunctionType {
    name : [dynamic]u8
    if ast.kind == .FunctionCall {
        name = ast.children[0].value.([dynamic]u8)
    } else {
        name = ast.value.([dynamic]u8)
    }
    fn := scope_function_lookup(name)
    if fn == nil {
        if must_be_function {
            append(&g_message_list,
                message_create(
                    .Error,
                    fmt.aprintf("Undefined function: %s", name),
                    ast.range,
                ),
            )
        }
        return nil
    }

    return fn
}

type_check :: proc(ast: ^Node, parent_ast: ^Node) {
    in_extern := false

    if ast == nil {
        return
    }

    #partial switch (ast.kind) {
    case .Integer: fallthrough
    case .Float: fallthrough
    case .String:
        infer_type(parent_ast, ast)
    case .Block:
        scope_enter()
        functions := find_function_definitions(ast)
        for fn, i in functions {
            scope_stack[len(scope_stack) - 1].function_definitions[get_character_sum_of_dyn_arr(&fn.name)] = fn
        }
        for child in ast.children {
            type_check(child, ast)
        }
        scope_leave()
    case .FunctionCall:
        type := scope_variable_lookup(ast.children[0].value.([dynamic]u8))
        if type != nil {
            name := ast.children[0].value.([dynamic]u8)
            free(ast.children[0])
            clear(&ast.children)
            ast.return_type = type
            ast.kind = .Identifier
            ast.value = name
        } else {
            fn := type_check_function_call(ast, parent_ast)
            if fn != nil {
                if len(fn.parameter_types) != len(ast.children) - 1 {
                    append(&g_message_list,
                        message_create(
                            .Error,
                            fmt.aprintf("Function call parameter count mismatch for function `%s`: {} and {}", fn.name, len(fn.parameter_types), len(ast.children) - 1),
                            ast.range,
                        ),
                    )
                    break
                }

                for param, i in fn.parameter_types {
                    type_check(ast.children[i + 1], ast)
                    if !compare_types(param, ast.children[i + 1].return_type) {
                        append(&g_message_list,
                            message_create(
                                .Error,
                                fmt.aprintf("Type mismatch in function call for `%s`: Wanted {}, got {}", fn.name, param, ast.children[i + 1].return_type),
                                ast.range,
                            ),
                        )
                    }
                }

                ast.return_type = fn.return_type
            }
        }
    case .Identifier:
        type := scope_variable_lookup(ast.value.([dynamic]u8))
        if type == nil {
            fn := type_check_function_call(ast, parent_ast, false)
            if fn == nil {
                append(&g_message_list,
                    message_create(
                        .Error,
                        fmt.aprintf("Undefined variable: %s", ast.value.([dynamic]u8)),
                        ast.range,
                    ),
                )
            } else {
                ast.kind = .FunctionCall
                append(&ast.children, node_create_value(.Identifier, ast.range, ast.value))
                ast.value = nil
                ast.return_type = fn.return_type
            }
        }
        ast.return_type = type
    case .BinaryExpression:
        type_check(ast.children[0], ast)
        type_check(ast.children[1], ast)

        if !compare_types(ast.children[0].return_type, ast.children[1].return_type) {
            append(&g_message_list,
                message_create(
                    .Error,
                    fmt.aprintf("Type mismatch: {} and {}", ast.children[0].return_type, ast.children[1].return_type),
                    ast.range,
                ),
            )
        }

        ast.return_type = ast.children[0].return_type

        if ast.value_token_kind == .Assign {
            ast.return_type = nil

            if !scope_variable_lookup_mutable(ast.children[0].value.([dynamic]u8)) {
                append(&g_message_list,
                    message_create(
                        .Error,
                        fmt.aprintf("Variable is not mutable: %s", ast.children[0].value.([dynamic]u8)),
                        ast.range,
                    ),
                )
            }
        } else if ast.value_token_kind == .Equals || ast.value_token_kind == .NotEquals ||
            ast.value_token_kind == .GreaterThan || ast.value_token_kind == .GreaterThanOrEqual ||
            ast.value_token_kind == .LessThan || ast.value_token_kind == .LessThanOrEqual {
            ast.return_type = type_create_integer(1, true)
        }

        // FIXME: Verify that the operation is possible
    case .UnaryExpression:
        // FIXME: Verify that the operation is possible
        type_check(ast.children[0], ast)
    case .Ret:
        function_return_type := scope_function_return_type_lookup()
        if function_return_type == nil {
            append(&g_message_list,
                message_create(
                    .Error,
                    fmt.aprintf("Return statement outside of function"),
                    ast.range,
                ),
            )
        } else {
            type_check(ast.children[0], ast)
            if !compare_types(function_return_type, ast.children[0].return_type) {
                append(&g_message_list,
                    message_create(
                        .Error,
                        fmt.aprintf("Type mismatch: {} and {}", function_return_type, ast.children[0].return_type),
                        ast.range,
                    ),
                )
            }
        }
    case .Cast:
        type_check(ast.children[0], ast)
        type_to := ast_to_type(ast.children[1])
        // FIXME: Check if compatible
        ast.return_type = type_to
    case .BitwiseCast:
        type_check(ast.children[0], ast)
        // FIXME: Check if they are both the same bit size
        ast.return_type = ast_to_type(ast.children[1])
    case .VariableDeclaration:
        if ast.children[2] != nil {
            type_check(ast.children[2], ast)
            if ast.children[1] == nil {
                ast.return_type = ast.children[2].return_type
            }
            if !compare_types(ast.return_type, ast.children[2].return_type) {
                append(&g_message_list,
                    message_create(
                        .Error,
                        fmt.aprintf("Type mismatch: {} and {}", ast.return_type, ast.children[2].return_type),
                        ast.range,
                    ),
                )
            }
        } else {
            ast.return_type = ast_to_type(ast.children[1])
        }
        scope_stack[len(scope_stack) - 1].variable_definitions[get_character_sum_of_dyn_arr(&ast.children[0].value.([dynamic]u8))] = ast.return_type
        scope_stack[len(scope_stack) - 1].variable_mutability_definitions[get_character_sum_of_dyn_arr(&ast.children[0].value.([dynamic]u8))] = !ast.value.(bool)
    case .If:
        type_check(ast.children[0], ast)
        if ast.children[0].return_type == nil || ast.children[0].return_type.kind != .Integer {
            append(&g_message_list,
                message_create(
                    .Error,
                    fmt.aprintf("If condition must be a signed/unsigned integer"),
                    ast.children[0].range,
                ),
            )
            break
        }
        type_check(ast.children[1], ast)
        if len(ast.children) == 3 {
            type_check(ast.children[2], ast)
        }
    case .ExternFunction:
        in_extern = true
        fallthrough
    case .Function:
        scope_enter()
        ast.return_type = ast_to_type(ast.children[0])
        scope_stack[len(scope_stack) - 1].function_return_type = ast.return_type
        for child, i in ast.children {
            if in_extern == false {
                if i < 2 {
                    continue
                }
            } else {
                if i < 1 {
                    continue
                }
            }
            type_check(child, ast)
            scope_stack[len(scope_stack) - 1].variable_definitions[get_character_sum_of_dyn_arr(&child.children[0].value.([dynamic]u8))] = child.return_type
            scope_stack[len(scope_stack) - 1].variable_mutability_definitions[get_character_sum_of_dyn_arr(&child.children[0].value.([dynamic]u8))] = true
        }
        if in_extern == false {
            type_check(ast.children[1], ast)
        }
        scope_leave()
    case .For:
        scope_enter()
        for child, i in ast.children {
            if child == nil {
                continue
            }
            if i == 1 {
                type_check(child, ast)
                should_error := false
                if child.return_type == nil {
                    should_error = true
                } else if child.return_type.kind != .Integer {
                    should_error = true
                }
                if should_error {
                    append(&g_message_list,
                        message_create(
                            .Error,
                            fmt.aprintf("For condition must be a signed/unsigned integer"),
                            child.range,
                        ),
                    )
                    break
                }
            } else {
                type_check(child, ast)
            }
        }
        scope_leave()
    case:
        fmt.panicf("Unhandled node kind in type_check: {}", ast.kind)
    }
}

find_function_definitions :: proc(ast_: ^Node) -> (ret: [dynamic]^FunctionType) {
    if ast_.kind != .Block {
        return
    }
    for ast in ast_.children {
        if ast == nil {
            continue
        }
        is_extern := false
        #partial switch (ast.kind) {
        case .ExternFunction:
            is_extern = true
            fallthrough
        case .Function:
            for fn in ret {
                if compare_dyn_arrs(&fn.name, &ast.value.([dynamic]u8)) {
                    append(&g_message_list,
                        message_create(
                            .Error,
                            fmt.aprintf("Function already defined: {}", ast.value.([dynamic]u8)),
                            ast.range,
                        ),
                    )
                    continue
                }
            }
            fn := function_type_create()
            fn.name = ast.value.([dynamic]u8)
            return_type : ^Type
            if ast.children[0] == nil {
                return_type = type_create_integer(0, false)
            } else {
                return_type = ast_to_type(ast.children[0])
            }
            fn.return_type = return_type
            for decl, i in ast.children {
                if is_extern == false {
                    if i < 2 {
                        continue
                    }
                } else {
                    if i < 1 {
                        continue
                    }

                }
                type := ast_to_type(decl.children[1])
                append(&fn.parameter_types, type)
            }
            append(&ret, fn)
        case:
        }
    }
    return
}
