#include "molasses/parser_primitives.h"
#include "molasses/generator_primitives.h"
#include "molasses/errors.h"

namespace molasses {


    namespace unix_system {
#ifdef linux
        struct syscall1 {
            int64_t operator()(int64_t syscall_id, int64_t p1) const {
                int64_t ret;
                asm volatile("syscall" : "=a"(ret) : "0"(syscall_id), "D"(p1) : "rcx", "r11", "memory");
                return ret;
            }
        };

        struct syscall2 {
            int64_t operator()(int64_t syscall_id, int64_t p1, int64_t p2) const {
                int64_t ret;
                asm volatile("syscall" : "=a"(ret) : "0"(syscall_id), "D"(p1), "S"(p2) : "rcx", "r11", "memory");
                return ret;
            }
        };

        struct syscall3 {
            int64_t operator()(int64_t syscall_id, int64_t p1, int64_t p2, int64_t p3) const {
                int64_t ret;
                asm volatile("syscall"
                             : "=a"(ret)
                             : "0"(syscall_id), "D"(p1), "S"(p2), "d"(p3)
                             : "rcx", "r11", "memory");
                return ret;
            }
        };

        struct syscall4 {
            int64_t operator()(int64_t syscall_id, int64_t p1, int64_t p2, int64_t p3, int64_t p4) const {
                int64_t ret;
                register long r10 asm("r10") = p4;
                asm volatile("syscall"
                             : "=a"(ret)
                             : "0"(syscall_id), "D"(p1), "S"(p2), "d"(p3), "r"(r10)
                             : "rcx", "r11", "memory");
                return ret;
            }
        };

        struct syscall5 {
            int64_t operator()(int64_t syscall_id, int64_t p1, int64_t p2, int64_t p3, int64_t p4, int64_t p5) const {
                int64_t ret;
                register long r10 asm("r10") = p4;
                register long r8 asm("r8") = p5;
                asm volatile("syscall"
                             : "=a"(ret)
                             : "0"(syscall_id), "D"(p1), "S"(p2), "d"(p3), "r"(r10), "r"(r8)
                             : "rcx", "r11", "memory");
                return ret;
            }
        };

        struct syscall6 {
            int64_t operator()(int64_t syscall_id, int64_t p1, int64_t p2, int64_t p3, int64_t p4, int64_t p5, int64_t p6) const {
                int64_t ret;
                register long r10 asm("r10") = p4;
                register long r8 asm("r8") = p5;
                register long r9 asm("r9") = p6;
                asm volatile("syscall"
                             : "=a"(ret)
                             : "0"(syscall_id), "D"(p1), "S"(p2), "d"(p3), "r"(r10), "r"(r8), "r"(r9)
                             : "rcx", "r11", "memory");
                return ret;
            }
        };
#endif
    }


    std::string marshal(const std::string& target) {
        std::stringstream builder;
        bool is_first = true;
        for(char character : target) {
            if(isalpha(character)) {
                builder << character;
            } else if(isdigit(character) and is_first) {
                builder << "___" << (int)character << "___";
            } else {
                builder << "__" << (int)character << "__";
            }
            is_first = false;
        }
        return builder.str();
    }

    std::string escape(const std::string& target) {
        std::stringstream builder;
        for(char character : target) {
            switch(character) {
                case 0:
                    builder << "\\0";
                    break;
                case '\n':
                    builder << "\\n";
                    break;
                case '"':
                    builder << "\\\"";
                    break;
                case '\t':
                    builder << "\\t";
                    break;
                case '\'':
                    builder << "\\'";
                    break;
                default:
                    builder << character;
            }
        }
        return builder.str();
    }

    template<>
    parser_context register_i32_operations<architecture_t::x86_64_linux>(parser_context ctx) {
        ctx.operations.emplace_back(
            std::make_shared<molasses::primitive_operation>(
                std::string{"+"},
                std::vector<std::string>({"i32", "i32"}),
                std::vector<std::string>({"i32"}),
                std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rbx\n",
                    "    addl %ebx, %eax\n",
                    "    andl $0xFFFFFFFF, %eax\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_a = current_stack.top();current_stack.pop();
                    auto value_b = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int32_t>(value_a) or not std::holds_alternative<int32_t>(value_b)) {
                        throw interpreter_error("+ expects i32 i32 as input");
                    }
                    current_stack.emplace(get<int32_t>(value_a) + get<int32_t>(value_b));
                }
            )
        );
        ctx.operations.emplace_back(
            std::make_shared<molasses::primitive_operation>(
                std::string{"+_i64"},
                std::vector<std::string>({"i64", "i64"}),
                std::vector<std::string>({"i64"}),
                std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rbx\n",
                    "    addq %rbx, %rax\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_a = current_stack.top();current_stack.pop();
                    auto value_b = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int64_t>(value_a) or not std::holds_alternative<int64_t>(value_b)) {
                        throw interpreter_error("+_i64 expects i64 i64 as input");
                    }
                    current_stack.emplace(get<int64_t>(value_a) + get<int64_t>(value_b));
                }
            )
        );
        ctx.operations.emplace_back(
            std::make_shared<molasses::primitive_operation>(
                std::string{"/%_i64"},
                std::vector<std::string>({"i64", "i64"}),
                std::vector<std::string>({"i64", "i64"}),
                std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rbx\n",
                    "    divq %rbx, %rax\n",
                    "    pushq %rax\n",
                    "    pushq %rdx\n"
                    // TODO: this is actually unsigned division, so it needs improvements on negative numbers
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_a = current_stack.top();current_stack.pop();
                    auto value_b = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int64_t>(value_a) or not std::holds_alternative<int64_t>(value_b)) {
                        throw interpreter_error("/%_i64 expects i64 i64 as input");
                    } else if(get<int64_t>(value_b) == 0) {
                        throw interpreter_error("/%_i64 division by zero");
                    }
                    current_stack.emplace(get<int64_t>(value_a) / get<int64_t>(value_b));
                    current_stack.emplace(get<int64_t>(value_a) % get<int64_t>(value_b));
                }
            )
        );
        ctx.operations.emplace_back(
            std::make_shared<molasses::primitive_operation>(
                std::string{"u8-ptr_to_i64"},
                std::vector<std::string>({"u8 ptr"}),
                std::vector<std::string>({"i64"}),
                std::vector<std::string>({
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<ptr_type>(value) && not (get<ptr_type>(value).pointed->name() == "u8 ptr")) {
                        throw interpreter_error("u8-ptr_to_i64 expects u8 ptr as input");
                    }
                    current_stack.emplace(intptr_t(get<ptr_type>(value).ptr));
                }
            )
        );
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"*"},
				std::vector<std::string>({"i32", "i32"}),
				std::vector<std::string>({"i32"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rbx\n",
                    "    imull %ebx, %eax\n",
                    "    andl $0xFFFFFFFF, %eax\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_a = current_stack.top();current_stack.pop();
                    auto value_b = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int32_t>(value_a) or not std::holds_alternative<int32_t>(value_b)) {
                        throw interpreter_error("* expects i32 i32 as input");
                    }
                    current_stack.emplace(get<int32_t>(value_a) * get<int32_t>(value_b));
                }
			)
		);
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"-"},
				std::vector<std::string>({"i32", "i32"}),
				std::vector<std::string>({"i32"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rbx\n",
                    "    subl %ebx, %eax\n",
                    "    andl $0xFFFFFFFF, %eax\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_a = current_stack.top();current_stack.pop();
                    auto value_b = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int32_t>(value_a) or not std::holds_alternative<int32_t>(value_b)) {
                        throw interpreter_error("- expects i32 i32 as input");
                    }
                    current_stack.emplace(get<int32_t>(value_a) - get<int32_t>(value_b));
                }
			)
		);
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"i32-to-i64"},
				std::vector<std::string>({"i32"}),
				std::vector<std::string>({"i64"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    movslq %eax, %rax\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int32_t>(value)) {
                        throw interpreter_error("i32-to-i64 expects i32 as input");
                    }
                    current_stack.emplace(int64_t(get<int32_t>(value)));
                }
			)
		);
		
        ctx.operations.emplace_back(
            std::make_shared<molasses::primitive_operation>(
                std::string{"drop_i64"},
                std::vector<std::string>({"i64"}),
                std::vector<std::string>({}),
                std::vector<std::string>({
                    "    popq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value = current_stack.top();current_stack.pop();
                    if(not std::holds_alternative<int64_t>(value)) {
                        throw interpreter_error("i32-to-i64 expects i32 as input");
                    }
                }
            )
        );

        ctx.operations.emplace_back(
            std::make_shared<molasses::primitive_operation>(
                std::string{"syscall1"},
                std::vector<std::string>({"i64", "i64"}),
                std::vector<std::string>({"i64"}),
                std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rdi\n",
                    "    syscall\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_0 = current_stack.top();current_stack.pop();
                    auto value_1 = current_stack.top();current_stack.pop();
                    if(
                        not std::holds_alternative<int64_t>(value_0) or
                        not std::holds_alternative<int64_t>(value_1)
                    ) {
                        throw interpreter_error("syscall1 expects i64 i64 as input");
                    }
#ifdef linux
                    unix_system::syscall1{}(get<int64_t>(value_0), get<int64_t>(value_1));
#endif
                }
            )
        );
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"syscall2"},
				std::vector<std::string>({"i64", "i64", "i64"}),
				std::vector<std::string>({"i64"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rdi\n",
                    "    popq %rsi\n",
                    "    syscall\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_0 = current_stack.top();current_stack.pop();
                    auto value_1 = current_stack.top();current_stack.pop();
                    auto value_2 = current_stack.top();current_stack.pop();
                    if(
                        not std::holds_alternative<int64_t>(value_0) or
                        not std::holds_alternative<int64_t>(value_1) or
                        not std::holds_alternative<int64_t>(value_2)
                    ) {
                        throw interpreter_error("syscall2 expects i64 i64 i64 as input");
                    }
#ifdef linux
                    unix_system::syscall2{}(get<int64_t>(value_0), get<int64_t>(value_1), get<int64_t>(value_2));
#endif
                }
			)
		);
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"syscall3"},
				std::vector<std::string>({"i64", "i64", "i64", "i64"}),
				std::vector<std::string>({"i64"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rdi\n",
                    "    popq %rsi\n",
                    "    popq %rdx\n",
                    "    syscall\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_0 = current_stack.top();current_stack.pop();
                    auto value_1 = current_stack.top();current_stack.pop();
                    auto value_2 = current_stack.top();current_stack.pop();
                    auto value_3 = current_stack.top();current_stack.pop();
                    if(
                        not std::holds_alternative<int64_t>(value_0) or
                        not std::holds_alternative<int64_t>(value_1) or
                        not std::holds_alternative<int64_t>(value_2) or
                        not std::holds_alternative<int64_t>(value_3)
                    ) {
                        throw interpreter_error("syscall3 expects i64 i64 i64 i64 as input");
                    }
#ifdef linux
                    unix_system::syscall3{}(get<int64_t>(value_0), get<int64_t>(value_1), get<int64_t>(value_2), get<int64_t>(value_3));
#endif
                }
			)
		);
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"syscall4"},
				std::vector<std::string>({"i64", "i64", "i64", "i64", "i64"}),
				std::vector<std::string>({"i64"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rdi\n",
                    "    popq %rsi\n",
                    "    popq %rdx\n",
                    "    popq %r10\n",
                    "    syscall\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_0 = current_stack.top();current_stack.pop();
                    auto value_1 = current_stack.top();current_stack.pop();
                    auto value_2 = current_stack.top();current_stack.pop();
                    auto value_3 = current_stack.top();current_stack.pop();
                    auto value_4 = current_stack.top();current_stack.pop();
                    if(
                        not std::holds_alternative<int64_t>(value_0) or
                        not std::holds_alternative<int64_t>(value_1) or
                        not std::holds_alternative<int64_t>(value_2) or
                        not std::holds_alternative<int64_t>(value_3) or
                        not std::holds_alternative<int64_t>(value_4)
                    ) {
                        throw interpreter_error("syscall4 expects i64 i64 i64 i64 i64 as input");
                    }
    #ifdef linux
                    unix_system::syscall4{}(get<int64_t>(value_0), get<int64_t>(value_1), get<int64_t>(value_2), get<int64_t>(value_3), get<int64_t>(value_4));
    #endif
                }
			)
		);
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"syscall5"},
				std::vector<std::string>({"i64", "i64", "i64", "i64", "i64", "i64"}),
				std::vector<std::string>({"i64"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rdi\n",
                    "    popq %rsi\n",
                    "    popq %rdx\n",
                    "    popq %r10\n",
                    "    popq %r8\n",
                    "    syscall\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_0 = current_stack.top();current_stack.pop();
                    auto value_1 = current_stack.top();current_stack.pop();
                    auto value_2 = current_stack.top();current_stack.pop();
                    auto value_3 = current_stack.top();current_stack.pop();
                    auto value_4 = current_stack.top();current_stack.pop();
                    auto value_5 = current_stack.top();current_stack.pop();
                    if(
                        not std::holds_alternative<int64_t>(value_0) or
                        not std::holds_alternative<int64_t>(value_1) or
                        not std::holds_alternative<int64_t>(value_2) or
                        not std::holds_alternative<int64_t>(value_3) or
                        not std::holds_alternative<int64_t>(value_4) or
                        not std::holds_alternative<int64_t>(value_5)
                    ) {
                        throw interpreter_error("syscall5 expects i64 i64 i64 i64 i64 i64 as input");
                    }
        #ifdef linux
                    unix_system::syscall5{}(get<int64_t>(value_0), get<int64_t>(value_1), get<int64_t>(value_2), get<int64_t>(value_3), get<int64_t>(value_4), get<int64_t>(value_5));
        #endif
                }
			)
		);
		ctx.operations.emplace_back(
			std::make_shared<molasses::primitive_operation>(
				std::string{"syscall6"},
				std::vector<std::string>({"i64", "i64", "i64", "i64", "i64", "i64", "i64"}),
				std::vector<std::string>({"i64"}),
				std::vector<std::string>({
                    "    popq %rax\n",
                    "    popq %rdi\n",
                    "    popq %rsi\n",
                    "    popq %rdx\n",
                    "    popq %r10\n",
                    "    popq %r8\n",
                    "    popq %r9\n",
                    "    syscall\n",
                    "    pushq %rax\n"
                }),
                [](const generate_context&, interpreter_stack& current_stack){
                    auto value_0 = current_stack.top();current_stack.pop();
                    auto value_1 = current_stack.top();current_stack.pop();
                    auto value_2 = current_stack.top();current_stack.pop();
                    auto value_3 = current_stack.top();current_stack.pop();
                    auto value_4 = current_stack.top();current_stack.pop();
                    auto value_5 = current_stack.top();current_stack.pop();
                    auto value_6 = current_stack.top();current_stack.pop();
                    if(
                        not std::holds_alternative<int64_t>(value_0) or
                        not std::holds_alternative<int64_t>(value_1) or
                        not std::holds_alternative<int64_t>(value_2) or
                        not std::holds_alternative<int64_t>(value_3) or
                        not std::holds_alternative<int64_t>(value_4) or
                        not std::holds_alternative<int64_t>(value_5) or
                        not std::holds_alternative<int64_t>(value_6)
                    ) {
                        throw interpreter_error("syscall6 expects i64 i64 i64 i64 i64 i64 i64 as input");
                    }
            #ifdef linux
                    unix_system::syscall6{}(
                        get<int64_t>(value_0),
                        get<int64_t>(value_1),
                        get<int64_t>(value_2),
                        get<int64_t>(value_3),
                        get<int64_t>(value_4),
                        get<int64_t>(value_5),
                        get<int64_t>(value_6)
                    );
            #endif
                }
			)
		);
		
		return ctx;
	}

    template<>
    std::vector<std::string> generate_call<architecture_t::x86_64_linux>(const std::string& target) {
        return {
            "    call "+marshal(target)+"\n",
        };
    }

    template<>
    std::vector<std::string> generate_string<architecture_t::x86_64_linux>(const symbol& representation, const std::string& string_value) {
        return {
            "__EMITED_STRING_____"+std::to_string(representation.id)+"___:\n",
            "    .asciz \""+escape(string_value)+"\"\n",
        };
    }

    template<>
    std::vector<std::string> generate_push_string_ptr<architecture_t::x86_64_linux>(const symbol& representation) {
        return {
            "    pushq $__EMITED_STRING_____"+std::to_string(representation.id)+"___\n"
        };
    }

    template<>
    std::vector<std::string> generate_push_int32<architecture_t::x86_64_linux>(int32_t target) {
        return {
            "    pushq $" +std::to_string(target)+ "\n"
        };
    }

    template<>
    std::vector<std::string> generate_push_int64<architecture_t::x86_64_linux>(int64_t target) {
        return {
            "    pushq $" +std::to_string(target)+ "\n"
        };
    }

    template<>
	std::vector<std::string> generate_label<architecture_t::x86_64_linux>(const std::string& target) {
		return {
            marshal(target)+":\n"
		};
	}

    template<>
	std::vector<std::string> generate_return<architecture_t::x86_64_linux>() {
		return {
			"    // Return to caller\n",
			"    addq $-8, %r10\n",
			"    pushq (%r10)\n",
			"    retq\n"
		};
	}

    template<>
	std::vector<std::string> generate_enter<architecture_t::x86_64_linux>() {
		return {
			"    // Prepare the function stack\n",
			"    popq (%r10)\n"
			"    addq $8, %r10\n",
		};
	}

    template<>
	std::vector<std::string> initialize_stack<architecture_t::x86_64_linux>() {
		std::vector<std::string> operations = {
			"code:\n",
			"    .skip 1000000\n",
			".text\n",
			"    .globl _start\n",
			"initialize_callstack:\n",
			"    movq $9, %rax\n",
			"    movq $0, %rdi\n",
			"    movq $8192, %rsi\n",
			"    movq $3, %rdx\n",
			"    movq $34, %r10\n",
			"    movq $-1, %r8\n",
			"    movq $0, %r9\n",
			"    syscall\n",
			"    movq %rax, %r10\n",
			"    retq\n",
			"_start:\n",
			"    call initialize_callstack\n"
		};
		
		for(const auto& op : generate_call("main")) {
			operations.push_back(op);
		}
		
		for(const auto& op : std::vector<std::string>{
				"    movq $0, %rdi\n",
				"    movq $60, %rax\n",
				"    syscall\n"
			}
		) {
			operations.push_back(op);
		}
		return operations;
	}
}