package main

import "core:fmt"
import "core:strings"

LLVMScope :: struct {
    name: cstring,
    basic_block: LLVMBasicBlockRef,
    definitions: map[int]LLVMValueRef,
    types: map[int]LLVMTypeRef,
}

g_llvm_scope_stack := [dynamic]LLVMScope{}

llvm_scope_find_definition :: proc(name: ^[dynamic]u8) -> LLVMValueRef {
    #reverse for &scope in g_llvm_scope_stack {
        value, ok := scope.definitions[get_character_sum_of_dyn_arr(name)]
        if ok {
            return value
        }
    }
    return nil
}

llvm_scope_find_type :: proc(name: ^[dynamic]u8) -> LLVMTypeRef {
    #reverse for &scope in g_llvm_scope_stack {
        value, ok := scope.types[get_character_sum_of_dyn_arr(name)]
        if ok {
            return value
        }
    }
    return nil
}

llvm_scope_enter :: proc(name: cstring, basic_block: LLVMBasicBlockRef) {
    append(&g_llvm_scope_stack, LLVMScope{
        name = name,
        basic_block = basic_block,
        definitions = make(map[int]LLVMValueRef),
        types = make(map[int]LLVMTypeRef),
    })
}

llvm_scope_leave :: proc() {
    delete(g_llvm_scope_stack[len(g_llvm_scope_stack)-1].definitions)
    delete(g_llvm_scope_stack[len(g_llvm_scope_stack)-1].types)
    pop(&g_llvm_scope_stack)
}

llvm_top_scope :: proc() -> ^LLVMScope {
    return &g_llvm_scope_stack[len(g_llvm_scope_stack)-1]
}

generate_llvm :: proc(ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, node: ^Node) {
    // FIXME: Everything leaks memory.
    scope_number : i32 = 0

	int_32_type := LLVMInt32TypeInContext(ctx)
	main_function_type := LLVMFunctionType(int_32_type, nil, 0, LLVMBool(0))
	main_function := LLVMAddFunction(mod, "main", main_function_type)

    get_basic_block_name :: proc(scope_num: i32) -> cstring {
        return fmt.caprintf("scope_%d", scope_num)
    }

    generate_llvm_integer :: proc(ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, node: ^Node) -> LLVMValueRef {
        bit_size := uint(node.return_type.bit_size)
        value := node.value.(u64)
        is_signed := node.return_type.is_signed
        return LLVMConstInt(LLVMIntTypeInContext(ctx, bit_size), value, LLVMBool(is_signed))
    }

    generate_llvm_float :: proc(ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, node: ^Node) -> LLVMValueRef {
        bit_size := uint(node.return_type.bit_size)
        value := node.value.(f64)
        if bit_size == 64 {
            return LLVMConstReal(LLVMDoubleTypeInContext(ctx), value)
        } else if bit_size == 32 {
            return LLVMConstReal(LLVMFloatTypeInContext(ctx), value)
        }

        append(
            &g_message_list,
            message_create(
                .Error,
                fmt.aprintf("Unsupported float bit size: %d", bit_size),
                node.range,
            ),
        )
        return nil
    }

    generate_llvm_type_from_node :: proc(ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, type: ^Type) -> LLVMTypeRef {
        if type == nil {
            return LLVMVoidTypeInContext(ctx)
        }

        switch type.kind {
        case .Integer:
            return LLVMIntTypeInContext(ctx, uint(type.bit_size))
        case .Float:
            if type.bit_size == 64 {
                return LLVMDoubleTypeInContext(ctx)
            } else if type.bit_size == 32 {
                return LLVMFloatTypeInContext(ctx)
            } else {
                fmt.panicf("LLVM-IR: Unsupported float type bit size: {}", type.bit_size)
            }
        case .Array:
            array_of := generate_llvm_type_from_node(ctx, mod, builder, type.array_of)
            return LLVMArrayType2(array_of, type.array_size)
        case .Pointer:
            pointer_of := generate_llvm_type_from_node(ctx, mod, builder, type.pointer_to)
            return LLVMPointerType(pointer_of, 0)
        }
        panic("LLVM-IR: Invalid type")
    }

    generate_llvm_function_extern :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef,
        fn: ^Node, scope_number: ^i32
    ) -> (type: LLVMTypeRef, value: LLVMValueRef) {
        function_args_type := [dynamic]LLVMTypeRef{}
        for arg in fn.children {
            if arg.kind != .VariableDeclaration {
                continue
            }
            append(&function_args_type, generate_llvm_type_from_node(ctx, mod, builder, arg.return_type))
        }
        function_return_type := generate_llvm_type_from_node(ctx, mod, builder, fn.return_type)

        type = LLVMFunctionType(function_return_type, raw_data(function_args_type[:]), len(function_args_type), LLVMBool(0))
        value = LLVMAddFunction(mod, strings.clone_to_cstring(string(fn.value.([dynamic]u8)[:])), type)
        return
    }

    generate_llvm_function_call :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef,
        node: ^Node,
    ) -> LLVMValueRef {
        // FIXME: Add support for dot access
        name := &node.children[0].value.([dynamic]u8)
        fn_type := llvm_scope_find_type(name)
        fn_value := llvm_scope_find_definition(name)
        if fn_type == nil || fn_value == nil {
            append(
                &g_message_list,
                message_create(
                    .Error,
                    fmt.aprintf("Function '%s' not found", name),
                    node.range,
                ),
            )
            return nil
        }

        fn_args := [dynamic]LLVMValueRef{}
        for &arg in node.children[1:] {
            value := generate_llvm_expression(ctx, mod, builder, arg)
            append(&fn_args, value)
        }

        call := LLVMBuildCall2(builder, fn_type, fn_value, raw_data(fn_args[:]), len(fn_args), cstring(raw_data(name[:])))
        return call
    }

    generate_llvm_value :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef,
        node: ^Node,
    ) -> LLVMValueRef {
        if node.kind == .Integer {
            return generate_llvm_integer(ctx, mod, builder, node)
        } else if node.kind == .Float {
            return generate_llvm_float(ctx, mod, builder, node)
        } else if node.kind == .FunctionCall {
            return generate_llvm_function_call(ctx, mod, builder, node)
        } else if node.kind == .Identifier {
            def := llvm_scope_find_definition(&node.value.([dynamic]u8))
            type := llvm_scope_find_type(&node.value.([dynamic]u8))
            //LLVMBuildLoad2 :: proc(Builder: LLVMBuilderRef, Ty: LLVMTypeRef, PointerVal: LLVMValueRef, Name: cstring) -> LLVMValueRef ---
            def_value := LLVMBuildLoad2(builder, type, def, "loadtmp")
            return def_value
        }

        fmt.panicf("FIXME: Implement other node kinds. Got: {}", node.kind)
    }

    generate_llvm_expression :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef,
        node: ^Node,
    ) -> LLVMValueRef {
        if node.kind == .BinaryExpression {
            return generate_llvm_binary_expression(ctx, mod, builder, node)
        }
        if node.kind == .Integer || node.kind == .Float || node.kind == .FunctionCall || node.kind == .Identifier {
            return generate_llvm_value(ctx, mod, builder, node)
        }
        if node.kind == .FunctionCall {
            return generate_llvm_function_call(ctx, mod, builder, node)
        }
        return nil
    }

    generate_llvm_binary_expression :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef,
        node: ^Node,
    ) -> LLVMValueRef {
        op := node.value_token_kind
        lhs_node := node.children[0]
        rhs_node := node.children[1]
        lhs := generate_llvm_expression(ctx, mod, builder, lhs_node)
        rhs := generate_llvm_expression(ctx, mod, builder, rhs_node)
        if lhs_node.return_type.kind != rhs_node.return_type.kind {
            fmt.panicf("LLVM-C: Binary expression types do not match: {} and {}", lhs_node.return_type.kind, rhs_node.return_type.kind)
        }
        is_float := lhs_node.return_type.kind == .Float
        is_signed := lhs_node.return_type.is_signed
        #partial switch op {
        case .Add:
            if is_float {
                return LLVMBuildFAdd(builder, lhs, rhs, "addtmp")
            } else {
                return LLVMBuildAdd(builder, lhs, rhs, "addtmp")
            }
        case .Subtract:
            if is_float {
                return LLVMBuildFSub(builder, lhs, rhs, "subtmp")
            } else {
                return LLVMBuildSub(builder, lhs, rhs, "subtmp")
            }
        case .Multiply:
            if is_float {
                return LLVMBuildFMul(builder, lhs, rhs, "multmp")
            } else {
                return LLVMBuildMul(builder, lhs, rhs, "multmp")
            }
        case .Divide:
            if is_float {
                return LLVMBuildFDiv(builder, lhs, rhs, "divtmp")
            } else {
                if is_signed {
                    return LLVMBuildSDiv(builder, lhs, rhs, "divtmp")
                } else {
                    return LLVMBuildUDiv(builder, lhs, rhs, "divtmp")
                }
            }
        case .Modulo:
            if is_float {
                return LLVMBuildFRem(builder, lhs, rhs, "modtmp")
            } else {
                if is_signed {
                    return LLVMBuildSRem(builder, lhs, rhs, "modtmp")
                } else {
                    return LLVMBuildURem(builder, lhs, rhs, "modtmp")
                }
            }
        case .BitwiseLeftShift:
            return LLVMBuildShl(builder, lhs, rhs, "shltmp")
        case .BitwiseRightShift:
            if is_signed {
                return LLVMBuildAShr(builder, lhs, rhs, "ashrtmp")
            } else {
                return LLVMBuildLShr(builder, lhs, rhs, "lshrtmp")
            }
        case .BitwiseAnd:
            return LLVMBuildAnd(builder, lhs, rhs, "andtmp")
        case .BitwiseOr:
            return LLVMBuildOr(builder, lhs, rhs, "ortmp")
        case .BitwiseXOR:
            return LLVMBuildXor(builder, lhs, rhs, "xortmp")
        case .Equals:
            if is_float {
                return LLVMBuildFCmp(builder, .LLVMRealOEQ, lhs, rhs, "eqtmp")
            } else {
                return LLVMBuildICmp(builder, .LLVMIntEQ, lhs, rhs, "eqtmp")
            }
        case .NotEquals:
            if is_float {
                return LLVMBuildFCmp(builder, .LLVMRealONE, lhs, rhs, "neqtmp")
            } else {
                return LLVMBuildICmp(builder, .LLVMIntNE, lhs, rhs, "neqtmp")
            }
        case .LessThan:
            if is_float {
                return LLVMBuildFCmp(builder, .LLVMRealOLT, lhs, rhs, "lttmp")
            } else {
                if is_signed {
                    return LLVMBuildICmp(builder, .LLVMIntSLT, lhs, rhs, "lttmp")
                } else {
                    return LLVMBuildICmp(builder, .LLVMIntULT, lhs, rhs, "lttmp")
                }
            }
        case .LessThanOrEqual:
            if is_float {
                return LLVMBuildFCmp(builder, .LLVMRealOLE, lhs, rhs, "letmp")
            } else {
                if is_signed {
                    return LLVMBuildICmp(builder, .LLVMIntSLE, lhs, rhs, "letmp")
                } else {
                    return LLVMBuildICmp(builder, .LLVMIntULE, lhs, rhs, "letmp")
                }
            }
        case .GreaterThan:
            if is_float {
                return LLVMBuildFCmp(builder, .LLVMRealOGT, lhs, rhs, "gttmp")
            } else {
                if is_signed {
                    return LLVMBuildICmp(builder, .LLVMIntSGT, lhs, rhs, "gttmp")
                } else {
                    return LLVMBuildICmp(builder, .LLVMIntUGT, lhs, rhs, "gttmp")
                }
            }
        case .GreaterThanOrEqual:
            if is_float {
                return LLVMBuildFCmp(builder, .LLVMRealOGE, lhs, rhs, "getmp")
            } else {
                if is_signed {
                    return LLVMBuildICmp(builder, .LLVMIntSGE, lhs, rhs, "getmp")
                } else {
                    return LLVMBuildICmp(builder, .LLVMIntUGE, lhs, rhs, "getmp")
                }
            }
        case .Assign:
            if node.children[0].kind != .Identifier {
                fmt.panicf("LLVM-C: Left hand side of assignment is not an identifier")
            }
            name := &node.children[0].value.([dynamic]u8)
            def := llvm_scope_find_definition(name)
            if def == nil {
                append(
                    &g_message_list,
                    message_create(
                        .Error,
                        fmt.aprintf("Variable '%s' not found", name),
                        node.range,
                    ),
                )
                return nil
            }
            return LLVMBuildStore(builder, rhs, def)
        case:
            fmt.panicf("LLVM-C: Unexpected binary expression kind: {}", op)
        }
    }

    generate_llvm_if :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, function: LLVMValueRef,
        node: ^Node, scope_number: ^i32
    ) {
        condition_node := node.children[0]
        true_node := node.children[1]
        false_node := node.children[2]

        bb := LLVMGetLastBasicBlock(function)

        condition := generate_llvm_expression(ctx, mod, builder, condition_node)
        if condition == nil {
            return
        }
        int_32_type : LLVMTypeRef
        if condition_node.return_type.kind == .Float {
            if condition_node.return_type.bit_size == 64 {
                int_32_type = LLVMDoubleTypeInContext(ctx)
            } else {
                int_32_type = LLVMFloatTypeInContext(ctx)
            }
        } else {
            int_32_type = LLVMIntTypeInContext(ctx, uint(condition_node.return_type.bit_size))
        }
        comparison_res := LLVMBuildICmp(builder, .LLVMIntNE, condition, LLVMConstInt(int_32_type, 0, LLVMBool(1)), "")

        true_block := LLVMAppendBasicBlockInContext(ctx, function, "")
        false_block := LLVMCreateBasicBlockInContext(ctx, "")
        end_block := LLVMCreateBasicBlockInContext(ctx, "")

        cond_br := LLVMBuildCondBr(builder, comparison_res, true_block, false_block)

        LLVMPositionBuilderAtEnd(builder, true_block)
        generate_llvm_scope(ctx, mod, builder, function, true_node, scope_number, true_block)
        LLVMBuildBr(builder, end_block)
        true_block = LLVMGetInsertBlock(builder)

        LLVMAppendExistingBasicBlock(function, false_block)

        LLVMPositionBuilderAtEnd(builder, false_block)
        generate_llvm_scope(ctx, mod, builder, function, false_node, scope_number, false_block)
        LLVMBuildBr(builder, end_block)
        false_block = LLVMGetInsertBlock(builder)

        LLVMPositionBuilderAtEnd(builder, end_block)
        LLVMAppendExistingBasicBlock(function, end_block)
    }

    generate_llvm_for :: proc(
        ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, function: LLVMValueRef,
        node: ^Node, scope_number: ^i32
    ) {
        sum_nils := 0
        for i in node.children {
            if i == nil {
                sum_nils += 1
            }
        }

        LoopType :: enum {
            Infinite,
            While,
        }

        loop_type : LoopType
        switch sum_nils {
        case 3: loop_type = LoopType.Infinite
        case 2: loop_type = LoopType.While
        case:
            panic("FIXME: Implement other for loop types")
        }

        if loop_type == LoopType.Infinite {
            basic_block := LLVMAppendBasicBlockInContext(ctx, function, "for")
            LLVMBuildBr(builder, basic_block)
            LLVMPositionBuilderAtEnd(builder, basic_block)
            scope_number^ += 1
            llvm_scope_enter("for_body", basic_block)
            generate_llvm_scope(ctx, mod, builder, function, node.children[3], scope_number, basic_block)
            llvm_scope_leave()
            LLVMBuildBr(builder, basic_block)
        } else if loop_type == LoopType.While {
            fmt.println("In while")
            condition_block := LLVMAppendBasicBlockInContext(ctx, function, "while_condition")
            body_block := LLVMAppendBasicBlockInContext(ctx, function, "while_body")
            end_block := LLVMAppendBasicBlockInContext(ctx, function, "while_end")

            LLVMBuildBr(builder, condition_block)

            LLVMPositionBuilderAtEnd(builder, condition_block)
            cond := generate_llvm_expression(ctx, mod, builder, node.children[1])
            LLVMBuildCondBr(builder, cond, body_block, end_block)

            LLVMPositionBuilderAtEnd(builder, body_block)
            llvm_scope_enter("for_body", body_block)
            scope_number^ += 1
            generate_llvm_scope(ctx, mod, builder, function, node.children[3], scope_number, body_block)
            llvm_scope_leave()
            LLVMBuildBr(builder, condition_block)

            LLVMPositionBuilderAtEnd(builder, end_block)
        }
    }

    generate_llvm_scope :: proc(ctx: LLVMContextRef, mod: LLVMModuleRef, builder: LLVMBuilderRef, function: LLVMValueRef, scope: ^Node, scope_number: ^i32, bb: LLVMBasicBlockRef = nil) {
        scope_name := get_basic_block_name(scope_number^)
        basic_block : LLVMBasicBlockRef
        if bb != nil {
            basic_block = bb
        } else {
            basic_block = LLVMAppendBasicBlockInContext(ctx, function, scope_name)
        }
        llvm_scope_enter(scope_name, basic_block)
        LLVMPositionBuilderAtEnd(builder, basic_block)
        scope_number^ += 1
        for node in scope.children {
            #partial switch node.kind {
            case .Block:
                generate_llvm_scope(ctx, mod, builder, function, node, scope_number)
            case .Function:
                panic("FIXME: Implement function delarations")
            case .ExternFunction:
                type, value := generate_llvm_function_extern(ctx, mod, builder, node, scope_number)
                //name := LLVMGetValueName2(value, nil)
                llvm_top_scope().definitions[get_character_sum_of_dyn_arr(&node.value.([dynamic]u8))] = value
                llvm_top_scope().types[get_character_sum_of_dyn_arr(&node.value.([dynamic]u8))] = type
            case .Integer: fallthrough
            case .Float: fallthrough
            case .Identifier: fallthrough
            case .FunctionCall: fallthrough
            case .BinaryExpression:
                generate_llvm_expression(ctx, mod, builder, node)
            case .VariableDeclaration:
                type := generate_llvm_type_from_node(ctx, mod, builder, node.return_type)
                var := LLVMBuildAlloca(builder, type, strings.clone_to_cstring(string(node.children[0].value.([dynamic]u8)[:])))
                if node.children[2] != nil {
                    value := generate_llvm_expression(ctx, mod, builder, node.children[2])
                    LLVMBuildStore(builder, value, var)
                }
                llvm_top_scope().definitions[get_character_sum_of_dyn_arr(&node.children[0].value.([dynamic]u8))] = var
                llvm_top_scope().types[get_character_sum_of_dyn_arr(&node.children[0].value.([dynamic]u8))] = type
            case .If:
                generate_llvm_if(ctx, mod, builder, function, node, scope_number)
            case .For:
                generate_llvm_for(ctx, mod, builder, function, node, scope_number)
            case:
                fmt.panicf("FIXME: Implement other node kinds. Got: {}", node.kind)
            }
        }
        llvm_scope_leave()
    }

    generate_llvm_scope(ctx, mod, builder, main_function, node, &scope_number)
    //basic_block := LLVMAppendBasicBlockInContext(ctx, main_function, "end")
    //LLVMPositionBuilderAtEnd(builder, basic_block)
    LLVMBuildRet(builder, LLVMConstInt(int_32_type, 0, LLVMBool(1)))
}
