#include <cstdio>
#include <cerrno>
#include <cstring>
#include <cassert>
#include <csignal>
#include <cmath>

#include <iostream>
#include <memory>

#include <sys/socket.h>
#include <sys/un.h>
#include <sys/mman.h>

#include <asm/ldt.h>

#include "common.hpp"
#include "dos_emu.hpp"

int loader_socket;

class mmap_deleter{
    size_t size;

public:
    mmap_deleter(size_t size)
    : size(size)
    {}

    void operator()(void* ptr) const {
        munmap(ptr, size);
    }
};

enum ErrorCode {
    ok,
    socket_creation_failed,
    connection_error
};

using mmap_pointer_t = std::unique_ptr<void, mmap_deleter>;

static ErrorCode connect_to_loader();

template<typename T>
static T receive_message();
static mmap_pointer_t map_memory(BinarySizeMessage const& message);
static void load_binary(mmap_pointer_t& binary, uint32_t binary_size);
static void register_signal_handlers();

int main(int, char*[]) {
    if(connect_to_loader() != ok) {
        return -1;
    }
    auto size_message = receive_message<BinarySizeMessage>();
    auto binary = map_memory(size_message);

    BinaryBaseOffsetMessage offset_message
    {
        .offset = (uint32_t)binary.get()
    };

    send(loader_socket, &offset_message, sizeof offset_message, 0);

    auto binary_load_message = receive_message<BinaryLoadMessage>();
    load_binary(binary, size_message.binary_size);

    register_signal_handlers();

    std::cout << std::hex << "Binary base: " << offset_message.offset << std::endl;

    asm (
        "jmp *%0        \n\t"
        :
        : "r" (binary_load_message.eip)
    );
}

static ErrorCode connect_to_loader() {
    if ((loader_socket = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) {
        perror("socket");
        return socket_creation_failed;
    }

    sockaddr_un addr;
    addr.sun_family = AF_UNIX;
    std::strncpy(addr.sun_path, socket_name.c_str(), sizeof addr.sun_path);

    size_t len = strnlen(addr.sun_path, sizeof addr.sun_path) + sizeof(addr.sun_family);
    if (connect(loader_socket, (struct sockaddr *)&addr, len) == -1) {
        perror("connect");
        return connection_error;
    }
    return ok;
}

template <typename T>
static T receive_message() {
    T message;
    uint32_t bytes_received = recv(loader_socket, &message, sizeof message, 0);
    assert((bytes_received == sizeof message) && "Didn't receive message");
    return message;
}

static mmap_pointer_t map_memory(BinarySizeMessage const& message) {
    void* binary = mmap(NULL, message.binary_size, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if(binary == MAP_FAILED) {
        std::cerr << "error during mmap: " << errno << " " << strerror(errno) << std::endl;
        exit(1);
    }
    return mmap_pointer_t(binary, message.binary_size);
}

static void load_binary(mmap_pointer_t& binary, uint32_t binary_size) {
    size_t bytes_received = recv(loader_socket, binary.get(), binary_size, 0);
    assert((bytes_received == binary_size) && "bytes_received did not match expected");
}

static void register_signal_handlers() {
    struct sigaction sa;

    memset(&sa, 0, sizeof sa);
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_SIGINFO;
    sa.sa_sigaction = dos_emu_handler;

    sigaction(SIGSEGV, &sa, NULL);
}