#include "souffle/SouffleInterface.h"
#include "souffle.h"
#include <algorithm>
#include <string>
#include <unordered_map>
#include <vector>
#include <memory>

#ifndef ESTIMATED_AVERAGE_STRING_SIZE
#define ESTIMATED_AVERAGE_STRING_SIZE 32
#endif
#ifndef GROW_FACTOR
#define GROW_FACTOR 2
#endif

extern "C"
{

struct buf_data
{
private:
    std::unique_ptr<char[]> m_data;
    size_t m_size;

public:
    buf_data(size_t size)
        : m_data(std::make_unique<char[]>(size))
        , m_size(size)
    {
        assert(size);
    };

    void resize(size_t num_bytes)
    {
        m_data = std::make_unique<char[]>(num_bytes);
        m_size = num_bytes;
    }

    auto size() const
    {
        return m_size;
    }

    auto data() const
    {
        return m_data.get();
    }
};

struct souffle_interface
{
    std::unique_ptr<souffle::SouffleProgram> m_prog;
    buf_data m_buf;

    souffle_interface(souffle::SouffleProgram *prog)
        : m_prog(prog)
        , m_buf(4)
    {
        assert(prog);
    };

    char *get_buf(size_t num_bytes)
    {
        if (num_bytes > m_buf.size()) m_buf.resize(num_bytes);

        return m_buf.data();
    }
};

}

namespace helpers
{
using souffle_type = char;

inline auto parse_signature(const souffle::Relation& relation)
{
    const auto arity = relation.getArity();

    std::vector<souffle_type> types;
    types.reserve(arity);

    for (size_t i = 0; i < arity; ++i)
    {
        types.push_back(*relation.getAttrType(i));
    }

    return types;
}

inline bool relation_contains_strings(const souffle::Relation& relation)
{
    const auto types = parse_signature(relation);
    return std::any_of(types.begin(), types.end(), [](auto x) { return x == 's'; });
}

using offset_t = uint32_t;

using number_t = int32_t;
using unsigned_t = uint32_t;
using float_t = float;

template <typename T>
inline void serialize_value(souffle::tuple& tuple, char* buf, offset_t& offset)
{
    auto ptr = reinterpret_cast<T*>(buf);
    tuple >> *ptr;
    offset += sizeof(T);
}

template <typename T>
inline void deserialize_value(souffle::tuple& tuple, char* buf, offset_t& offset)
{
    auto ptr = reinterpret_cast<T*>(buf);
    tuple << *ptr;
    offset += sizeof(T);
}

inline void deserialize_symbol(souffle::tuple& tuple, char* buf, offset_t& offset)
{
    auto ptr = reinterpret_cast<uint32_t*>(buf);
    const auto num_bytes = *ptr;
    if (num_bytes == 0) {
        tuple << "";
        offset += sizeof(uint32_t);
        return;
    }

    auto string_ptr = reinterpret_cast<const char*>(buf) + sizeof(uint32_t);
    std::string str(string_ptr, num_bytes);
    tuple << str;
    offset += sizeof(uint32_t) + num_bytes;
}

using deserializer_t = void(*)(souffle::tuple&, char*, offset_t&);
using deserializer_map = std::unordered_map<souffle_type, deserializer_t>;

static const deserializer_map deserializers_map = {
    {'s', deserialize_symbol},
    {'i', deserialize_value<number_t>},
    {'u', deserialize_value<unsigned_t>},
    {'f', deserialize_value<float_t>}
};

using serializer_t = void(*)(souffle::tuple&, char*, offset_t&);
using serializer_map = std::unordered_map<souffle_type, serializer_t>;

static const serializer_map serializers_map = {
    {'i', serialize_value<number_t>},
    {'u', serialize_value<unsigned_t>},
    {'f', serialize_value<float_t>}
};

inline std::string unknown_souffle_type(souffle_type ty)
{
    std::string base_message = "Found unknown Souffle primitive type: ";
    return base_message + std::string(1, ty);
}

inline auto types_to_deserializer(const std::vector<souffle_type>& types)
{
    std::vector<deserializer_t> deserializers;
    deserializers.reserve(types.size());

    for (const auto& type: types)
    {
        const auto match = deserializers_map.find(type);
        assert(match != deserializers_map.end() && unknown_souffle_type(match->first).c_str());
        deserializers.push_back(match->second);
    }

    return [deserializers = std::move(deserializers)](souffle::tuple& tuple, char* buf, offset_t& offset)
    {
        for (const auto& deserializer : deserializers)
        {
            deserializer(tuple, buf + offset, offset);
        }
    };
}

inline auto types_to_serializer(const std::vector<souffle_type>& types)
{
    std::vector<serializer_t> serializers;
    serializers.reserve(types.size());

    for (const auto& type: types)
    {
        const auto match = serializers_map.find(type);
        assert(match != serializers_map.end() && unknown_souffle_type(match->first).c_str());
        serializers.push_back(match->second);
    }

    return [serializers = std::move(serializers)](souffle::tuple& tuple, char* buf, offset_t& offset)
    {
        for (const auto& serializer : serializers)
        {
            serializer(tuple, buf + offset, offset);
        }
    };
}

inline auto guess_tuple_size(const std::vector<souffle_type>& types)
{
    size_t size = 0;

    for (const auto& type : types)
    {
        size += type == 's'
             ? ESTIMATED_AVERAGE_STRING_SIZE
             : sizeof(number_t);
    }

    return size;
}

struct Serializer
{
public:
    inline Serializer(souffle_t *prog, const souffle::Relation& relation)
        : m_relation(relation)
        , m_types(parse_signature(relation))
        , m_buf(prog->m_buf)
    {
        auto tuple_size = guess_tuple_size(m_types);

        m_fact_count = relation.size();
        m_num_bytes = sizeof(uint32_t) + m_fact_count * tuple_size;
        m_offset = 0;

        // NOTE: we need to have atleast `m_num_bytes` large buffer, to make
        // memcpy later not write beyond the buffer.
        if (m_num_bytes > m_buf.size()) m_buf.resize(m_num_bytes);
    }

    inline const Serializer& serialize()
    {
        using serializer_t = void(*)(Serializer*, souffle::tuple&);
        using serializer_map_t = std::unordered_map<souffle_type, serializer_t>;

        serializer_t do_serialize_symbol = [](auto s, auto& t) {
            s->serialize_symbol(t);
        };
        serializer_t do_serialize_number = [](auto s, auto& t) {
            s->serialize_number(t);
        };
        serializer_t do_serialize_unsigned = [](auto s, auto& t) {
            s->serialize_unsigned(t);
        };
        serializer_t do_serialize_float = [](auto s, auto& t) {
            s->serialize_float(t);
        };

        static const serializer_map_t serializers_map = {
            {'s', do_serialize_symbol},
            {'i', do_serialize_number},
            {'u', do_serialize_unsigned},
            {'f', do_serialize_float},
        };

        std::vector<serializer_t> serializers;
        serializers.reserve(m_types.size());

        for (const auto& type: m_types)
        {
            const auto match = serializers_map.find(type);
            assert(match != serializers_map.end() && unknown_souffle_type(match->first).c_str());
            serializers.push_back(match->second);
        }

        const auto serialize = [this, serializers = std::move(serializers)](auto& tuple)
        {
            for (const auto& serializer : serializers)
            {
                serializer(this, tuple);
            }
        };

        auto buf = reinterpret_cast<uint32_t*>(m_buf.data());
        *buf = m_fact_count;
        m_offset += sizeof(uint32_t);

        for (auto& tuple: m_relation)
        {
            serialize(tuple);
        }

        return *this;
    }

    inline void serialize_number(souffle::tuple& tuple)
    {
        constexpr auto byte_count = sizeof(number_t);
        if (!has_remaining_bytes(byte_count)) {
            resize_buf(byte_count);
        }

        serialize_value<number_t>(tuple, m_buf.data() + m_offset, m_offset);
    }

    inline void serialize_unsigned(souffle::tuple& tuple)
    {
        constexpr auto byte_count = sizeof(unsigned_t);
        if (!has_remaining_bytes(byte_count)) {
            resize_buf(byte_count);
        }

        serialize_value<unsigned_t>(tuple, m_buf.data() + m_offset, m_offset);
    }

    inline void serialize_float(souffle::tuple& tuple)
    {
        constexpr auto byte_count = sizeof(float_t);
        if (!has_remaining_bytes(byte_count)) {
            resize_buf(byte_count);
        }

        serialize_value<float_t>(tuple, m_buf.data() + m_offset, m_offset);
    }

    inline void serialize_symbol(souffle::tuple& tuple)
    {
        std::string str;
        tuple >> str;
        const uint32_t num_bytes = str.length();

        auto total_byte_count = sizeof(uint32_t) + num_bytes;
        if (!has_remaining_bytes(total_byte_count)) {
            resize_buf(total_byte_count);
        }

        auto buf = m_buf.data() + m_offset;
        auto ptr = reinterpret_cast<uint32_t*>(buf);
        *ptr = num_bytes;

        // TODO: check if we can directly write into byte buf?
        auto string_ptr = reinterpret_cast<char*>(buf) + sizeof(uint32_t);
        std::copy(str.begin(), str.end(), string_ptr);
        m_offset += sizeof(uint32_t) + num_bytes;
    }

    inline bool has_remaining_bytes(size_t count) const
    {
        return m_num_bytes >= m_offset + count;
    }

    inline void resize_buf(size_t byte_count)
    {
        size_t grow_factor = GROW_FACTOR;
        while (m_offset + byte_count > m_num_bytes * grow_factor) {
            grow_factor *= 2;
        }
        const auto new_num_bytes = m_num_bytes * grow_factor;
        m_num_bytes = new_num_bytes;

        buf_data new_buf(new_num_bytes);
        memcpy(new_buf.data(), m_buf.data(), m_offset);
        std::swap(m_buf, new_buf);
    }

    inline byte_buf_t *to_buf() const
    {
        return reinterpret_cast<byte_buf_t*>(m_buf.data());
    }

private:
    const souffle::Relation& m_relation;
    std::vector<souffle_type> m_types;
    size_t m_fact_count;
    buf_data& m_buf;
    size_t m_num_bytes;
    offset_t m_offset;
};

inline byte_buf_t *serialize_slow(souffle_t *prog, const souffle::Relation& relation)
{
    Serializer s(prog, relation);
    return s.serialize().to_buf();
}

inline byte_buf_t *serialize_fast(souffle_t *prog, const souffle::Relation& relation)
{
    const auto types = parse_signature(relation);
    const auto serialize = types_to_serializer(types);

    const auto fact_count = relation.size();
    const auto tuple_size = guess_tuple_size(types);
    const auto num_bytes = sizeof(uint32_t) + fact_count * tuple_size;
    auto buf = prog->get_buf(num_bytes);
    const auto start_ptr = buf;

    offset_t offset = 0;

    auto ptr = reinterpret_cast<uint32_t*>(buf);
    *ptr = fact_count;
    offset += 4;

    for (auto& tuple: relation)
    {
        serialize(tuple, buf, offset);
    }

    return reinterpret_cast<byte_buf_t*>(start_ptr);
}

}  // namespace helpers

extern "C"
{
    souffle_t *souffle_init(const char *progName)
    {
        auto prog = souffle::ProgramFactory::newInstance(progName);
        return prog ? new souffle_interface(prog) : nullptr;
    }

    void souffle_free(souffle_t *program)
    {
        assert(program);
        delete program;
    }

    void souffle_set_num_threads(souffle_t *program, size_t num_cores)
    {
        assert(program);
        program->m_prog->setNumThreads(num_cores);
    }

    size_t souffle_get_num_threads(souffle_t *program)
    {
        assert(program);
        return program->m_prog->getNumThreads();
    }

    void souffle_run(souffle_t *program)
    {
        assert(program);
        program->m_prog->run();
    }

    void souffle_load_all(souffle_t *program, const char *input_directory)
    {
        assert(program);
        assert(input_directory);
        program->m_prog->loadAll(input_directory);
    }

    void souffle_print_all(souffle_t *program, const char *output_directory)
    {
        assert(program);
        assert(output_directory);
        program->m_prog->printAll(output_directory);
    }

    relation_t *souffle_relation(souffle_t *program, const char *relation_name)
    {
        assert(program);
        assert(relation_name);
        auto relation = program->m_prog->getRelation(relation_name);
        assert(relation);
        return reinterpret_cast<relation_t *>(relation);
    }

    bool souffle_contains_tuple(relation_t *rel, byte_buf_t *buf)
    {
        auto relation = reinterpret_cast<souffle::Relation *>(rel);
        auto data = reinterpret_cast<char*>(buf);
        assert(relation && "Relation is NULL in souffle_contains_tuple");
        assert(data && "byte buf is NULL in souffle_contains_tuple");

        auto& r = *relation;
        const auto types = helpers::parse_signature(r);
        const auto deserialize_tuple = helpers::types_to_deserializer(types);

        souffle::tuple tuple(relation);
        helpers::offset_t offset = 0;
        deserialize_tuple(tuple, data, offset);
        return r.contains(tuple);
    }

    void souffle_tuple_push_many(relation_t *rel, byte_buf_t *buf, size_t size)
    {
        auto relation = reinterpret_cast<souffle::Relation*>(rel);
        auto data = reinterpret_cast<char*>(buf);
        assert(data && "byte buf is NULL in souffle_tuple_push_many");
        assert(relation && "Relation is NULL in souffle_tuple_push_many");

        auto& r = *relation;
        const auto types = helpers::parse_signature(r);
        const auto deserialize_tuple = helpers::types_to_deserializer(types);

        helpers::offset_t offset = 0;
        for (size_t i = 0; i < size; ++i)
        {
            souffle::tuple tuple(relation);
            deserialize_tuple(tuple, data, offset);
            r.insert(tuple);
        }
    }

    byte_buf_t *souffle_tuple_pop_many(souffle_t *prog, relation_t *rel)
    {
        auto relation = reinterpret_cast<souffle::Relation*>(rel);
        assert(prog && "Program is NULL in souffle_tuple_pop_many");
        assert(relation && "Relation is NULL in souffle_tuple_pop_many");
        auto& r = *relation;
        return helpers::relation_contains_strings(r)
            ? helpers::serialize_slow(prog, r)
            : helpers::serialize_fast(prog, r);
    }
}