#include "gem/tokenizer.h"

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static void bytes(FILE *file, const void *data, size_t size)
{
    if (fwrite(data, 1, size, file) != size)
        abort();
}

static void u32(FILE *file, uint32_t value)
{
    uint8_t data[4];
    for (unsigned i = 0; i < 4; ++i)
        data[i] = (uint8_t)(value >> (i * 8));
    bytes(file, data, sizeof(data));
}

static void u64(FILE *file, uint64_t value)
{
    uint8_t data[8];
    for (unsigned i = 0; i < 8; ++i)
        data[i] = (uint8_t)(value >> (i * 8));
    bytes(file, data, sizeof(data));
}

static void string(FILE *file, const char *value)
{
    u64(file, strlen(value));
    bytes(file, value, strlen(value));
}

static void key(FILE *file, const char *value, uint32_t type)
{
    string(file, value);
    u32(file, type);
}

static void string_value(FILE *file, const char *name, const char *value)
{
    key(file, name, GEM_GGUF_STRING);
    string(file, value);
}

static void u32_value(FILE *file, const char *name, uint32_t value)
{
    key(file, name, GEM_GGUF_UINT32);
    u32(file, value);
}

static void bool_value(FILE *file, const char *name, int value)
{
    uint8_t byte = value != 0;
    key(file, name, GEM_GGUF_BOOL);
    bytes(file, &byte, 1);
}

static void array_header(FILE *file, const char *name, uint32_t type, uint64_t count)
{
    key(file, name, GEM_GGUF_ARRAY);
    u32(file, type);
    u64(file, count);
}

static void write_tokenizer(const char *path)
{
    static const char *const tokens[] = {
        "<unk>", "<bos>", "<eos>", "h", "e", "l", "o", "he", "ll", "hell",
        "hello", "\xe2\x96\x81", "w", "\xe2\x96\x81w", "\n", "\n\n", "<0x21>",
        "<|turn>",
    };
    static const char *const merges[] = {"h e", "l l", "he ll", "hell o", "\xe2\x96\x81 w"};
    const uint32_t token_count = sizeof(tokens) / sizeof(tokens[0]);
    FILE *file = fopen(path, "wb");
    if (!file)
        abort();
    bytes(file, "GGUF", 4);
    u32(file, 3);
    u64(file, 0);
    u64(file, 8);
    string_value(file, "tokenizer.ggml.model", "gemma4");
    array_header(file, "tokenizer.ggml.tokens", GEM_GGUF_STRING, token_count);
    for (uint32_t i = 0; i < token_count; ++i)
        string(file, tokens[i]);
    array_header(file, "tokenizer.ggml.scores", GEM_GGUF_FLOAT32, token_count);
    for (uint32_t i = 0; i < token_count; ++i)
        u32(file, 0);
    array_header(file, "tokenizer.ggml.token_type", GEM_GGUF_INT32, token_count);
    for (uint32_t i = 0; i < token_count; ++i)
        u32(file, i < 3 ? 3 : i == 17 ? 4 : 1);
    array_header(file, "tokenizer.ggml.merges", GEM_GGUF_STRING,
                 sizeof(merges) / sizeof(merges[0]));
    for (size_t i = 0; i < sizeof(merges) / sizeof(merges[0]); ++i)
        string(file, merges[i]);
    u32_value(file, "tokenizer.ggml.bos_token_id", 1);
    u32_value(file, "tokenizer.ggml.eos_token_id", 2);
    bool_value(file, "tokenizer.ggml.add_bos_token", 1);
    if (fclose(file) != 0)
        abort();
}

static int check_tokens(const uint32_t *actual, size_t actual_count,
                        const uint32_t *expected, size_t expected_count)
{
    if (actual_count != expected_count)
        return -1;
    return memcmp(actual, expected, expected_count * sizeof(*expected));
}

int main(void)
{
    char path[] = "/tmp/gem-tokenizer-XXXXXX";
    int descriptor = mkstemp(path);
    if (descriptor < 0)
        return 2;
    close(descriptor);
    write_tokenizer(path);

    struct gem_gguf gguf;
    struct gem_tokenizer tokenizer;
    char error[256];
    if (gem_gguf_open(&gguf, path, error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        return 1;
    }
    if (gem_tokenizer_init(&tokenizer, &gguf, error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        gem_gguf_close(&gguf);
        return 1;
    }

    uint32_t *ids = NULL;
    size_t count = 0;
    const uint32_t expected[] = {1, 10, 13, 15};
    if (gem_tokenizer_encode(&tokenizer, "hello w\n\n", true, true, &ids, &count,
                             error, sizeof(error)) != 0 ||
        check_tokens(ids, count, expected, sizeof(expected) / sizeof(expected[0])) != 0) {
        fprintf(stderr, "tokenization mismatch: %s (count %zu)\n", error, count);
        free(ids);
        gem_tokenizer_destroy(&tokenizer);
        gem_gguf_close(&gguf);
        unlink(path);
        return 1;
    }
    free(ids);

    uint8_t piece[16];
    size_t piece_size;
    if (gem_tokenizer_decode_token(&tokenizer, 13, false, piece, sizeof(piece),
                                   &piece_size, error, sizeof(error)) != 0 ||
        piece_size != 2 || memcmp(piece, " w", 2) != 0 ||
        gem_tokenizer_decode_token(&tokenizer, 16, false, piece, sizeof(piece),
                                   &piece_size, error, sizeof(error)) != 0 ||
        piece_size != 1 || piece[0] != '!' ||
        gem_tokenizer_decode_token(&tokenizer, 17, true, piece, sizeof(piece),
                                   &piece_size, error, sizeof(error)) != 0 || piece_size != 0) {
        fprintf(stderr, "token decode mismatch: %s\n", error);
        gem_tokenizer_destroy(&tokenizer);
        gem_gguf_close(&gguf);
        unlink(path);
        return 1;
    }

    const uint32_t fallback_expected[] = {16};
    if (gem_tokenizer_encode(&tokenizer, "!", false, false, &ids, &count,
                             error, sizeof(error)) != 0 ||
        check_tokens(ids, count, fallback_expected, 1) != 0) {
        fprintf(stderr, "byte fallback mismatch\n");
        free(ids);
        gem_tokenizer_destroy(&tokenizer);
        gem_gguf_close(&gguf);
        unlink(path);
        return 1;
    }
    free(ids);

    const uint32_t special_expected[] = {17, 10};
    if (gem_tokenizer_encode(&tokenizer, "<|turn>hello", false, false, &ids, &count,
                             error, sizeof(error)) != 0 ||
        check_tokens(ids, count, special_expected, 2) != 0) {
        fprintf(stderr, "user-defined special token mismatch\n");
        free(ids);
        gem_tokenizer_destroy(&tokenizer);
        gem_gguf_close(&gguf);
        unlink(path);
        return 1;
    }
    free(ids);
    gem_tokenizer_destroy(&tokenizer);
    gem_gguf_close(&gguf);
    unlink(path);
    return 0;
}