#include "gmachine.hpp"
#include <iostream>

namespace lily {
    std::ostream& operator<<(std::ostream& os, instruction& inst) {
        inst.to_stream(os);
        return os;
    }

    std::ostream& instruction_slide::to_stream(std::ostream& os) {
        os << "slide(" << amount << ")";
        return os;
    }

    std::ostream& instruction_alloc::to_stream(std::ostream& os) {
        os << "alloc(" << amount << ")";
        return os;
    }

    std::ostream& instruction_pop::to_stream(std::ostream& os) {
        os << "pop(" << amount << ")";
        return os;
    }

    std::ostream& instruction_unwind::to_stream(std::ostream& os) {
        os << "unwind";
        return os;
    }

    std::ostream& instruction_push_global::to_stream(std::ostream& os) {
        os << "pushglobal(" << name << ")";
        return os;
    }

    std::ostream& instruction_push_int::to_stream(std::ostream& os) {
        os << "pushint(" << value << ")";
        return os;
    }

    std::ostream& instruction_push_str::to_stream(std::ostream& os) {
        os << "pushstr(" << str << ")";
        return os;
    }

    std::ostream& instruction_push::to_stream(std::ostream& os) {
        os << "push(" << offset << ")";
        return os;
    }

    std::ostream& instruction_mkapp::to_stream(std::ostream& os) {
        os << "mkapp";
        return os;
    }

    std::ostream& instruction_eval::to_stream(std::ostream& os) {
        os << "eval";
        return os;
    }

    std::ostream& instruction_op::to_stream(std::ostream& os) {
        os << "op(" << op_supercombinator(op) << ")";
        return os;
    }

    std::ostream& instruction_cond::to_stream(std::ostream& os) {
        os << "cond";
        return os;
    }

    std::ostream& instruction_update::to_stream(std::ostream& os) {
        os << "update(" << offset << ")";
        return os;
    }

    std::ostream& instruction_pack::to_stream(std::ostream& os) {
        os << "pack(" << constructor << ", " << arity << ")";
        return os;
    }

    std::ostream& instruction_split::to_stream(std::ostream& os) {
        os << "split(" << arity << ")";
        return os;
    }

    std::ostream& instruction_jump::to_stream(std::ostream& os) {
        os << "jump";
        return os;
    }

    void instruction_slide::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        builder.CreateCall(stack_slide_func, { stack, get_int32_constant(amount) });
    }

    void instruction_alloc::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        builder.CreateCall(stack_alloc_func, { stack, get_int32_constant(amount) });
    }

    void instruction_pop::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        builder.CreateCall(stack_popn_func, { stack, get_int32_constant(amount) });
    }

    void instruction_unwind::gen_llvm(llvm_context& ctx) {

    }

    void instruction_push_global::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* new_node = builder.CreateCall(malloc_node_global_func,
                { get_int8_constant(ctx.get_supercombinator_arity(name)), ctx.get_supercombinator_function(name) }, "temp");
        builder.CreateCall(stack_push_func, { stack, new_node });
    }

    void instruction_push_int::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* new_node = builder.CreateCall(malloc_node_num_func,
                { get_int32_constant(value) });
        builder.CreateCall(stack_push_func, { stack, new_node });
    }

    void instruction_push_str::gen_llvm(llvm_context& ctx) {
        
    }

    void instruction_push::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* peeked = builder.CreateCall(stack_peek_func, { stack, get_int32_constant(offset) }, "temp");
        builder.CreateCall(stack_push_func, { stack, peeked });
    }

    void instruction_mkapp::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* func = builder.CreateCall(stack_pop_func, { stack }, "temp");
        llvm::Value* param = builder.CreateCall(stack_pop_func, { stack }, "temp");
        llvm::Value* app_node = builder.CreateCall(malloc_node_app_func, { func, param }, "temp");
        builder.CreateCall(stack_push_func, { stack, app_node });
    }

    void instruction_eval::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* top = builder.CreateCall(stack_pop_func, { stack }, "temp");
        llvm::Value* evaluated = builder.CreateCall(eval_func, { top }, "temp");
        builder.CreateCall(stack_push_func, { stack, evaluated });
    }

    void instruction_op::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* param2 = builder.CreateCall(stack_pop_func, { stack }, "temp");
        llvm::Value* param1 = builder.CreateCall(stack_pop_func, { stack }, "temp");
        llvm::Value* param1_node_num = builder.CreatePointerCast(param1, llvm::PointerType::getUnqual(node_num_type), "temp");
        llvm::Value* param2_node_num = builder.CreatePointerCast(param2, llvm::PointerType::getUnqual(node_num_type), "temp");
        llvm::Value* param1_intptr = builder.CreateGEP(param1_node_num, { get_int32_constant(0), get_int32_constant(1) }, "temp");
        llvm::Value* param2_intptr = builder.CreateGEP(param2_node_num, { get_int32_constant(0), get_int32_constant(1) }, "temp");
        llvm::Value* param1_int = builder.CreateLoad(llvm::IntegerType::get(context, 32), param1_intptr, "temp");
        llvm::Value* param2_int = builder.CreateLoad(llvm::IntegerType::get(context, 32), param2_intptr, "temp");
        llvm::Value* op_result;
        switch(op) {
            case binop::add: op_result = builder.CreateAdd(param2_int, param1_int, "temp"); break;
            case binop::subtract: op_result = builder.CreateSub(param2_int, param1_int, "temp"); break;
            case binop::times: op_result = builder.CreateMul(param2_int, param1_int, "temp"); break;
            case binop::divide: op_result = builder.CreateSDiv(param2_int, param1_int, "temp"); break;
        }
        llvm::Value* new_node = builder.CreateCall(malloc_node_num_func, { op_result }, "temp");
        builder.CreateCall(stack_push_func, { stack, new_node });
    }

    void instruction_cond::gen_llvm(llvm_context& ctx) {
        
    }

    void instruction_update::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        builder.CreateCall(stack_update_func, { stack, get_int32_constant(offset) });
    }

    void instruction_pack::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* packed = builder.CreateCall(pack_func, { stack, get_int8_constant(constructor), get_int32_constant(arity) }, "temp");
        builder.CreateCall(stack_push_func, { stack, packed });
    }

    void instruction_split::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::Value* popped = builder.CreateCall(stack_pop_func, { stack }, "temp");
        builder.CreateCall(split_func, { stack, popped, get_int32_constant(arity) });
    }

    void instruction_jump::gen_llvm(llvm_context& ctx) {
        llvm::Value* stack = ctx.get_current_function()->arg_begin();
        llvm::BasicBlock* safety_block = llvm::BasicBlock::Create(context, "safety", ctx.get_current_function());
        llvm::Value* top = builder.CreateCall(stack_peek_func, { stack, get_int32_constant(0) }, "temp");
        llvm::Value* top_data = builder.CreatePointerCast(top, llvm::PointerType::getUnqual(node_data_type), "temp");
        llvm::Value* top_data_tag_ptr = builder.CreateGEP(top_data, { get_int32_constant(0), get_int32_constant(1) }, "temp");
        llvm::Value* top_data_tag = builder.CreateLoad(top_data_tag_ptr, "temp");
        llvm::SwitchInst* swtch = builder.CreateSwitch(top_data_tag, safety_block);

        std::vector<llvm::BasicBlock*> blocks;
        for(auto& branch : instructions) {
            llvm::BasicBlock* new_block = llvm::BasicBlock::Create(context, "branch", ctx.get_current_function());
            builder.SetInsertPoint(new_block);
            for(auto& i : branch) {
                i->gen_llvm(ctx);
            }
            builder.CreateBr(safety_block);
            blocks.push_back(new_block);
        }

        for(auto& pair : const_instructions) {
            swtch->addCase(get_int8_constant(pair.first), blocks[pair.second]);
        }

        builder.SetInsertPoint(safety_block);
    }
}