#include "gem/model.h"

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

static int failures;

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 kv_string(FILE *file, const char *key, const char *value)
{
    string(file, key);
    u32(file, GEM_GGUF_STRING);
    string(file, value);
}

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

static void kv_f32(FILE *file, const char *key, float value)
{
    uint32_t bits;
    memcpy(&bits, &value, sizeof(bits));
    string(file, key);
    u32(file, GEM_GGUF_FLOAT32);
    u32(file, bits);
}

static void kv_u32_array(FILE *file, const char *key, uint32_t blocks, uint32_t swa, uint32_t full)
{
    string(file, key);
    u32(file, GEM_GGUF_ARRAY);
    u32(file, GEM_GGUF_UINT32);
    u64(file, blocks);
    for (uint32_t i = 0; i < blocks; ++i)
        u32(file, i % 6 == 5 ? full : swa);
}

static void kv_swa_pattern(FILE *file, const char *key, uint32_t blocks)
{
    string(file, key);
    u32(file, GEM_GGUF_ARRAY);
    u32(file, GEM_GGUF_BOOL);
    u64(file, blocks);
    for (uint32_t i = 0; i < blocks; ++i) {
        uint8_t value = i % 6 != 5;
        bytes(file, &value, 1);
    }
}

static void write_model(const char *path, const char *architecture, uint32_t blocks)
{
    FILE *file = fopen(path, "wb");
    if (!file)
        abort();
    bytes(file, "GGUF", 4);
    u32(file, 3);
    u64(file, 0);
    u64(file, 19);
    kv_string(file, "general.architecture", architecture);
    kv_u32(file, "gemma4.block_count", blocks);
    kv_u32(file, "gemma4.context_length", 262144);
    kv_u32(file, "gemma4.embedding_length", blocks == 48 ? 3840 : 5376);
    kv_u32(file, "gemma4.attention.head_count", blocks == 48 ? 16 : 32);
    kv_u32_array(file, "gemma4.attention.head_count_kv", blocks,
                 blocks == 48 ? 8 : 16, blocks == 48 ? 1 : 4);
    kv_swa_pattern(file, "gemma4.attention.sliding_window_pattern", blocks);
    kv_u32(file, "gemma4.attention.sliding_window", 1024);
    kv_u32(file, "gemma4.feed_forward_length", blocks == 48 ? 15360 : 21504);
    kv_u32(file, "gemma4.attention.key_length", 512);
    kv_u32(file, "gemma4.attention.key_length_swa", 256);
    kv_u32(file, "gemma4.attention.value_length", 512);
    kv_u32(file, "gemma4.attention.value_length_swa", 256);
    kv_u32(file, "gemma4.rope.dimension_count", 512);
    kv_u32(file, "gemma4.rope.dimension_count_swa", 256);
    kv_f32(file, "gemma4.attention.layer_norm_rms_epsilon", 1e-6f);
    kv_f32(file, "gemma4.rope.freq_base", 1000000.0f);
    kv_f32(file, "gemma4.rope.freq_base_swa", 10000.0f);
    kv_f32(file, "gemma4.final_logit_softcapping", 30.0f);
    if (fclose(file) != 0)
        abort();
}

static void expect_open(const char *path, enum gem_model_size size, uint32_t blocks)
{
    struct gem_model model;
    char error[256];
    if (gem_model_open_metadata(&model, path, error, sizeof(error)) != 0) {
        fprintf(stderr, "expected load success: %s\n", error);
        failures++;
        return;
    }
    if (model.size != size || model.block_count != blocks || model.context_length != 262144 ||
        model.sliding_window != 1024) {
        fprintf(stderr, "loaded metadata did not match\n");
        failures++;
    }
    gem_model_close(&model);
}

static void expect_failure(const char *path, const char *needle)
{
    struct gem_model model;
    char error[256];
    if (gem_model_open_metadata(&model, path, error, sizeof(error)) == 0) {
        fprintf(stderr, "expected load failure\n");
        gem_model_close(&model);
        failures++;
    } else if (!strstr(error, needle)) {
        fprintf(stderr, "unexpected error '%s', expected '%s'\n", error, needle);
        failures++;
    }
}

static void write_tensor_gguf(const char *path, uint64_t offset, size_t data_bytes)
{
    FILE *file = fopen(path, "wb");
    if (!file)
        abort();
    bytes(file, "GGUF", 4);
    u32(file, 3);
    u64(file, 1);
    u64(file, 0);
    string(file, "weight");
    u32(file, 2);
    u64(file, 2);
    u64(file, 2);
    u32(file, 0); /* F32 */
    u64(file, offset);
    while ((size_t)ftell(file) % 32)
        fputc(0, file);
    for (size_t i = 0; i < data_bytes; ++i)
        fputc(0, file);
    fclose(file);
}

static void expect_gguf(const char *path, int success, const char *needle)
{
    struct gem_gguf gguf;
    char error[256];
    int result = gem_gguf_open(&gguf, path, error, sizeof(error));
    if ((result == 0) != success) {
        fprintf(stderr, "unexpected GGUF result: %s\n", result ? error : "success");
        failures++;
    } else if (result != 0 && !strstr(error, needle)) {
        fprintf(stderr, "unexpected GGUF error '%s'\n", error);
        failures++;
    }
    if (result == 0)
        gem_gguf_close(&gguf);
}

static void expect_tensor_data(const char *path)
{
    struct gem_gguf gguf;
    struct gem_gguf_tensor foreign = {0};
    char error[256];
    if (gem_gguf_open(&gguf, path, error, sizeof(error)) != 0) {
        fprintf(stderr, "could not reopen tensor GGUF: %s\n", error);
        failures++;
        return;
    }
    const struct gem_gguf_tensor *tensor = gem_gguf_find_tensor(&gguf, "weight");
    const uint8_t *data = gem_gguf_tensor_data(&gguf, tensor);
    long page_size = sysconf(_SC_PAGESIZE);
    if (!tensor || !data || data != gguf.mapping + gguf.data_offset ||
        gem_gguf_tensor_data(&gguf, &foreign) != NULL || page_size <= 0 ||
        gguf.mapping_size < gguf.file_size || gguf.mapping_size % (size_t)page_size) {
        fprintf(stderr, "tensor data binding did not enforce mmap ownership\n");
        failures++;
    }
    gem_gguf_close(&gguf);
}

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

    write_model(path, "gemma4", 48);
    expect_open(path, GEM_MODEL_12B, 48);
    write_model(path, "gemma4", 60);
    expect_open(path, GEM_MODEL_31B, 60);
    write_model(path, "llama", 60);
    expect_failure(path, "architecture");
    write_model(path, "gemma4", 42);
    expect_failure(path, "block count");

    FILE *file = fopen(path, "wb");
    if (!file)
        abort();
    bytes(file, "GGUF", 4);
    u32(file, 3);
    fclose(file);
    expect_failure(path, "truncated");

    write_tensor_gguf(path, 0, 16);
    expect_gguf(path, 1, "");
    expect_tensor_data(path);
    write_tensor_gguf(path, 32, 48);
    expect_gguf(path, 0, "layout");
    write_tensor_gguf(path, 0, 15);
    expect_gguf(path, 0, "truncated");

    unlink(path);
    return failures ? 1 : 0;
}