#include "gem/runtime.h"

#include <stdio.h>
#include <string.h>

static int check_plan(enum gem_model_size size, uint32_t layers, uint64_t embedding,
                      uint64_t query_heads, uint64_t swa_kv, uint64_t full_kv,
                      uint64_t ff, uint64_t expected_bytes, uint64_t expected_mixed_bytes,
                      uint64_t expected_mixed_128k_bytes)
{
    struct gem_model model;
    memset(&model, 0, sizeof(model));
    model.size = size;
    model.block_count = layers;
    model.context_length = 262144;
    model.sliding_window = 1024;
    model.embedding_length = embedding;
    model.head_count = query_heads;
    model.head_count_kv_swa = swa_kv;
    model.head_count_kv_full = full_kv;
    model.head_dimension_swa = 256;
    model.head_dimension_full = 512;
    model.feed_forward_length = ff;

    struct gem_runtime_plan plan;
    char error[256];
    if (gem_runtime_plan_init(&plan, &model, 32768, 2, error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        return 1;
    }
    if (plan.kv_bytes != expected_bytes || plan.layer_count != layers ||
        plan.swa_kv_element_size != 2 || plan.full_kv_element_size != 2 ||
        !plan.layers[0].sliding || plan.layers[5].sliding ||
        plan.layers[0].cache_capacity != 1024 || plan.layers[5].cache_capacity != 32768 ||
        plan.layers[0].kv_heads != swa_kv || plan.layers[5].kv_heads != full_kv ||
        plan.layers[0].head_dimension != 256 || plan.layers[5].head_dimension != 512 ||
        plan.layers[0].kv_element_size != 2 || plan.layers[5].kv_element_size != 2) {
        fprintf(stderr, "runtime plan mismatch for %uB\n", (unsigned)size);
        gem_runtime_plan_destroy(&plan);
        return 1;
    }
    for (uint32_t i = 0; i < plan.layer_count; ++i)
        if (plan.layers[i].key_offset % 256 || plan.layers[i].value_offset % 256) {
            fprintf(stderr, "unaligned layer %u\n", i);
            gem_runtime_plan_destroy(&plan);
            return 1;
    }
    gem_runtime_plan_destroy(&plan);
    if (gem_runtime_plan_init_types(&plan, &model, 32768, 4, 2,
                                    error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        return 1;
    }
    if (plan.kv_bytes != expected_mixed_bytes || plan.swa_kv_element_size != 4 ||
        plan.full_kv_element_size != 2 || plan.layers[0].kv_element_size != 4 ||
        plan.layers[5].kv_element_size != 2) {
        fprintf(stderr, "mixed runtime plan mismatch for %uB\n", (unsigned)size);
        gem_runtime_plan_destroy(&plan);
        return 1;
    }
    gem_runtime_plan_destroy(&plan);
    if (gem_runtime_plan_init_types(&plan, &model, 131072, 4, 2,
                                    error, sizeof(error)) != 0) {
        fprintf(stderr, "%s\n", error);
        return 1;
    }
    if (plan.kv_bytes != expected_mixed_128k_bytes ||
        plan.layers[0].cache_capacity != 1024 ||
        plan.layers[5].cache_capacity != 131072) {
        fprintf(stderr, "128K mixed runtime plan mismatch for %uB\n", (unsigned)size);
        gem_runtime_plan_destroy(&plan);
        return 1;
    }
    gem_runtime_plan_destroy(&plan);
    if (gem_runtime_plan_init_types(&plan, &model, 32768, 4, 3,
                                    error, sizeof(error)) == 0) {
        gem_runtime_plan_destroy(&plan);
        fprintf(stderr, "accepted invalid mixed KV element size\n");
        return 1;
    }
    if (gem_runtime_plan_init(&plan, &model, 262145, 2, error, sizeof(error)) == 0) {
        gem_runtime_plan_destroy(&plan);
        fprintf(stderr, "accepted oversized context\n");
        return 1;
    }
    return 0;
}

int main(void)
{
    if (check_plan(GEM_MODEL_12B, 48, 3840, 16, 8, 1, 15360,
                   UINT64_C(872415232), UINT64_C(1207959552),
                   UINT64_C(2818572288)) != 0)
        return 1;
    if (check_plan(GEM_MODEL_31B, 60, 5376, 32, 16, 4, 21504,
                   UINT64_C(3523215360), UINT64_C(4362076160),
                   UINT64_C(12415139840)) != 0)
        return 1;
    return 0;
}