#include "wrapper.h"
#include <marisa.h>
#include <string>

extern "C" {

MarisaKeyset* marisa_keyset_new() {
    return reinterpret_cast<MarisaKeyset*>(new marisa::Keyset());
}

void marisa_keyset_delete(MarisaKeyset* keyset) {
    delete reinterpret_cast<marisa::Keyset*>(keyset);
}

void marisa_keyset_push_back(MarisaKeyset* keyset, const char* key, size_t length, float weight) {
    marisa::Keyset* ks = reinterpret_cast<marisa::Keyset*>(keyset);
    ks->push_back(std::string(key, length), weight);
}

size_t marisa_keyset_size(const MarisaKeyset* keyset) {
    const marisa::Keyset* ks = reinterpret_cast<const marisa::Keyset*>(keyset);
    return ks->size();
}

MarisaTrie* marisa_trie_new() {
    return reinterpret_cast<MarisaTrie*>(new marisa::Trie());
}

void marisa_trie_delete(MarisaTrie* trie) {
    delete reinterpret_cast<marisa::Trie*>(trie);
}

int marisa_trie_build(MarisaTrie* trie, MarisaKeyset* keyset) {
    try {
        marisa::Trie* tr = reinterpret_cast<marisa::Trie*>(trie);
        marisa::Keyset* ks = reinterpret_cast<marisa::Keyset*>(keyset);
        tr->build(*ks);
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_lookup(const MarisaTrie* trie, MarisaAgent* agent) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        marisa::Agent* ag = reinterpret_cast<marisa::Agent*>(agent);
        return tr->lookup(*ag) ? 1 : 0;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_reverse_lookup(const MarisaTrie* trie, MarisaAgent* agent) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        marisa::Agent* ag = reinterpret_cast<marisa::Agent*>(agent);
        tr->reverse_lookup(*ag);
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_common_prefix_search(const MarisaTrie* trie, MarisaAgent* agent) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        marisa::Agent* ag = reinterpret_cast<marisa::Agent*>(agent);
        return tr->common_prefix_search(*ag) ? 1 : 0;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_predictive_search(const MarisaTrie* trie, MarisaAgent* agent) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        marisa::Agent* ag = reinterpret_cast<marisa::Agent*>(agent);
        return tr->predictive_search(*ag) ? 1 : 0;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_size(const MarisaTrie* trie, size_t* size) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        *size = tr->size();
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_save(const MarisaTrie* trie, const char* filename) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        tr->save(filename);
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_load(MarisaTrie* trie, const char* filename) {
    try {
        marisa::Trie* tr = reinterpret_cast<marisa::Trie*>(trie);
        tr->load(filename);
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_mmap(MarisaTrie* trie, const char* filename) {
    try {
        marisa::Trie* tr = reinterpret_cast<marisa::Trie*>(trie);
        tr->mmap(filename);
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_io_size(const MarisaTrie* trie, size_t* size) {
    try {
        const marisa::Trie* tr = reinterpret_cast<const marisa::Trie*>(trie);
        *size = tr->io_size();
        return 1;
    } catch (...) {
        return 0;
    }
}

int marisa_trie_clear(MarisaTrie* trie) {
    try {
        marisa::Trie* tr = reinterpret_cast<marisa::Trie*>(trie);
        tr->clear();
        return 1;
    } catch (...) {
        return 0;
    }
}

MarisaAgent* marisa_agent_new() {
    return reinterpret_cast<MarisaAgent*>(new marisa::Agent());
}

void marisa_agent_delete(MarisaAgent* agent) {
    delete reinterpret_cast<marisa::Agent*>(agent);
}

void marisa_agent_set_query(MarisaAgent* agent, const char* query, size_t length) {
    marisa::Agent* ag = reinterpret_cast<marisa::Agent*>(agent);
    ag->set_query(query, length);
}

void marisa_agent_set_query_by_id(MarisaAgent* agent, size_t id) {
    marisa::Agent* ag = reinterpret_cast<marisa::Agent*>(agent);
    ag->set_query(id);
}

const char* marisa_agent_key_ptr(const MarisaAgent* agent) {
    const marisa::Agent* ag = reinterpret_cast<const marisa::Agent*>(agent);
    return ag->key().str().data();
}

size_t marisa_agent_key_length(const MarisaAgent* agent) {
    const marisa::Agent* ag = reinterpret_cast<const marisa::Agent*>(agent);
    return ag->key().str().length();
}

size_t marisa_agent_key_id(const MarisaAgent* agent) {
    const marisa::Agent* ag = reinterpret_cast<const marisa::Agent*>(agent);
    return ag->key().id();
}

}