#include "sp_wrap.h"

#include <sentencepiece_processor.h>

#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>

struct SPProcessor {
    sentencepiece::SentencePieceProcessor proc;
};

static char *copy_err(const std::string &msg) {
    char *out = static_cast<char *>(std::malloc(msg.size() + 1));
    if (out != nullptr) {
        std::memcpy(out, msg.c_str(), msg.size() + 1);
    }
    return out;
}

SPProcessor *sp_load(const char *path, char **err) {
    if (err != nullptr) {
        *err = nullptr;
    }
    auto *p = new SPProcessor();
    const auto status = p->proc.Load(path);
    if (!status.ok()) {
        if (err != nullptr) {
            *err = copy_err(status.ToString());
        }
        delete p;
        return nullptr;
    }
    return p;
}

void sp_free(SPProcessor *p) { delete p; }

int sp_bos_id(const SPProcessor *p) { return p->proc.bos_id(); }

int sp_eos_id(const SPProcessor *p) { return p->proc.eos_id(); }

int sp_pad_id(const SPProcessor *p) { return p->proc.pad_id(); }

int sp_encode(const SPProcessor *p, const char *text, int **out_ids, int *out_len, char **err) {
    if (err != nullptr) {
        *err = nullptr;
    }
    if (out_ids != nullptr) {
        *out_ids = nullptr;
    }
    if (out_len != nullptr) {
        *out_len = 0;
    }
    std::vector<int> ids;
    const auto status = p->proc.Encode(text, &ids);
    if (!status.ok()) {
        if (err != nullptr) {
            *err = copy_err(status.ToString());
        }
        return 0;
    }
    if (ids.empty()) {
        return 1;
    }
    int *buf = static_cast<int *>(std::malloc(sizeof(int) * ids.size()));
    if (buf == nullptr) {
        if (err != nullptr) {
            *err = copy_err("malloc failed");
        }
        return 0;
    }
    for (size_t i = 0; i < ids.size(); i++) {
        buf[i] = ids[i];
    }
    *out_ids = buf;
    *out_len = static_cast<int>(ids.size());
    return 1;
}

char *sp_id_to_piece(const SPProcessor *p, int id, char **err) {
    if (err != nullptr) {
        *err = nullptr;
    }
    const std::string piece = p->proc.IdToPiece(id);
    return copy_err(piece);
}